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

import (
4
	"context"
Jesse Gross's avatar
Jesse Gross committed
5
	"errors"
Michael Yang's avatar
Michael Yang committed
6
7
8
9
10
11
12
13
14
15
16
17
18
	"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"

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

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

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

	Backend() ml.Backend
Jesse Gross's avatar
Jesse Gross committed
35
	Config() config
Michael Yang's avatar
Michael Yang committed
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
63
64
65
// 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.
66
	PostTokenize([]input.Input) ([]input.Input, error)
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
}

// 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
86
87
}

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

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

	models[name] = f
}

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

107
	b, err := ml.NewBackend(ctx, r, params)
Michael Yang's avatar
Michael Yang committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
	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
123
124
	base := Base{b: b, config: m.Config()}

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

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

143
func getTextProcessor(kv fsggml.KV) (TextProcessor, error) {
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
	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
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
175
176
177
178
	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
179
				vv.Set(reflect.ValueOf(base))
Michael Yang's avatar
Michael Yang committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
			} 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
205
					if tensor := base.Backend().Get(strings.Join(name, ".")); tensor != nil {
206
						slog.Log(context.TODO(), logutil.LevelTrace, "found tensor", "", tensor)
Michael Yang's avatar
Michael Yang committed
207
208
209
210
						vv.Set(reflect.ValueOf(tensor))
						break
					}
				}
211
			} else if tt.Kind() == reflect.Pointer || tt.Kind() == reflect.Interface {
Jesse Gross's avatar
Jesse Gross committed
212
				setPointer(base, vv, tagsCopy)
Michael Yang's avatar
Michael Yang committed
213
214
			} else if tt.Kind() == reflect.Slice || tt.Kind() == reflect.Array {
				for i := range vv.Len() {
215
216
					vvv := vv.Index(i)
					if vvv.Kind() == reflect.Pointer || vvv.Kind() == reflect.Interface {
Jesse Gross's avatar
Jesse Gross committed
217
						setPointer(base, vvv, append(tagsCopy, Tag{Name: strconv.Itoa(i)}))
218
					} else {
Jesse Gross's avatar
Jesse Gross committed
219
						vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{Name: strconv.Itoa(i)})...))
220
					}
Michael Yang's avatar
Michael Yang committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
				}
			}

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

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

	return v
}

Jesse Gross's avatar
Jesse Gross committed
237
func setPointer(base Base, v reflect.Value, tags []Tag) {
238
239
240
241
242
243
244
245
246
247
248
249
250
251
	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
252
	if f := populateFields(base, vv, tags...); f.CanAddr() {
253
254
255
256
		v.Set(f.Addr())
	}
}

Michael Yang's avatar
Michael Yang committed
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
283
284
285
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
}

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

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

295
296
297
298
299
300
	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
301
302
	cache := m.Config().Cache
	if cache != nil {
303
		err := cache.StartForward(ctx, batch, false)
Jesse Gross's avatar
Jesse Gross committed
304
305
306
		if err != nil {
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
307
308
	}

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

314
	ctx.Forward(t).Compute(t)
315
316

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