prompt_test.go 10.9 KB
Newer Older
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
	"bytes"
5
6
	"image"
	"image/png"
7
8
	"testing"

Michael Yang's avatar
Michael Yang committed
9
	"github.com/google/go-cmp/cmp"
Michael Yang's avatar
lint  
Michael Yang committed
10

11
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
12
	"github.com/ollama/ollama/template"
13
14
15
)

func TestChatPrompt(t *testing.T) {
Michael Yang's avatar
Michael Yang committed
16
	type expect struct {
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
		prompt        string
		images        [][]byte
		aspectRatioID int
		error         error
	}

	tmpl, err := template.Parse(`
{{- if .System }}{{ .System }} {{ end }}
{{- if .Prompt }}{{ .Prompt }} {{ end }}
{{- if .Response }}{{ .Response }} {{ end }}`)
	if err != nil {
		t.Fatal(err)
	}
	visionModel := Model{Template: tmpl, ProjectorPaths: []string{"vision"}}
	mllamaModel := Model{Template: tmpl, ProjectorPaths: []string{"vision"}, Config: ConfigV2{ModelFamilies: []string{"mllama"}}}

	createImg := func(width, height int) ([]byte, error) {
34
		img := image.NewRGBA(image.Rect(0, 0, width, height))
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
		var buf bytes.Buffer

		if err := png.Encode(&buf, img); err != nil {
			return nil, err
		}

		return buf.Bytes(), nil
	}

	imgBuf, err := createImg(5, 5)
	if err != nil {
		t.Fatal(err)
	}

	imgBuf2, err := createImg(6, 6)
	if err != nil {
		t.Fatal(err)
Michael Yang's avatar
Michael Yang committed
52
53
54
55
	}

	cases := []struct {
		name  string
56
		model Model
Michael Yang's avatar
Michael Yang committed
57
58
59
		limit int
		msgs  []api.Message
		expect
60
61
	}{
		{
Michael Yang's avatar
Michael Yang committed
62
			name:  "messages",
63
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
64
65
66
67
68
69
70
71
			limit: 64,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "You're a test, Harry! I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
72
73
74
			},
		},
		{
Michael Yang's avatar
Michael Yang committed
75
			name:  "truncate messages",
76
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
77
78
79
80
81
82
83
84
85
			limit: 1,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "A test. And a thumping good one at that, I'd wager. ",
			},
86
87
		},
		{
Michael Yang's avatar
Michael Yang committed
88
			name:  "truncate messages with image",
89
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
90
91
92
93
94
95
96
			limit: 64,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{[]byte("something")}},
			},
			expect: expect{
97
				prompt: "[img-0]A test. And a thumping good one at that, I'd wager. ",
Michael Yang's avatar
Michael Yang committed
98
99
100
101
				images: [][]byte{
					[]byte("something"),
				},
			},
102
103
		},
		{
Michael Yang's avatar
Michael Yang committed
104
			name:  "truncate messages with images",
105
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
106
107
108
109
110
111
112
			limit: 64,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!", Images: []api.ImageData{[]byte("something")}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{[]byte("somethingelse")}},
			},
			expect: expect{
113
				prompt: "[img-0]A test. And a thumping good one at that, I'd wager. ",
Michael Yang's avatar
Michael Yang committed
114
115
116
117
				images: [][]byte{
					[]byte("somethingelse"),
				},
			},
118
119
		},
		{
Michael Yang's avatar
Michael Yang committed
120
			name:  "messages with images",
121
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
122
123
124
125
126
127
128
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!", Images: []api.ImageData{[]byte("something")}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{[]byte("somethingelse")}},
			},
			expect: expect{
129
				prompt: "[img-0]You're a test, Harry! I-I'm a what? [img-1]A test. And a thumping good one at that, I'd wager. ",
Michael Yang's avatar
Michael Yang committed
130
131
132
133
134
				images: [][]byte{
					[]byte("something"),
					[]byte("somethingelse"),
				},
			},
135
136
		},
		{
Michael Yang's avatar
Michael Yang committed
137
			name:  "message with image tag",
138
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
139
140
141
142
143
144
145
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry! [img]", Images: []api.ImageData{[]byte("something")}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{[]byte("somethingelse")}},
			},
			expect: expect{
146
				prompt: "You're a test, Harry! [img-0] I-I'm a what? [img-1]A test. And a thumping good one at that, I'd wager. ",
Michael Yang's avatar
Michael Yang committed
147
148
149
150
151
				images: [][]byte{
					[]byte("something"),
					[]byte("somethingelse"),
				},
			},
152
153
		},
		{
Michael Yang's avatar
Michael Yang committed
154
			name:  "messages with interleaved images",
155
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "user", Images: []api.ImageData{[]byte("something")}},
				{Role: "user", Images: []api.ImageData{[]byte("somethingelse")}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "You're a test, Harry!\n\n[img-0]\n\n[img-1] I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
				images: [][]byte{
					[]byte("something"),
					[]byte("somethingelse"),
				},
			},
171
172
		},
		{
Michael Yang's avatar
Michael Yang committed
173
			name:  "truncate message with interleaved images",
174
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
175
176
177
178
179
180
181
182
183
184
185
186
187
188
			limit: 1024,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "user", Images: []api.ImageData{[]byte("something")}},
				{Role: "user", Images: []api.ImageData{[]byte("somethingelse")}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "[img-0] I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
				images: [][]byte{
					[]byte("somethingelse"),
				},
			},
