Browse Source

preserve last assistant message (#5802)

Jeffrey Morgan 9 months ago
parent
commit
20090f3172
2 changed files with 22 additions and 1 deletions
  1. 2 1
      template/template.go
  2. 20 0
      template/template_test.go

+ 2 - 1
template/template.go

@@ -264,6 +264,7 @@ func (t *Template) Execute(w io.Writer, v Values) error {
 	nodes := deleteNode(t.Template.Root.Copy(), func(n parse.Node) bool {
 	nodes := deleteNode(t.Template.Root.Copy(), func(n parse.Node) bool {
 		if field, ok := n.(*parse.FieldNode); ok && slices.Contains(field.Ident, "Response") {
 		if field, ok := n.(*parse.FieldNode); ok && slices.Contains(field.Ident, "Response") {
 			cut = true
 			cut = true
+			return false
 		}
 		}
 
 
 		return cut
 		return cut
@@ -273,7 +274,7 @@ func (t *Template) Execute(w io.Writer, v Values) error {
 	if err := template.Must(template.New("").AddParseTree("", &tree)).Execute(&b, map[string]any{
 	if err := template.Must(template.New("").AddParseTree("", &tree)).Execute(&b, map[string]any{
 		"System":   system,
 		"System":   system,
 		"Prompt":   prompt,
 		"Prompt":   prompt,
-		"Response": "",
+		"Response": response,
 	}); err != nil {
 	}); err != nil {
 		return err
 		return err
 	}
 	}

+ 20 - 0
template/template_test.go

@@ -260,6 +260,26 @@ func TestExecuteWithMessages(t *testing.T) {
 
 
 Hello friend![/INST] Hello human![INST] What is your name?[/INST] `,
 Hello friend![/INST] Hello human![INST] What is your name?[/INST] `,
 		},
 		},
+		{
+			"mistral assistant",
+			[]template{
+				{"no response", `[INST] {{ .Prompt }}[/INST] `},
+				{"response", `[INST] {{ .Prompt }}[/INST] {{ .Response }}`},
+				{"messages", `
+{{- range $i, $m := .Messages }}
+{{- if eq .Role "user" }}[INST] {{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}{{ end }}
+{{- end }}`},
+			},
+			Values{
+				Messages: []api.Message{
+					{Role: "user", Content: "Hello friend!"},
+					{Role: "assistant", Content: "Hello human!"},
+					{Role: "user", Content: "What is your name?"},
+					{Role: "assistant", Content: "My name is Ollama and I"},
+				},
+			},
+			`[INST] Hello friend![/INST] Hello human![INST] What is your name?[/INST] My name is Ollama and I`,
+		},
 		{
 		{
 			"chatml",
 			"chatml",
 			[]template{
 			[]template{