feat: Read system prompt from config file

macie · Jan 2, 2024 · 6839423 · 6839423
1 parent 8be172c
commit 6839423
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/cmd/boludo/config.go b/cmd/boludo/config.go
@@ -191,6 +191,7 @@ type ConfigFile map[string]ModelSpec
 // ModelSpec represents a model specification in the configuration file.
 type ModelSpec struct {
 	Model         string
+	SystemPrompt  string
 	InitialPrompt string
 	Format        string
 	Creativity    float32
@@ -229,6 +230,7 @@ func (c *ConfigFile) UnmarshalTOML(data interface{}) error {
 	for configId := range definedConfigs {
 		defaultSpec := ModelSpec{
 			Model:         "",
+			SystemPrompt:  "",
 			InitialPrompt: "",
 			Format:        "",
 			Creativity:    llama.DefaultOptions.Temp,
@@ -244,6 +246,8 @@ func (c *ConfigFile) UnmarshalTOML(data interface{}) error {
 				defaultSpec.Cutoff = (float32)(v.(float64))
 			case "format":
 				defaultSpec.Format = v.(string)
+			case "system-prompt":
+				defaultSpec.SystemPrompt = v.(string)
 			case "initial-prompt":
 				defaultSpec.InitialPrompt = v.(string)
 			}
@@ -274,6 +278,7 @@ func (c *ConfigFile) Prompt(configId string) llama.Prompt {
 	if spec, ok := (*c)[configId]; ok {
 		return llama.Prompt{
 			Format: spec.Format,
+			System: spec.SystemPrompt,
 		}
 	}
 	return llama.Prompt{}

diff --git a/cmd/boludo/config_test.go b/cmd/boludo/config_test.go
@@ -119,6 +119,15 @@ func TestParseFile(t *testing.T) {
 				Cutoff:     0.0,
 			},
 		}},
+		{"[assistant]\nmodel = \"model.gguf\"\ninitial-prompt = 'Reword:'\nsystem-prompt = 'You are an assistant.'", ConfigFile{
+			"assistant": ModelSpec{
+				Model:         "model.gguf",
+				Format:        "",
+				InitialPrompt: "Reword:",
+				SystemPrompt:  "You are an assistant.",
+				Creativity:    1.0,
+			},
+		}},
 	}
 	for _, tc := range testcases {
 		tc := tc