template draft to test on the mac

This commit is contained in:
Dave Lee 2023-07-20 19:51:17 -04:00
parent 5ea869bfc3
commit 72f3398211
No known key found for this signature in database
5 changed files with 101 additions and 66 deletions

View file

@ -60,10 +60,11 @@ type Functions struct {
}
type TemplateConfig struct {
Completion string `yaml:"completion"`
Functions string `yaml:"function"`
Chat string `yaml:"chat"`
Edit string `yaml:"edit"`
Completion string `yaml:"completion"`
Functions string `yaml:"function"`
Chat string `yaml:"chat"`
ChatMessage string `yaml:"chat_message"`
Edit string `yaml:"edit"`
}
type ConfigLoader struct {

View file

@ -112,7 +112,7 @@ func ChatEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx)
var predInput string
mess := []string{}
for _, i := range input.Messages {
for messageIndex, i := range input.Messages {
var content string
role := i.Role
if role == "system" {
@ -130,31 +130,47 @@ func ChatEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx)
}
r := config.Roles[role]
contentExists := i.Content != nil && *i.Content != ""
if r != "" {
if contentExists {
content = fmt.Sprint(r, " ", *i.Content)
if config.TemplateConfig.ChatMessage != "" {
chatMessageData := model.ChatMessageTemplateData{
SystemPrompt: requestSystemPrompt,
Role: r,
RoleName: role,
Content: *i.Content,
MessageIndex: messageIndex,
}
if i.FunctionCall != nil {
j, err := json.Marshal(i.FunctionCall)
if err == nil {
if contentExists {
content += "\n" + fmt.Sprint(r, " ", string(j))
} else {
content = fmt.Sprint(r, " ", string(j))
chatMessageTemplate, err := o.Loader.TemplateForChatMessage(config.TemplateConfig.ChatMessage, chatMessageData)
if err != nil {
log.Error().Msgf("error processing message %+v using template \"%s\". Skipping!", chatMessageData, config.TemplateConfig.ChatMessage)
continue
}
content = chatMessageTemplate
} else {
if r != "" {
if contentExists {
content = fmt.Sprint(r, " ", *i.Content)
}
if i.FunctionCall != nil {
j, err := json.Marshal(i.FunctionCall)
if err == nil {
if contentExists {
content += "\n" + fmt.Sprint(r, " ", string(j))
} else {
content = fmt.Sprint(r, " ", string(j))
}
}
}
}
} else {
if contentExists {
content = fmt.Sprint(*i.Content)
}
if i.FunctionCall != nil {
j, err := json.Marshal(i.FunctionCall)
if err == nil {
if contentExists {
content += "\n" + string(j)
} else {
content = string(j)
} else {
if contentExists {
content = fmt.Sprint(*i.Content)
}
if i.FunctionCall != nil {
j, err := json.Marshal(i.FunctionCall)
if err == nil {
if contentExists {
content += "\n" + string(j)
} else {
content = string(j)
}
}
}
}
@ -190,7 +206,6 @@ func ChatEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx)
templatedInput, err := o.Loader.TemplatePrefix(templateFile, model.PromptTemplateData{
Input: predInput,
Functions: funcs,
System: requestSystemPrompt,
})
if err == nil {
predInput = templatedInput

View file

@ -77,8 +77,7 @@ func CompletionEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fibe
// A model can have a "file.bin.tmpl" file associated with a prompt template prefix
templatedInput, err := o.Loader.TemplatePrefix(templateFile, model.PromptTemplateData{
Input: predInput,
System: config.SystemPrompt,
Input: predInput,
})
if err == nil {
predInput = templatedInput
@ -124,8 +123,7 @@ func CompletionEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fibe
for _, i := range config.PromptStrings {
// A model can have a "file.bin.tmpl" file associated with a prompt template prefix
templatedInput, err := o.Loader.TemplatePrefix(templateFile, model.PromptTemplateData{
Input: i,
System: config.SystemPrompt,
Input: i,
})
if err == nil {
i = templatedInput

View file

@ -37,7 +37,6 @@ func EditEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx)
templatedInput, err := o.Loader.TemplatePrefix(templateFile, model.PromptTemplateData{
Input: i,
Instruction: input.Instruction,
System: config.SystemPrompt,
})
if err == nil {
i = templatedInput

View file

@ -21,35 +21,50 @@ import (
// These are the definitions of all possible variables LocalAI will currently populate for use in a prompt template file
// Please note: Not all of these are populated on every endpoint - your template should either be tested for each endpoint you map it to, or tolerant of zero values.
type PromptTemplateData struct {
System string
Input string
Instruction string
Functions []grammar.Function
Input string
Instruction string
Functions []grammar.Function
MessageIndex int
}
// TODO: Ask mudler about FunctionCall stuff being useful at the message level?
type ChatMessageTemplateData struct {
SystemPrompt string
Role string
RoleName string
Content string
MessageIndex int
}
type ModelLoader struct {
ModelPath string
mu sync.Mutex
// TODO: this needs generics
models map[string]*grpc.Client
grpcProcesses map[string]*process.Process
promptsTemplates map[string]*template.Template
models map[string]*grpc.Client
grpcProcesses map[string]*process.Process
promptsTemplates map[string]*template.Template
chatMessageTemplates map[string]*template.Template
}
func NewModelLoader(modelPath string) *ModelLoader {
return &ModelLoader{
ModelPath: modelPath,
models: make(map[string]*grpc.Client),
promptsTemplates: make(map[string]*template.Template),
grpcProcesses: make(map[string]*process.Process),
ModelPath: modelPath,
models: make(map[string]*grpc.Client),
promptsTemplates: make(map[string]*template.Template),
chatMessageTemplates: make(map[string]*template.Template),
grpcProcesses: make(map[string]*process.Process),
}
}
func (ml *ModelLoader) ExistsInModelPath(s string) bool {
_, err := os.Stat(filepath.Join(ml.ModelPath, s))
func existsInModelPath(modelPath string, s string) bool {
_, err := os.Stat(filepath.Join(modelPath, s))
return err == nil
}
func (ml *ModelLoader) ExistsInModelPath(s string) bool {
return existsInModelPath(ml.ModelPath, s)
}
func (ml *ModelLoader) ListModels() ([]string, error) {
files, err := ioutil.ReadDir(ml.ModelPath)
if err != nil {
@ -69,24 +84,23 @@ func (ml *ModelLoader) ListModels() ([]string, error) {
return models, nil
}
func (ml *ModelLoader) TemplatePrefix(modelName string, in PromptTemplateData) (string, error) {
ml.mu.Lock()
defer ml.mu.Unlock()
// TODO: Split ModelLoader and TemplateLoader? Just to keep things more organized. Left together to share a mutex until I look into that.
m, ok := ml.promptsTemplates[modelName]
func evaluateTemplate[T any](templateName string, in T, modelPath string, templateMap *(map[string]*template.Template), mutex *sync.Mutex) (string, error) {
mutex.Lock()
defer mutex.Unlock()
m, ok := (*templateMap)[templateName]
if !ok {
modelFile := filepath.Join(ml.ModelPath, modelName)
if err := ml.loadTemplateIfExists(modelName, modelFile); err != nil {
return "", err
}
t, exists := ml.promptsTemplates[modelName]
if exists {
m = t
// return "", fmt.Errorf("template not loaded: %s", templateName)
loadErr := loadTemplateIfExists(templateName, modelPath, templateMap)
if loadErr != nil {
return "", loadErr
}
m = (*templateMap)[templateName] // ok is not important since we check m on the next line
}
if m == nil {
return "", fmt.Errorf("failed loading any template")
return "", fmt.Errorf("failed loading a template for %s", templateName)
}
var buf bytes.Buffer
@ -97,21 +111,29 @@ func (ml *ModelLoader) TemplatePrefix(modelName string, in PromptTemplateData) (
return buf.String(), nil
}
func (ml *ModelLoader) loadTemplateIfExists(modelName, modelFile string) error {
func (ml *ModelLoader) TemplatePrefix(templateName string, in PromptTemplateData) (string, error) {
return evaluateTemplate[PromptTemplateData](templateName, in, ml.ModelPath, &(ml.promptsTemplates), &(ml.mu))
}
func (ml *ModelLoader) TemplateForChatMessage(templateName string, messageData ChatMessageTemplateData) (string, error) {
return evaluateTemplate[ChatMessageTemplateData](templateName, messageData, ml.ModelPath, &(ml.chatMessageTemplates), &(ml.mu))
}
func loadTemplateIfExists(templateName, modelPath string, templateMap *(map[string]*template.Template)) error {
// Check if the template was already loaded
if _, ok := ml.promptsTemplates[modelName]; ok {
if _, ok := (*templateMap)[templateName]; ok {
return nil
}
// Check if the model path exists
// skip any error here - we run anyway if a template does not exist
modelTemplateFile := fmt.Sprintf("%s.tmpl", modelName)
modelTemplateFile := fmt.Sprintf("%s.tmpl", templateName)
if !ml.ExistsInModelPath(modelTemplateFile) {
if !existsInModelPath(modelPath, modelTemplateFile) {
return nil
}
dat, err := os.ReadFile(filepath.Join(ml.ModelPath, modelTemplateFile))
dat, err := os.ReadFile(filepath.Join(modelPath, modelTemplateFile))
if err != nil {
return err
}
@ -121,7 +143,7 @@ func (ml *ModelLoader) loadTemplateIfExists(modelName, modelFile string) error {
if err != nil {
return err
}
ml.promptsTemplates[modelName] = tmpl
(*templateMap)[templateName] = tmpl
return nil
}
@ -146,7 +168,7 @@ func (ml *ModelLoader) LoadModel(modelName string, loader func(string) (*grpc.Cl
}
// If there is a prompt template, load it
if err := ml.loadTemplateIfExists(modelName, modelFile); err != nil {
if err := loadTemplateIfExists(modelName, ml.ModelPath, &(ml.promptsTemplates)); err != nil {
return nil, err
}