189
190
		},
		{
Michael Yang's avatar
Michael Yang committed
191
			name:  "message with system prompt",
192
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
193
194
195
196
197
198
199
200
			limit: 2048,
			msgs: []api.Message{
				{Role: "system", Content: "You are the Test Who Lived."},
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
201
				prompt: "You are the Test Who Lived. You're a test, Harry! I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
202
203
			},
		},
Michael Yang's avatar
Michael Yang committed
204
205
		{
			name:  "out of order system",
206
			model: visionModel,
Michael Yang's avatar
Michael Yang committed
207
208
209
210
211
212
213
214
215
216
217
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "system", Content: "You are the Test Who Lived."},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "You're a test, Harry! I-I'm a what? You are the Test Who Lived. A test. And a thumping good one at that, I'd wager. ",
			},
		},
218
219
220
221
222
223
224
225
		{
			name:  "multiple images same prompt",
			model: visionModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "Compare these two pictures of hotdogs", Images: []api.ImageData{[]byte("one hotdog"), []byte("two hotdogs")}},
			},
			expect: expect{
226
				prompt: "[img-0][img-1]Compare these two pictures of hotdogs ",
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
				images: [][]byte{[]byte("one hotdog"), []byte("two hotdogs")},
			},
		},
		{
			name:  "messages with mllama (no images)",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager."},
			},
			expect: expect{
				prompt: "You're a test, Harry! I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
			},
		},
		{
			name:  "messages with mllama single prompt",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "How many hotdogs are in this image?", Images: []api.ImageData{imgBuf}},
			},
			expect: expect{
251
				prompt:        "[img-0]<|image|>How many hotdogs are in this image? ",
252
253
254
255
256
257
258
259
260
261
262
263
264
265
				images:        [][]byte{imgBuf},
				aspectRatioID: 1,
			},
		},
		{
			name:  "messages with mllama",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{imgBuf}},
			},
			expect: expect{
266
				prompt:        "You're a test, Harry! I-I'm a what? [img-0]<|image|>A test. And a thumping good one at that, I'd wager. ",
267
268
269
270
271
272
273
274
275
276
277
278
279
280
				images:        [][]byte{imgBuf},
				aspectRatioID: 1,
			},
		},
		{
			name:  "multiple messages with mllama",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!", Images: []api.ImageData{imgBuf}},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{imgBuf2}},
			},
			expect: expect{
281
282
				prompt:        "[img-0]<|image|>You're a test, Harry! I-I'm a what? [img-1]<|image|>A test. And a thumping good one at that, I'd wager. ",
				images:        [][]byte{imgBuf, imgBuf2},
283
284
285
286
287
288
289
290
291
292
293
294
295
				aspectRatioID: 1,
			},
		},
		{
			name:  "earlier image with mllama",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "How many hotdogs are in this image?", Images: []api.ImageData{imgBuf}},
				{Role: "assistant", Content: "There are four hotdogs."},
				{Role: "user", Content: "Which ones have mustard?"},
			},
			expect: expect{
296
				prompt:        "[img-0]<|image|>How many hotdogs are in this image? There are four hotdogs. Which ones have mustard? ",
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
				images:        [][]byte{imgBuf},
				aspectRatioID: 1,
			},
		},
		{
			name:  "too many images with mllama",
			model: mllamaModel,
			limit: 2048,
			msgs: []api.Message{
				{Role: "user", Content: "You're a test, Harry!"},
				{Role: "assistant", Content: "I-I'm a what?"},
				{Role: "user", Content: "A test. And a thumping good one at that, I'd wager.", Images: []api.ImageData{imgBuf, imgBuf}},
			},
			expect: expect{
				error: errTooManyImages,
			},
		},
314
315
	}

Michael Yang's avatar
Michael Yang committed
316
317
	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
318
			model := tt.model
Michael Yang's avatar
Michael Yang committed
319
			opts := api.Options{Runner: api.Runner{NumCtx: tt.limit}}
320
			prompt, images, err := chatPrompt(t.Context(), &model, mockRunner{}.Tokenize, &opts, tt.msgs, nil)
321
			if tt.error == nil && err != nil {
Michael Yang's avatar
Michael Yang committed
322
				t.Fatal(err)
323
324
			} else if tt.error != nil && err != tt.error {
				t.Fatalf("expected err '%q', got '%q'", tt.error, err)
Michael Yang's avatar
Michael Yang committed
325
326
			}

Michael Yang's avatar
Michael Yang committed
327
328
329
330
			if diff := cmp.Diff(prompt, tt.prompt); diff != "" {
				t.Errorf("mismatch (-got +want):\n%s", diff)
			}

Michael Yang's avatar
Michael Yang committed
331
332
333
334
335
336
337
338
339
			if len(images) != len(tt.images) {
				t.Fatalf("expected %d images, got %d", len(tt.images), len(images))
			}

			for i := range images {
				if images[i].ID != i {
					t.Errorf("expected ID %d, got %d", i, images[i].ID)
				}

340
341
342
343
344
345
346
347
				if len(model.Config.ModelFamilies) == 0 {
					if !bytes.Equal(images[i].Data, tt.images[i]) {
						t.Errorf("expected %q, got %q", tt.images[i], images[i].Data)
					}
				} else {
					if images[i].AspectRatioID != tt.aspectRatioID {
						t.Errorf("expected aspect ratio %d, got %d", tt.aspectRatioID, images[i].AspectRatioID)
					}
Michael Yang's avatar
Michael Yang committed
348
				}
349
350
351
352
			}
		})
	}
}