model.go 8.47 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
package model

import (
Jesse Gross's avatar
Jesse Gross committed
4
	"errors"
Michael Yang's avatar
Michael Yang committed
5
6
7
	"fmt"
	_ "image/jpeg"
	_ "image/png"
Michael Yang's avatar
Michael Yang committed
8
	"log/slog"
Michael Yang's avatar
Michael Yang committed
9
10
11
12
13
14
15
16
17
	"os"
	"reflect"
	"strconv"
	"strings"

	_ "golang.org/x/image/bmp"
	_ "golang.org/x/image/tiff"
	_ "golang.org/x/image/webp"

18
19
	"github.com/ollama/ollama/fs"
	fsggml "github.com/ollama/ollama/fs/ggml"
Jesse Gross's avatar
Jesse Gross committed
20
	"github.com/ollama/ollama/kvcache"
21
	"github.com/ollama/ollama/logutil"
Michael Yang's avatar
Michael Yang committed
22
23
	"github.com/ollama/ollama/ml"
	_ "github.com/ollama/ollama/ml/backend"
Michael Yang's avatar
Michael Yang committed
24
	"github.com/ollama/ollama/ml/nn/pooling"
25
	"github.com/ollama/ollama/model/input"
Michael Yang's avatar
Michael Yang committed
26
27
)

28
29
30
31
32
var (
	ErrNoVisionModel        = errors.New("this model is missing data required for image input")
	ErrUnsupportedModel     = errors.New("model not supported")
	ErrUnsupportedTokenizer = errors.New("tokenizer not supported")
)
33

34
// Model implements a specific model architecture, defining the forward pass and any model-specific configuration
Michael Yang's avatar
Michael Yang committed
35
type Model interface {
Jesse Gross's avatar
Jesse Gross committed
36
	Forward(ml.Context, input.Batch) (ml.Tensor, error)
Michael Yang's avatar
Michael Yang committed
37
38

	Backend() ml.Backend
Jesse Gross's avatar
Jesse Gross committed
39
	Config() config
Michael Yang's avatar
Michael Yang committed
40
41
}

42
43
44
45
46
// MultimodalProcessor must be implemented by multimodal models.
type MultimodalProcessor interface {
	// EncodeMultimodal processes a single input (such as an image) and
	// generates an output (typically an embedding) that can be used by the model.
	//
47
48
49
50
	// The return value is one or more tensors, each with optional model-specific
	// opaque metadata. Typically, the tensors might be views into an embedding
	// with each view representing a chunk of data that can be processed independently
	// in different batches.
51
52
	//
	// The result may be cached by the runner.
53
	EncodeMultimodal(ml.Context, []byte) ([]input.Multimodal, error)
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

	// PostTokenize is called after tokenization to allow the model to edit the
	// input stream to correctly arrange multimodal elements.
	//
	// The input is a slice of tokens with the results of EncodeMultimodal interleaved
	// in the order that the user provided them. Each element of the slice will be
	// either a single token or single multimodal object.
	//
	// The model must ensure that inputs are stored according to how they will be
	// processed and stored in the cache. For example, Llava-style models should insert
	// placeholder tokens equal to the feature size of the corresponding image with
	// the image itself attached to and split across these tokens. When Forward is called
	// a partial subset of these tokens may be submitted according to the batch size.
	//
	// This function is also responsible for updating MultimodalHash for any Multimodal
	// that is modified to ensure that there is a unique hash value that accurately
	// represents the contents.
71
	PostTokenize([]*input.Input) ([]*input.Input, error)
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
}

// Base implements the common fields and methods for all models
type Base struct {
	b ml.Backend
	config
}

type config struct {
	Cache kvcache.Cache
}

// Backend returns the underlying backend that will run the model
func (m *Base) Backend() ml.Backend {
	return m.b
}

func (m *Base) Config() config {
	return m.config
91
92
}

93
var models = make(map[string]func(fs.Config) (Model, error))
Michael Yang's avatar
Michael Yang committed
94

95
// Register registers a model constructor for the given architecture
96
func Register(name string, f func(fs.Config) (Model, error)) {
Michael Yang's avatar
Michael Yang committed
97
98
99
100
101
102
103
	if _, ok := models[name]; ok {
		panic("model: model already registered")
	}

	models[name] = f
}

