diff --git a/template/template.go b/template/template.go index b5bfb16c..f7453791 100644 --- a/template/template.go +++ b/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 { if field, ok := n.(*parse.FieldNode); ok && slices.Contains(field.Ident, "Response") { cut = true + return false } 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{ "System": system, "Prompt": prompt, - "Response": "", + "Response": response, }); err != nil { return err } diff --git a/template/template_test.go b/template/template_test.go index ae0db80b..b46e1df5 100644 --- a/template/template_test.go +++ b/template/template_test.go @@ -260,6 +260,26 @@ func TestExecuteWithMessages(t *testing.T) { 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", []template{