model.go 7.56 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
8
9
10
11
12
13
14
15
16
17
	"fmt"
	_ "image/jpeg"
	_ "image/png"
	"log/slog"
	"os"
	"reflect"
	"strconv"
	"strings"

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

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

Michael Yang's avatar
Michael Yang committed
25
var ErrNoVisionModel = errors.New("this model is missing data required for image input")
26

27
// Model implements a specific model architecture, defining the forward pass and any model-specific configuration
Michael Yang's avatar
Michael Yang committed
28
type Model interface {
Jesse Gross's avatar
Jesse Gross committed
29
	Forward(ml.Context, input.Batch) (ml.Tensor, error)
Michael Yang's avatar
Michael Yang committed
30
31

	Backend() ml.Backend
Jesse Gross's avatar
Jesse Gross committed
32
	Config() config
Michael Yang's avatar
Michael Yang committed
33
34
}

35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
// 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.
	//
	// The return value is most typically an ml.Tensor, however, different
	// type are possible, such as an object containing a tensor plus
	// additional metadata, a slice of tensors or even just the original input.
	//
	// The result may be cached by the runner.
	EncodeMultimodal(ml.Context, []byte) (any, error)

	// 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.
63
	PostTokenize([]input.Input) ([]input.Input, error)
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
}

// 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
83
84
}

Michael Yang's avatar
Michael Yang committed
85
86
var models = make(map[string]func(ml.Config) (Model, error))

87
// Register registers a model constructor for the given architecture
Michael Yang's avatar
Michael Yang committed
88
89
90
91
92
93
94
95
func Register(name string, f func(ml.Config) (Model, error)) {
	if _, ok := models[name]; ok {
		panic("model: model already registered")
	}

	models[name] = f
}

96
// New initializes a new model instance with the provided configuration based on the metadata in the model file
97
func New(modelPath string, params ml.BackendParams) (Model, error) {
98
	r, err := os.Open(modelPath)
Michael Yang's avatar
Michael Yang committed
99
100
101
102
103
	if err != nil {
		return nil, err
	}
	defer r.Close()

104
	b, err := ml.NewBackend(r, params)
Michael Yang's avatar
Michael Yang committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
	if err != nil {
		return nil, err
	}

	arch := b.Config().Architecture()
	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
120
121
	base := Base{b: b, config: m.Config()}

Michael Yang's avatar
Michael Yang committed
122
	v := reflect.ValueOf(m)
Jesse Gross's avatar
Jesse Gross committed
123
	v.Elem().Set(populateFields(base, v.Elem()))
Michael Yang's avatar
Michael Yang committed
124
125
126
	return m, nil
}

127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
func NewTextProcessor(s string) (TextProcessor, error) {
	r, err := os.Open(s)
	if err != nil {
		return nil, err
	}
	defer r.Close()
	meta, _, err := fs.Decode(r, -1)
	if err != nil {
		return nil, err
	}
	return getTextProcessor(meta.KV())
}

func getTextProcessor(kv fs.KV) (TextProcessor, error) {
	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
157
func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
Michael Yang's avatar
Michael Yang committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
	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
176
				vv.Set(reflect.ValueOf(base))
Michael Yang's avatar
Michael Yang committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
			} else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() {
				var fn func([]Tag) [][]string
				fn = func(tags []Tag) (values [][]string) {
					if len(tags) < 1 {
						return nil
					}

					values = [][]string{{tags[0].Name}}
					for _, alt := range tags[0].Alternate {
						values = append(values, []string{alt})
					}

					for i, value := range values {
						for _, rest := range fn(tags[1:]) {
							value = append(value, rest...)
						}

						values[i] = value
					}

					return values
				}

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

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

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

	return v
}

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

		vv = vv.Elem()
	}

	vv = vv.Elem()
	if v.IsNil() {
		vv = reflect.New(v.Type().Elem()).Elem()
	}

Jesse Gross's avatar
Jesse Gross committed
249
	if f := populateFields(base, vv, tags...); f.CanAddr() {
250
251
252
253
		v.Set(f.Addr())
	}
}

Michael Yang's avatar
Michael Yang committed
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
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
}

283
func Forward(ctx ml.Context, m Model, inputs []int32, batch input.Batch) (ml.Tensor, error) {
Jesse Gross's avatar
Jesse Gross committed
284
285
	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
286
287
	}

Jesse Gross's avatar
Jesse Gross committed
288
	if len(batch.Positions) < 1 {
Jesse Gross's avatar
Jesse Gross committed
289
290
291
		return nil, errors.New("batch size cannot be less than 1")
	}

292
293
294
295
296
297
	var err error
	batch.Inputs, err = ctx.Input().FromIntSlice(inputs, len(inputs))
	if err != nil {
		return nil, err
	}

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

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

311
	ctx.Forward(t).Compute(t)
312
313

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