model.go 7.83 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
	"math"
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"
24
	"github.com/ollama/ollama/model/input"
Michael Yang's avatar
Michael Yang committed
25
26
)

27
28
29
30
31
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")
)
32

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

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

41
42
43
44
45
// 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.
	//
46
47
48
49
	// 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.
50
51
	//
	// The result may be cached by the runner.
52
	EncodeMultimodal(ml.Context, []byte) ([]input.Multimodal, error)
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

	// 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.
70
	PostTokenize([]*input.Input) ([]*input.Input, error)
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
}

// 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
90
91
}

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

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

	models[name] = f
}

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

	arch := b.Config().Architecture()
Michael Yang's avatar
Michael Yang committed
111
112
113
114
	if b.Config().Uint("pooling_type", math.MaxUint32) != math.MaxUint32 {
		arch = arch + "_embed"
	}

Michael Yang's avatar
Michael Yang committed
115
116
117
118
119
120
121
122
123
124
	f, ok := models[arch]
	if !ok {
		return nil, fmt.Errorf("unsupported model architecture %q", arch)
	}

	m, err := f(b.Config())
	if err != nil {
		return nil, err
	}

Jesse Gross's avatar
Jesse Gross committed
125
126
	base := Base{b: b, config: m.Config()}

Michael Yang's avatar
Michael Yang committed
127
	v := reflect.ValueOf(m)
Jesse Gross's avatar
Jesse Gross committed
128
	v.Elem().Set(populateFields(base, v.Elem()))
Michael Yang's avatar
Michael Yang committed
129
130
131
	return m, nil
}

132
133
134
135
136
137
func NewTextProcessor(s string) (TextProcessor, error) {
	r, err := os.Open(s)
	if err != nil {
		return nil, err
	}
	defer r.Close()
138
	meta, err := fsggml.Decode(r, -1)
139
140
141
142
143
144
	if err != nil {
		return nil, err
	}
	return getTextProcessor(meta.KV())
}

145
func getTextProcessor(kv fsggml.KV) (TextProcessor, error) {
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
	arch := kv.Architecture()
	f, ok := models[arch]
	if !ok {
		return nil, fmt.Errorf("unsupported model architecture %q", arch)
	}
	m, err := f(kv)
	if err != nil {
		return nil, err
	}
	tp, ok := m.(TextProcessor)
	if !ok {
		return nil, fmt.Errorf("%v is not a TextProcessor", m)
	}
	return tp, nil
}

Jesse Gross's avatar
Jesse Gross committed
162
func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
Michael Yang's avatar
Michael Yang committed
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
	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 != "" {
				tagsCopy = append(tagsCopy, ParseTags(tag))
			}

			if tt == reflect.TypeOf((*Base)(nil)).Elem() {
Jesse Gross's avatar
Jesse Gross committed
181
				vv.Set(reflect.ValueOf(base))
Michael Yang's avatar
Michael Yang committed
182
183
			} else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() {
				var fn func([]Tag) [][]string
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
				fn = func(tags []Tag) (names [][]string) {
					if len(tags) > 0 {
						localNames := []string{tags[0].Name}
						localNames = append(localNames, tags[0].Alternate...)

						for _, localName := range localNames {
							fullName := []string{localName}
							nested := fn(tags[1:])
							if len(nested) > 0 {
								for _, rest := range nested {
									names = append(names, append(fullName, rest...))
								}
							} else {
								names = append(names, fullName)
							}
Michael Yang's avatar
Michael Yang committed
199
200
201
						}
					}

202
					return names
Michael Yang's avatar
Michael Yang committed
203
204
205
206
				}

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

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

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

	return v
}

Jesse Gross's avatar
Jesse Gross committed
239
func setPointer(base Base, v reflect.Value, tags []Tag) {
240
241
242
243
244
245
246
247
248
	vv := v
	if v.Kind() == reflect.Interface {
		if v.IsNil() {
			return
		}

		vv = vv.Elem()
	}

249
	vv = reflect.Indirect(vv)
250
251
252
253
	if v.IsNil() {
		vv = reflect.New(v.Type().Elem()).Elem()
	}

Jesse Gross's avatar
Jesse Gross committed
254
	if f := populateFields(base, vv, tags...); f.CanAddr() {
255
256
257
258
		v.Set(f.Addr())
	}
}

Michael Yang's avatar
Michael Yang committed
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
type Tag struct {
	Name      string
	Alternate []string
}

func ParseTags(s string) (tag Tag) {
	parts := strings.Split(s, ",")
	if len(parts) > 0 {
		tag.Name = parts[0]

		for _, part := range parts[1:] {
			if value, ok := strings.CutPrefix(part, "alt:"); ok {
				tag.Alternate = append(tag.Alternate, value)
			}
		}
	}

	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
}

288
func Forward(ctx ml.Context, m Model, batch input.Batch) (ml.Tensor, error) {
Jesse Gross's avatar
Jesse Gross committed
289
290
	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
291
292
	}

Jesse Gross's avatar
Jesse Gross committed
293
	if len(batch.Positions) < 1 {
Jesse Gross's avatar
Jesse Gross committed
294
295
296
297
298
		return nil, errors.New("batch size cannot be less than 1")
	}

	cache := m.Config().Cache
	if cache != nil {
299
		err := cache.StartForward(ctx, batch, false)
Jesse Gross's avatar
Jesse Gross committed
300
301
302
		if err != nil {
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
303
304
	}

Jesse Gross's avatar
Jesse Gross committed
305
	t, err := m.Forward(ctx, batch)
Michael Yang's avatar
Michael Yang committed
306
307
308
309
	if err != nil {
		return nil, err
	}

310
	ctx.Forward(t)
311
312

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