template_test.go 10.7 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
lint  
Michael Yang committed
15

Michael Yang's avatar
Michael Yang committed
16
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
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
51
	"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
52
				tmpl, err := Parse(b.String())
Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
58
59
60
61
62
63
64
				if err != nil {
					t.Fatal(err)
				}

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

Michael Yang's avatar
Michael Yang committed
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
108
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 {
109
				var actual bytes.Buffer
Michael Yang's avatar
Michael Yang committed
110
111
112
113
114
115
116
117
118
119
				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)
					}

120
121
122
123
124
125
126
127
					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
128
129
130
						t.Errorf("mismatch (-got +want):\n%s", diff)
					}
				})
131
132

				t.Run("legacy", func(t *testing.T) {
133
					t.Skip("legacy outputs are currently default outputs")
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
					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
151
152
153
154
155
			}
		})
	}
}

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

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

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

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

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

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

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

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

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

247
248
249
{{ end }}
{{- range .Messages }}
{{- if eq .Role "user" }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}[INST] {{ end }}
Michael Yang's avatar
Michael Yang committed
250
{{- end }}`},
Michael Yang's avatar
Michael Yang committed
251
252
253
254
255
256
			},
			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
257
					{Role: "user", Content: "What is your name?"},
Michael Yang's avatar
Michael Yang committed
258
259
				},
			},
260
			`[INST] You are a helpful assistant!
Michael Yang's avatar
Michael Yang committed
261

262
Hello friend![/INST] Hello human![INST] What is your name?[/INST] `,
Michael Yang's avatar
Michael Yang committed
263
		},
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
		{
			"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`,
		},
Michael Yang's avatar
Michael Yang committed
284
		{
Michael Yang's avatar
Michael Yang committed
285
286
287
288
			"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
289
290
291
292
293
{{ .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
294
295
`},
				{"messages", `
296
297
298
{{- range $index, $_ := .Messages }}<|im_start|>{{ .Role }}
{{ .Content }}<|im_end|>
{{ end }}<|im_start|>assistant
Michael Yang's avatar
Michael Yang committed
299
`},
Michael Yang's avatar
Michael Yang committed
300
301
302
303
304
305
			},
			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
306
					{Role: "user", Content: "What is your name?"},
Michael Yang's avatar
Michael Yang committed
307
308
				},
			},
309
310
311
			`<|im_start|>system
You are a helpful assistant!<|im_end|>
<|im_start|>user
Michael Yang's avatar
Michael Yang committed
312
313
314
315
Hello friend!<|im_end|>
<|im_start|>assistant
Hello human!<|im_end|>
<|im_start|>user
Michael Yang's avatar
Michael Yang committed
316
What is your name?<|im_end|>
Michael Yang's avatar
Michael Yang committed
317
318
319
320
<|im_start|>assistant
`,
		},
		{
Michael Yang's avatar
Michael Yang committed
321
322
323
324
			"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
325
326
327

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

Michael Yang's avatar
Michael Yang committed
328
329
`},
				{"messages", `
Michael Yang's avatar
Michael Yang committed
330
{{- range .Messages }}
331
332
333
334
335
{{- if eq .Role "user" }}Question: {{ .Content }}

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

{{ end }}
Michael Yang's avatar
Michael Yang committed
336
{{- end }}Answer: `},
Michael Yang's avatar
Michael Yang committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
			},
			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
362
363
364
365
		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
366
367
368
369
370
371
372
373
374
					if err != nil {
						t.Fatal(err)
					}

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

375
376
					if diff := cmp.Diff(b.String(), tt.expected); diff != "" {
						t.Errorf("mismatch (-got +want):\n%s", diff)
Michael Yang's avatar
Michael Yang committed
377
378
379
380
381
382
					}
				})
			}
		})
	}
}
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417

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)
			}
		})
	}
}