template_test.go 9.97 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
8
9
10
package template

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"os"
	"path/filepath"
	"slices"
Michael Yang's avatar
Michael Yang committed
11
	"strings"
Michael Yang's avatar
Michael Yang committed
12
13
	"testing"

Michael Yang's avatar
Michael Yang committed
14
	"github.com/google/go-cmp/cmp"
Michael Yang's avatar
Michael Yang committed
15
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
	"github.com/ollama/ollama/llm"
)

func TestNamed(t *testing.T) {
	f, err := os.Open(filepath.Join("testdata", "templates.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		var ss map[string]string
		if err := json.Unmarshal(scanner.Bytes(), &ss); err != nil {
			t.Fatal(err)
		}

		for k, v := range ss {
			t.Run(k, func(t *testing.T) {
				kv := llm.KV{"tokenizer.chat_template": v}
				s := kv.ChatTemplate()
				r, err := Named(s)
				if err != nil {
					t.Fatal(err)
				}

				if r.Name != k {
					t.Errorf("expected %q, got %q", k, r.Name)
				}

				var b bytes.Buffer
				if _, err := io.Copy(&b, r.Reader()); err != nil {
					t.Fatal(err)
				}

Michael Yang's avatar
Michael Yang committed
51
				tmpl, err := Parse(b.String())
Michael Yang's avatar
Michael Yang committed
52
53
54
55
56
57
58
59
60
61
62
63
				if err != nil {
					t.Fatal(err)
				}

				if tmpl.Tree.Root.String() == "" {
					t.Errorf("empty %s template", k)
				}
			})
		}
	}
}

Michael Yang's avatar
Michael Yang committed
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
func TestTemplate(t *testing.T) {
	cases := make(map[string][]api.Message)
	for _, mm := range [][]api.Message{
		{
			{Role: "user", Content: "Hello, how are you?"},
		},
		{
			{Role: "user", Content: "Hello, how are you?"},
			{Role: "assistant", Content: "I'm doing great. How can I help you today?"},
			{Role: "user", Content: "I'd like to show off how chat templating works!"},
		},
		{
			{Role: "system", Content: "You are a helpful assistant."},
			{Role: "user", Content: "Hello, how are you?"},
			{Role: "assistant", Content: "I'm doing great. How can I help you today?"},
			{Role: "user", Content: "I'd like to show off how chat templating works!"},
		},
	} {
		var roles []string
		for _, m := range mm {
			roles = append(roles, m.Role)
		}

		cases[strings.Join(roles, "-")] = mm
	}

	matches, err := filepath.Glob("*.gotmpl")
	if err != nil {
		t.Fatal(err)
	}

	for _, match := range matches {
		t.Run(match, func(t *testing.T) {
			bts, err := os.ReadFile(match)
			if err != nil {
				t.Fatal(err)
			}

			tmpl, err := Parse(string(bts))
			if err != nil {
				t.Fatal(err)
			}

			for n, tt := range cases {
108
				var actual bytes.Buffer
Michael Yang's avatar
Michael Yang committed
109
110
111
112
113
114
115
116
117
118
				t.Run(n, func(t *testing.T) {
					if err := tmpl.Execute(&actual, Values{Messages: tt}); err != nil {
						t.Fatal(err)
					}

					expect, err := os.ReadFile(filepath.Join("testdata", match, n))
					if err != nil {
						t.Fatal(err)
					}

119
120
121
122
123
124
125
126
					bts := actual.Bytes()

					if slices.Contains([]string{"chatqa.gotmpl", "llama2-chat.gotmpl", "mistral-instruct.gotmpl", "openchat.gotmpl", "vicuna.gotmpl"}, match) && bts[len(bts)-1] == ' ' {
						t.Log("removing trailing space from output")
						bts = bts[:len(bts)-1]
					}

					if diff := cmp.Diff(bts, expect); diff != "" {
Michael Yang's avatar
Michael Yang committed
127
128
129
						t.Errorf("mismatch (-got +want):\n%s", diff)
					}
				})
130
131

				t.Run("legacy", func(t *testing.T) {
132
					t.Skip("legacy outputs are currently default outputs")
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
					var legacy bytes.Buffer
					if err := tmpl.Execute(&legacy, Values{Messages: tt, forceLegacy: true}); err != nil {
						t.Fatal(err)
					}

					legacyBytes := legacy.Bytes()
					if slices.Contains([]string{"chatqa.gotmpl", "openchat.gotmpl", "vicuna.gotmpl"}, match) && legacyBytes[len(legacyBytes)-1] == ' ' {
						t.Log("removing trailing space from legacy output")
						legacyBytes = legacyBytes[:len(legacyBytes)-1]
					} else if slices.Contains([]string{"codellama-70b-instruct.gotmpl", "llama2-chat.gotmpl", "mistral-instruct.gotmpl"}, match) {
						t.Skip("legacy outputs cannot be compared to messages outputs")
					}

					if diff := cmp.Diff(legacyBytes, actual.Bytes()); diff != "" {
						t.Errorf("mismatch (-got +want):\n%s", diff)
					}
				})
Michael Yang's avatar
Michael Yang committed
150
151
152
153
154
			}
		})
	}
}

Michael Yang's avatar
Michael Yang committed
155
156
func TestParse(t *testing.T) {
	cases := []struct {
Michael Yang's avatar
Michael Yang committed
157
158
		template string
		vars     []string
Michael Yang's avatar
Michael Yang committed
159
	}{
Michael Yang's avatar
Michael Yang committed
160
161
		{"{{ .Prompt }}", []string{"prompt", "response"}},
		{"{{ .System }} {{ .Prompt }}", []string{"prompt", "response", "system"}},
Michael Yang's avatar
Michael Yang committed
162
		{"{{ .System }} {{ .Prompt }} {{ .Response }}", []string{"prompt", "response", "system"}},
Michael Yang's avatar
Michael Yang committed
163
		{"{{ with .Tools }}{{ . }}{{ end }} {{ .System }} {{ .Prompt }}", []string{"prompt", "response", "system", "tools"}},
Michael Yang's avatar
Michael Yang committed
164
		{"{{ range .Messages }}{{ .Role }} {{ .Content }}{{ end }}", []string{"content", "messages", "role"}},
165
166
167
168
169
170
		{`{{- range .Messages }}
{{- if eq .Role "system" }}SYSTEM:
{{- else if eq .Role "user" }}USER:
{{- else if eq .Role "assistant" }}ASSISTANT:
{{- end }} {{ .Content }}
{{- end }}`, []string{"content", "messages", "role"}},
171
172
173
174
175
176
177
178
179
180
181
182
		{`{{- if .Messages }}
{{- range .Messages }}<|im_start|>{{ .Role }}
{{ .Content }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ else -}}
{{ if .System }}<|im_start|>system
{{ .System }}<|im_end|>
{{ end }}{{ if .Prompt }}<|im_start|>user
{{ .Prompt }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ .Response }}<|im_end|>
{{- end -}}`, []string{"content", "messages", "prompt", "response", "role", "system"}},
Michael Yang's avatar
Michael Yang committed
183
184
185
186
187
188
189
190
191
	}

	for _, tt := range cases {
		t.Run("", func(t *testing.T) {
			tmpl, err := Parse(tt.template)
			if err != nil {
				t.Fatal(err)
			}

192
193
			if diff := cmp.Diff(tmpl.Vars(), tt.vars); diff != "" {
				t.Errorf("mismatch (-got +want):\n%s", diff)
Michael Yang's avatar
Michael Yang committed
194
195
196
197
			}
		})
	}
}
Michael Yang's avatar
Michael Yang committed
198
199

func TestExecuteWithMessages(t *testing.T) {
Michael Yang's avatar
Michael Yang committed
200
201
202
203
	type template struct {
		name     string
		template string
	}
Michael Yang's avatar
Michael Yang committed
204
	cases := []struct {
Michael Yang's avatar
Michael Yang committed
205
206
		name      string
		templates []template
Michael Yang's avatar
Michael Yang committed
207
208
209
210
		values    Values
		expected  string
	}{
		{
Michael Yang's avatar
Michael Yang committed
211
212
			"mistral",
			[]template{
213
214
215
216
217
218
				{"no response", `[INST] {{ if .System }}{{ .System }}

{{ end }}{{ .Prompt }}[/INST] `},
				{"response", `[INST] {{ if .System }}{{ .System }}

{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`},
219
				{"messages", `[INST] {{ if .System }}{{ .System }}
220

221
222
223
{{ end }}
{{- range .Messages }}
{{- if eq .Role "user" }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}[INST] {{ end }}
Michael Yang's avatar
Michael Yang committed
224
{{- end }}`},
Michael Yang's avatar
Michael Yang committed
225
226
227
228
229
			},
			Values{
				Messages: []api.Message{
					{Role: "user", Content: "Hello friend!"},
					{Role: "assistant", Content: "Hello human!"},
Michael Yang's avatar
Michael Yang committed
230
					{Role: "user", Content: "What is your name?"},
Michael Yang's avatar
Michael Yang committed
231
232
				},
			},
Michael Yang's avatar
Michael Yang committed
233
			`[INST] Hello friend![/INST] Hello human![INST] What is your name?[/INST] `,
Michael Yang's avatar
Michael Yang committed
234
235
		},
		{
Michael Yang's avatar
Michael Yang committed
236
237
			"mistral system",
			[]template{
238
239
240
241
242
243
				{"no response", `[INST] {{ if .System }}{{ .System }}

{{ end }}{{ .Prompt }}[/INST] `},
				{"response", `[INST] {{ if .System }}{{ .System }}

{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`},
244
				{"messages", `[INST] {{ if .System }}{{ .System }}
245

246
247
248
{{ end }}
{{- range .Messages }}
{{- if eq .Role "user" }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}[INST] {{ end }}
Michael Yang's avatar
Michael Yang committed
249
{{- end }}`},
Michael Yang's avatar
Michael Yang committed
250
251
252
253
254
255
			},
			Values{
				Messages: []api.Message{
					{Role: "system", Content: "You are a helpful assistant!"},
					{Role: "user", Content: "Hello friend!"},
					{Role: "assistant", Content: "Hello human!"},
Michael Yang's avatar
Michael Yang committed
256
					{Role: "user", Content: "What is your name?"},
Michael Yang's avatar
Michael Yang committed
257
258
				},
			},
259
			`[INST] You are a helpful assistant!
Michael Yang's avatar
Michael Yang committed
260

261
Hello friend![/INST] Hello human![INST] What is your name?[/INST] `,
Michael Yang's avatar
Michael Yang committed
262
263
		},
		{
Michael Yang's avatar
Michael Yang committed
264
265
266
267
			"chatml",
			[]template{
				// this does not have a "no response" test because it's impossible to render the same output
				{"response", `{{ if .System }}<|im_start|>system
Michael Yang's avatar
Michael Yang committed
268
269
270
271
272
{{ .System }}<|im_end|>
{{ end }}{{ if .Prompt }}<|im_start|>user
{{ .Prompt }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ .Response }}<|im_end|>
Michael Yang's avatar
Michael Yang committed
273
274
`},
				{"messages", `
275
276
277
{{- range $index, $_ := .Messages }}<|im_start|>{{ .Role }}
{{ .Content }}<|im_end|>
{{ end }}<|im_start|>assistant
Michael Yang's avatar
Michael Yang committed
278
`},
Michael Yang's avatar
Michael Yang committed
279
280
281
282
283
284
			},
			Values{
				Messages: []api.Message{
					{Role: "system", Content: "You are a helpful assistant!"},
					{Role: "user", Content: "Hello friend!"},
					{Role: "assistant", Content: "Hello human!"},
Michael Yang's avatar
Michael Yang committed
285
					{Role: "user", Content: "What is your name?"},
Michael Yang's avatar
Michael Yang committed
286
287
				},
			},
288
289
290
			`<|im_start|>system
You are a helpful assistant!<|im_end|>
<|im_start|>user
Michael Yang's avatar
Michael Yang committed
291
292
293
294
Hello friend!<|im_end|>
<|im_start|>assistant
Hello human!<|im_end|>
<|im_start|>user
Michael Yang's avatar
Michael Yang committed
295
What is your name?<|im_end|>
Michael Yang's avatar
Michael Yang committed
296
297
298
299
<|im_start|>assistant
`,
		},
		{
Michael Yang's avatar
Michael Yang committed
300
301
302
303
			"moondream",
			[]template{
				// this does not have a "no response" test because it's impossible to render the same output
				{"response", `{{ if .Prompt }}Question: {{ .Prompt }}
Michael Yang's avatar
Michael Yang committed
304
305
306

{{ end }}Answer: {{ .Response }}

Michael Yang's avatar
Michael Yang committed
307
308
`},
				{"messages", `
Michael Yang's avatar
Michael Yang committed
309
{{- range .Messages }}
310
311
312
313
314
{{- if eq .Role "user" }}Question: {{ .Content }}

{{ else if eq .Role "assistant" }}Answer: {{ .Content }}

{{ end }}
Michael Yang's avatar
Michael Yang committed
315
{{- end }}Answer: `},
Michael Yang's avatar
Michael Yang committed
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
			},
			Values{
				Messages: []api.Message{
					{Role: "user", Content: "What's in this image?", Images: []api.ImageData{[]byte("")}},
					{Role: "assistant", Content: "It's a hot dog."},
					{Role: "user", Content: "What's in _this_ image?"},
					{Role: "user", Images: []api.ImageData{[]byte("")}},
					{Role: "user", Content: "Is it a hot dog?"},
				},
			},
			`Question: [img-0] What's in this image?

Answer: It's a hot dog.

Question: What's in _this_ image?

[img-1]

Is it a hot dog?

Answer: `,
		},
	}

	for _, tt := range cases {
Michael Yang's avatar
Michael Yang committed
341
342
343
344
		t.Run(tt.name, func(t *testing.T) {
			for _, ttt := range tt.templates {
				t.Run(ttt.name, func(t *testing.T) {
					tmpl, err := Parse(ttt.template)
Michael Yang's avatar
Michael Yang committed
345
346
347
348
349
350
351
352
353
					if err != nil {
						t.Fatal(err)
					}

					var b bytes.Buffer
					if err := tmpl.Execute(&b, tt.values); err != nil {
						t.Fatal(err)
					}

354
355
					if diff := cmp.Diff(b.String(), tt.expected); diff != "" {
						t.Errorf("mismatch (-got +want):\n%s", diff)
Michael Yang's avatar
Michael Yang committed
356
357
358
359
360
361
					}
				})
			}
		})
	}
}
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396

func TestExecuteWithSuffix(t *testing.T) {
	tmpl, err := Parse(`{{- if .Suffix }}<PRE> {{ .Prompt }} <SUF>{{ .Suffix }} <MID>
{{- else }}{{ .Prompt }}
{{- end }}`)
	if err != nil {
		t.Fatal(err)
	}

	cases := []struct {
		name   string
		values Values
		expect string
	}{
		{
			"message", Values{Messages: []api.Message{{Role: "user", Content: "hello"}}}, "hello",
		},
		{
			"prompt suffix", Values{Prompt: "def add(", Suffix: "return x"}, "<PRE> def add( <SUF>return x <MID>",
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			var b bytes.Buffer
			if err := tmpl.Execute(&b, tt.values); err != nil {
				t.Fatal(err)
			}

			if diff := cmp.Diff(b.String(), tt.expect); diff != "" {
				t.Errorf("mismatch (-got +want):\n%s", diff)
			}
		})
	}
}