104
// New initializes a new model instance with the provided configuration based on the metadata in the model file
105
106
func New(modelPath string, params ml.BackendParams) (Model, error) {
	b, err := ml.NewBackend(modelPath, params)
Michael Yang's avatar
Michael Yang committed
107
108
109
110
	if err != nil {
		return nil, err
	}

111
	m, err := modelForArch(b.Config())
Michael Yang's avatar
Michael Yang committed
112
113
114
115
	if err != nil {
		return nil, err
	}

Jesse Gross's avatar
Jesse Gross committed
116
	base := Base{b: b, config: m.Config()}
Michael Yang's avatar
Michael Yang committed
117
	v := reflect.ValueOf(m)
Jesse Gross's avatar
Jesse Gross committed
118
	v.Elem().Set(populateFields(base, v.Elem()))
Michael Yang's avatar
Michael Yang committed
119
120
121
	return m, nil
}

122
123
124
125
126
127
func NewTextProcessor(s string) (TextProcessor, error) {
	r, err := os.Open(s)
	if err != nil {
		return nil, err
	}
	defer r.Close()
128

129
	meta, err := fsggml.Decode(r, -1)
130
131
132
133
	if err != nil {
		return nil, err
	}

134
	m, err := modelForArch(meta.KV())
135
136
137
	if err != nil {
		return nil, err
	}
138

139
140
	tp, ok := m.(TextProcessor)
	if !ok {
141
		return nil, ErrUnsupportedTokenizer
142
143
144
145
	}
	return tp, nil
}

146
147
148
149
150
151
152
153
154
155
156
157
158
159
func modelForArch(c fs.Config) (Model, error) {
	arch := c.Architecture()
	if pooling.Type(c.Uint("pooling_type")) != pooling.TypeNone {
		arch = arch + "_embed"
	}

	f, ok := models[arch]
	if !ok {
		return nil, ErrUnsupportedModel
	}

	return f(c)
}

Jesse Gross's avatar
Jesse Gross committed
160
func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
Michael Yang's avatar
Michael Yang committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
	t := v.Type()

	if t.Kind() == reflect.Struct {
		allNil := true
		for i := range t.NumField() {
			tt := t.Field(i).Type
			vv := v.Field(i)
			if !vv.CanSet() {
				continue
			}

			// make a copy
			tagsCopy := tags
			if tag := t.Field(i).Tag.Get("gguf"); tag != "" {
Michael Yang's avatar
Michael Yang committed
175
				tagsCopy = append(tagsCopy, parseTag(tag))
Michael Yang's avatar
Michael Yang committed
176
177
178
			}

			if tt == reflect.TypeOf((*Base)(nil)).Elem() {
Jesse Gross's avatar
Jesse Gross committed
179
				vv.Set(reflect.ValueOf(base))
Michael Yang's avatar
Michael Yang committed
180
			} else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() {
Michael Yang's avatar
Michael Yang committed
181
182
				var fn func([]Tag, string, string) [][]string
				fn = func(tags []Tag, prefix, suffix string) (fullNames [][]string) {
183
					if len(tags) > 0 {
Michael Yang's avatar
Michael Yang committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
						var names []string
						if tags[0].name != "" {
							for _, n := range append([]string{tags[0].name}, tags[0].alternatives...) {
								names = append(names, prefix+n+suffix)
							}
						}

						if childNames := fn(tags[1:], tags[0].prefix, tags[0].suffix); len(childNames) == 0 {
							// no child names, append current names
							fullNames = append(fullNames, names)
						} else if len(names) == 0 {
							// no current names, append child names
							fullNames = append(fullNames, childNames...)
						} else {
							// combine current and child names
							for _, name := range names {
								for _, childName := range childNames {
									fullNames = append(fullNames, append([]string{name}, childName...))
202
203
								}
							}
Michael Yang's avatar
Michael Yang committed
204
205
206
						}
					}

Michael Yang's avatar
Michael Yang committed
207
					return fullNames
Michael Yang's avatar
Michael Yang committed
208
209
				}

Michael Yang's avatar
Michael Yang committed
210
				names := fn(tagsCopy, "", "")
Michael Yang's avatar
Michael Yang committed
211
				for _, name := range names {
Jesse Gross's avatar
Jesse Gross committed
212
					if tensor := base.Backend().Get(strings.Join(name, ".")); tensor != nil {
213
						logutil.Trace("found tensor", "", tensor)
Michael Yang's avatar
Michael Yang committed
214
215
216
217
						vv.Set(reflect.ValueOf(tensor))
						break
					}
				}
218
			} else if tt.Kind() == reflect.Pointer || tt.Kind() == reflect.Interface {
Jesse Gross's avatar
Jesse Gross committed
219
				setPointer(base, vv, tagsCopy)
Michael Yang's avatar
Michael Yang committed
220
221
			} else if tt.Kind() == reflect.Slice || tt.Kind() == reflect.Array {
				for i := range vv.Len() {
222
223
					vvv := vv.Index(i)
					if vvv.Kind() == reflect.Pointer || vvv.Kind() == reflect.Interface {
Michael Yang's avatar
Michael Yang committed
224
						setPointer(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)}))
225
					} else {
Michael Yang's avatar
Michael Yang committed
226
						vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)})...))
