feat(defaults): add defaults for Command-R models (#2529)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
mudler · Jun 9, 2024 · d9109ff · d9109ff
1 parent d7e1372
commit d9109ff
Showing 1 changed file with 47 additions and 7 deletions.
diff --git a/core/config/guesser.go b/core/config/guesser.go
@@ -13,14 +13,48 @@ type familyType uint8
 
 const (
 	Unknown familyType = iota
-	LLaMa3             = iota
-	LLama2             = iota
+	LLaMa3
+	LLama2
+	CommandR
 )
 
-var defaultsTemplate map[familyType]TemplateConfig = map[familyType]TemplateConfig{
+type settingsConfig struct {
+	StopWords      []string
+	TemplateConfig TemplateConfig
+}
+
+var defaultsSettings map[familyType]settingsConfig = map[familyType]settingsConfig{
 	LLaMa3: {
-		Chat:        "<|begin_of_text|>{{.Input }}\n<|start_header_id|>assistant<|end_header_id|>",
-		ChatMessage: "<|start_header_id|>{{ .RoleName }}<|end_header_id|>\n\n{{.Content }}<|eot_id|>",
+		StopWords: []string{"<|eot_id|>"},
+		TemplateConfig: TemplateConfig{
+			Chat:        "<|begin_of_text|>{{.Input }}\n<|start_header_id|>assistant<|end_header_id|>",
+			ChatMessage: "<|start_header_id|>{{ .RoleName }}<|end_header_id|>\n\n{{.Content }}<|eot_id|>",
+		},
+	},
+	CommandR: {
+		TemplateConfig: TemplateConfig{
+			Chat: "{{.Input -}}<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+			Functions: `<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>
+You are a function calling AI model, you can call the following functions:
+## Available Tools
+{{range .Functions}}
+- {"type": "function", "function": {"name": "{{.Name}}", "description": "{{.Description}}", "parameters": {{toJson .Parameters}} }}
+{{end}}
+When using a tool, reply with JSON, for instance {"name": "tool_name", "arguments": {"param1": "value1", "param2": "value2"}}
+<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{{.Input -}}`,
+			ChatMessage: `{{if eq .RoleName "user" -}}
+<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{.Content}}<|END_OF_TURN_TOKEN|>
+{{- else if eq .RoleName "system" -}}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{.Content}}<|END_OF_TURN_TOKEN|>
+{{- else if eq .RoleName "assistant" -}}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{{.Content}}<|END_OF_TURN_TOKEN|>
+{{- else if eq .RoleName "tool" -}}
+<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{.Content}}<|END_OF_TURN_TOKEN|>
+{{- else if .FunctionCall -}}
+<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{{toJson .FunctionCall}}}<|END_OF_TURN_TOKEN|>
+{{- end -}}`,
+		},
+		StopWords: []string{"<|END_OF_TURN_TOKEN|>"},
 	},
 }
 
@@ -68,10 +102,14 @@ func guessDefaultsFromFile(cfg *BackendConfig, modelPath string) {
 		return
 	}
 
-	templ, ok := defaultsTemplate[family]
+	// identify template
+	settings, ok := defaultsSettings[family]
 	if ok {
-		cfg.TemplateConfig = templ
+		cfg.TemplateConfig = settings.TemplateConfig
 		log.Debug().Any("family", family).Msgf("guessDefaultsFromFile: guessed template %+v", cfg.TemplateConfig)
+		if len(cfg.StopWords) == 0 {
+			cfg.StopWords = settings.StopWords
+		}
 	} else {
 		log.Debug().Any("family", family).Msgf("guessDefaultsFromFile: no template found for family")
 	}
@@ -81,6 +119,8 @@ func identifyFamily(f *gguf.GGUFFile) familyType {
 	switch {
 	case f.Architecture().Architecture == "llama" && f.Tokenizer().EOSTokenID == 128009:
 		return LLaMa3
+	case f.Architecture().Architecture == "command-r" && f.Tokenizer().EOSTokenID == 255001:
+		return CommandR
 	}
 
 	return Unknown