template: preprocess message and collect system

This commit is contained in:
Michael Yang
2024-07-12 11:48:06 -07:00
parent 179737feb7
commit 36c87c433b
2 changed files with 23 additions and 67 deletions

View File

@@ -216,13 +216,11 @@ func TestExecuteWithMessages(t *testing.T) {
{"response", `[INST] {{ if .System }}{{ .System }}
{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`},
{"messages", `{{- $system := contents .Messages "system" -}}
{{- range $index, $_ := .Messages }}
{{- if eq .Role "user" }}[INST] {{ if $system }}{{ $system }}
{{- $system = "" }}
{"messages", `[INST] {{ if .System }}{{ .System }}
{{ end }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}
{{- end }}
{{ end }}
{{- range .Messages }}
{{- if eq .Role "user" }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}[INST] {{ end }}
{{- end }}`},
},
Values{
@@ -243,13 +241,11 @@ func TestExecuteWithMessages(t *testing.T) {
{"response", `[INST] {{ if .System }}{{ .System }}
{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`},
{"messages", `{{- $system := contents .Messages "system" -}}
{{- range $index, $_ := .Messages }}
{{- if eq .Role "user" }}[INST] {{ if $system }}{{ $system }}
{{- $system = "" }}
{"messages", `[INST] {{ if .System }}{{ .System }}
{{ end }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}
{{- end }}
{{ end }}
{{- range .Messages }}
{{- if eq .Role "user" }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}[INST] {{ end }}
{{- end }}`},
},
Values{
@@ -363,36 +359,3 @@ Answer: `,
})
}
}
func TestFuncs(t *testing.T) {
t.Run("contents", func(t *testing.T) {
cases := map[string]string{
"": "A\n\nB\n\nC\n\nD\n\nE\n\nF",
"system": "A\n\nF",
"user": "B\n\nE",
"assistant": "C\n\nD",
}
s := []*api.Message{
{Role: "system", Content: "A"},
{Role: "user", Content: "B"},
{Role: "assistant", Content: "C"},
{Role: "assistant", Content: "D"},
{Role: "user", Content: "E"},
{Role: "system", Content: "F"},
}
fn, ok := funcs["contents"].(func([]*api.Message, ...string) string)
if !ok {
t.Fatal("contents is not a function")
}
for k, v := range cases {
t.Run(k, func(t *testing.T) {
if diff := cmp.Diff(fn(s, k), v); diff != "" {
t.Errorf("mismatch (-got +want):\n%s", diff)
}
})
}
})
}