227
					}
Michael Yang's avatar
Michael Yang committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
				}
			}

			if !canNil(tt) || !vv.IsNil() {
				allNil = false
			}
		}

		if allNil {
			return reflect.Zero(t)
		}
	}

	return v
}

Jesse Gross's avatar
Jesse Gross committed
244
func setPointer(base Base, v reflect.Value, tags []Tag) {
245
246
247
248
249
250
251
252
253
	vv := v
	if v.Kind() == reflect.Interface {
		if v.IsNil() {
			return
		}

		vv = vv.Elem()
	}

254
	vv = reflect.Indirect(vv)
255
256
257
258
	if v.IsNil() {
		vv = reflect.New(v.Type().Elem()).Elem()
	}

Jesse Gross's avatar
Jesse Gross committed
259
	if f := populateFields(base, vv, tags...); f.CanAddr() {
260
261
262
263
		v.Set(f.Addr())
	}
}

Michael Yang's avatar
Michael Yang committed
264
type Tag struct {
Michael Yang's avatar
Michael Yang committed
265
266
267
268
269
	name,
	// prefix and suffix are applied to child tags
	prefix,
	suffix string
	alternatives []string
Michael Yang's avatar
Michael Yang committed
270
271
}

Michael Yang's avatar
Michael Yang committed
272
func parseTag(s string) (tag Tag) {
Michael Yang's avatar
Michael Yang committed
273
274
	parts := strings.Split(s, ",")
	if len(parts) > 0 {
Michael Yang's avatar
Michael Yang committed
275
		tag.name = parts[0]
Michael Yang's avatar
Michael Yang committed
276
277

		for _, part := range parts[1:] {
Michael Yang's avatar
Michael Yang committed
278
279
280
281
282
283
284
285
286
287
288
289
			if value, ok := strings.CutPrefix(part, "alt:"); ok && tag.name == "" {
				// elevate alternative to primary if no primary given
				tag.name = value
				slog.Warn("gguf tag has alt: but no primary name", "tag", s)
			} else if ok {
				tag.alternatives = append(tag.alternatives, value)
			}
			if value, ok := strings.CutPrefix(part, "pre:"); ok {
				tag.prefix = value
			}
			if value, ok := strings.CutPrefix(part, "suf:"); ok {
				tag.suffix = value
Michael Yang's avatar
Michael Yang committed
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
			}
		}
	}

	return
}

func canNil(t reflect.Type) bool {
	return t.Kind() == reflect.Chan ||
		t.Kind() == reflect.Func ||
		t.Kind() == reflect.Interface ||
		t.Kind() == reflect.Map ||
		t.Kind() == reflect.Pointer ||
		t.Kind() == reflect.Slice
}

306
func Forward(ctx ml.Context, m Model, batch input.Batch) (ml.Tensor, error) {
Jesse Gross's avatar
Jesse Gross committed
307
308
	if len(batch.Positions) != len(batch.Sequences) {
		return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(batch.Positions), len(batch.Sequences))
Jesse Gross's avatar
Jesse Gross committed
309
310
	}

Jesse Gross's avatar
Jesse Gross committed
311
	if len(batch.Positions) < 1 {
Jesse Gross's avatar
Jesse Gross committed
312
313
314
315
316
		return nil, errors.New("batch size cannot be less than 1")
	}

	cache := m.Config().Cache
	if cache != nil {
317
		err := cache.StartForward(ctx, batch, false)
Jesse Gross's avatar
Jesse Gross committed
318
319
320
		if err != nil {
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
321
322
	}

Jesse Gross's avatar
Jesse Gross committed
323
	t, err := m.Forward(ctx, batch)
Michael Yang's avatar
Michael Yang committed
324
325
326
327
	if err != nil {
		return nil, err
	}

328
	ctx.Forward(t)
329
330

	return t, nil
Michael Yang's avatar
Michael Yang committed
331
}