"src/diffusers/models/unets/unet_3d_blocks.py" did not exist on "b62d9a1fdc0910bb864340b7bc29e86f6aa31d47"
model.go 8.07 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
package server

import (
	"archive/zip"
	"bytes"
	"context"
Michael Yang's avatar
tools  
Michael Yang committed
7
	"encoding/json"
Michael Yang's avatar
Michael Yang committed
8
9
10
	"errors"
	"fmt"
	"io"
11
	"log/slog"
Michael Yang's avatar
Michael Yang committed
12
13
14
	"net/http"
	"os"
	"path/filepath"
Michael Yang's avatar
tools  
Michael Yang committed
15
16
17
	"slices"
	"strings"
	"text/template/parse"
Michael Yang's avatar
Michael Yang committed
18
19
20
21

	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/convert"
	"github.com/ollama/ollama/llm"
Michael Yang's avatar
Michael Yang committed
22
	"github.com/ollama/ollama/template"
Michael Yang's avatar
Michael Yang committed
23
24
25
	"github.com/ollama/ollama/types/model"
)

Michael Yang's avatar
Michael Yang committed
26
var intermediateBlobs map[string]string = make(map[string]string)
27

28
type layerGGML struct {
Michael Yang's avatar
Michael Yang committed
29
30
31
32
	*Layer
	*llm.GGML
}

33
func parseFromModel(ctx context.Context, name model.Name, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
34
	m, err := ParseNamedManifest(name)
Michael Yang's avatar
Michael Yang committed
35
36
	switch {
	case errors.Is(err, os.ErrNotExist):
Michael Yang's avatar
Michael Yang committed
37
		if err := PullModel(ctx, name.String(), &registryOptions{}, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
38
39
40
			return nil, err
		}

41
		m, err = ParseNamedManifest(name)
Michael Yang's avatar
Michael Yang committed
42
43
44
		if err != nil {
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
45
46
47
48
	case err != nil:
		return nil, err
	}

49
50
	for _, layer := range m.Layers {
		layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, name.DisplayShortest())
Michael Yang's avatar
Michael Yang committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
		if err != nil {
			return nil, err
		}

		switch layer.MediaType {
		case "application/vnd.ollama.image.model",
			"application/vnd.ollama.image.projector",
			"application/vnd.ollama.image.adapter":
			blobpath, err := GetBlobsPath(layer.Digest)
			if err != nil {
				return nil, err
			}

			blob, err := os.Open(blobpath)
			if err != nil {
				return nil, err
			}
			defer blob.Close()

70
			ggml, _, err := llm.DecodeGGML(blob, 0)
Michael Yang's avatar
Michael Yang committed
71
72
73
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
74

75
			layers = append(layers, &layerGGML{layer, ggml})
Michael Yang's avatar
Michael Yang committed
76
		default:
77
			layers = append(layers, &layerGGML{layer, nil})
Michael Yang's avatar
Michael Yang committed
78
79
80
81
82
83
		}
	}

	return layers, nil
}

84
85
func parseFromZipFile(_ context.Context, f *os.File, digest string, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
	fi, err := f.Stat()
Michael Yang's avatar
Michael Yang committed
86
	if err != nil {
87
		return nil, err
Michael Yang's avatar
Michael Yang committed
88
89
	}

90
	r, err := zip.NewReader(f, fi.Size())
91
92
93
94
	if err != nil {
		return nil, err
	}

95
96
	p, err := os.MkdirTemp(filepath.Dir(f.Name()), "")
	if err != nil {
97
98
		return nil, err
	}
99
	defer os.RemoveAll(p)
100

Michael Yang's avatar
Michael Yang committed
101
102
103
	fn(api.ProgressResponse{Status: "converting model"})
	// TODO(mxyng): this should write directly into a layer
	// e.g. NewLayer(arch.Reader(), "application/vnd.ollama.image.model")
104
	t, err := os.CreateTemp(p, "fp16")
Michael Yang's avatar
Michael Yang committed
105
106
107
	if err != nil {
		return nil, err
	}
108
109
	defer t.Close()
	defer os.Remove(t.Name())
Michael Yang's avatar
Michael Yang committed
110

111
112
	fn(api.ProgressResponse{Status: "converting model"})
	if err := convert.Convert(convert.NewZipReader(r, p, 32<<20), t); err != nil {
Michael Yang's avatar
Michael Yang committed
113
114
115
		return nil, err
	}

116
	if _, err := t.Seek(0, io.SeekStart); err != nil {
Michael Yang's avatar
Michael Yang committed
117
118
119
		return nil, err
	}

120
	layer, err := NewLayer(t, "application/vnd.ollama.image.model")
Michael Yang's avatar
Michael Yang committed
121
	if err != nil {
Michael Yang's avatar
Michael Yang committed
122
		return nil, err
Michael Yang's avatar
Michael Yang committed
123
124
	}

125
	bin, err := layer.Open()
Michael Yang's avatar
Michael Yang committed
126
127
128
129
130
	if err != nil {
		return nil, err
	}
	defer bin.Close()

131
	ggml, _, err := llm.DecodeGGML(bin, 0)
Michael Yang's avatar
Michael Yang committed
132
133
134
135
	if err != nil {
		return nil, err
	}

136
	layers = append(layers, &layerGGML{layer, ggml})
137

Michael Yang's avatar
Michael Yang committed
138
	intermediateBlobs[digest] = layer.Digest
139
	return detectChatTemplate(layers)
Michael Yang's avatar
Michael Yang committed
140
141
}

142
func parseFromFile(ctx context.Context, file *os.File, digest string, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
Michael Yang's avatar
Michael Yang committed
143
144
145
146
147
148
149
150
151
152
	sr := io.NewSectionReader(file, 0, 512)
	contentType, err := detectContentType(sr)
	if err != nil {
		return nil, err
	}

	switch contentType {
	case "gguf", "ggla":
		// noop
	case "application/zip":
153
		return parseFromZipFile(ctx, file, digest, fn)
Michael Yang's avatar
Michael Yang committed
154
155
156
157
158
159
160
161
162
163
164
	default:
		return nil, fmt.Errorf("unsupported content type: %s", contentType)
	}

	stat, err := file.Stat()
	if err != nil {
		return nil, err
	}

	var offset int64
	for offset < stat.Size() {
165
		ggml, n, err := llm.DecodeGGML(file, 0)
Michael Yang's avatar
Michael Yang committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return nil, err
		}

		mediatype := "application/vnd.ollama.image.model"
		if ggml.Name() == "ggla" {
			mediatype = "application/vnd.ollama.image.adapter"
		} else if ggml.KV().Architecture() == "clip" {
			mediatype = "application/vnd.ollama.image.projector"
		}

		layer, err := NewLayer(io.NewSectionReader(file, offset, n), mediatype)
		if err != nil {
			return nil, err
		}

184
		layers = append(layers, &layerGGML{layer, ggml})
Michael Yang's avatar
Michael Yang committed
185
186
187
		offset = n
	}

188
189
190
191
192
193
	return detectChatTemplate(layers)
}

func detectChatTemplate(layers []*layerGGML) ([]*layerGGML, error) {
	for _, layer := range layers {
		if s := layer.GGML.KV().ChatTemplate(); s != "" {
Michael Yang's avatar
Michael Yang committed
194
			if t, err := template.Named(s); err != nil {
195
196
				slog.Debug("template detection", "error", err)
			} else {
197
				layer, err := NewLayer(t.Reader(), "application/vnd.ollama.image.template")
198
199
200
201
				if err != nil {
					return nil, err
				}

202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
				layer.status = fmt.Sprintf("using autodetected template %s", t.Name)
				layers = append(layers, &layerGGML{layer, nil})

				if t.Parameters != nil {
					var b bytes.Buffer
					if err := json.NewEncoder(&b).Encode(t.Parameters); err != nil {
						return nil, err
					}

					layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
					if err != nil {
						return nil, err
					}

					layers = append(layers, &layerGGML{layer, nil})
				}
218
219
220
221
			}
		}
	}

Michael Yang's avatar
Michael Yang committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
	return layers, nil
}

func detectContentType(r io.Reader) (string, error) {
	var b bytes.Buffer
	if _, err := io.Copy(&b, r); err != nil {
		return "", err
	}

	if contentType := llm.DetectGGMLType(b.Bytes()); contentType != "" {
		return contentType, nil
	}

	if contentType := http.DetectContentType(b.Bytes()); contentType != "application/octet-stream" {
		return contentType, nil
	}

	return "unknown", nil
}
Michael Yang's avatar
tools  
Michael Yang committed
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258

// parseToolCalls attempts to parse a JSON string into a slice of ToolCalls.
// mxyng: this only really works if the input contains tool calls in some JSON format
func (m *Model) parseToolCalls(s string) ([]api.ToolCall, bool) {
	// create a subtree from the node that ranges over .ToolCalls
	tmpl := m.Template.Subtree(func(n parse.Node) bool {
		if t, ok := n.(*parse.RangeNode); ok {
			return slices.Contains(template.Identifiers(t.Pipe), "ToolCalls")
		}

		return false
	})

	if tmpl == nil {
		return nil, false
	}

	var b bytes.Buffer
259
	if err := tmpl.Execute(&b, map[string][]api.ToolCall{
Michael Yang's avatar
tools  
Michael Yang committed
260
261
		"ToolCalls": {
			{
262
263
264
265
266
				Function: api.ToolCallFunction{
					Name: "@@name@@",
					Arguments: api.ToolCallFunctionArguments{
						"@@argument@@": 1,
					},
Michael Yang's avatar
tools  
Michael Yang committed
267
268
269
270
271
272
273
				},
			},
		},
	}); err != nil {
		return nil, false
	}

274
	var kv map[string]any
Michael Yang's avatar
tools  
Michael Yang committed
275
	// execute the subtree with placeholders to identify the keys
Michael Yang's avatar
Michael Yang committed
276
277
	// trim any commands that might exist in the template
	if err := json.Unmarshal(bytes.TrimSuffix(b.Bytes(), []byte(",")), &kv); err != nil {
Michael Yang's avatar
tools  
Michael Yang committed
278
279
280
281
282
283
		return nil, false
	}

	// find the keys that correspond to the name and arguments fields
	var name, arguments string
	for k, v := range kv {
284
285
		switch v.(type) {
		case string:
Michael Yang's avatar
tools  
Michael Yang committed
286
			name = k
287
		case map[string]any:
Michael Yang's avatar
tools  
Michael Yang committed
288
289
290
291
			arguments = k
		}
	}

292
293
294
295
	if name == "" || arguments == "" {
		return nil, false
	}

Michael Yang's avatar
Michael Yang committed
296
297
	var objs []map[string]any
	for offset := 0; offset < len(s); {
298
299
		var obj map[string]any
		decoder := json.NewDecoder(strings.NewReader(s[offset:]))
Michael Yang's avatar
Michael Yang committed
300
		if err := decoder.Decode(&obj); errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
Michael Yang's avatar
Michael Yang committed
301
302
303
304
305
306
307
308
			break
		} else if syntax := &(json.SyntaxError{}); errors.As(err, &syntax) {
			// skip over any syntax errors
			offset += int(syntax.Offset)
		} else if unmarshalType := &(json.UnmarshalTypeError{}); errors.As(err, &unmarshalType) {
			// skip over any unmarshalable types
			offset += int(unmarshalType.Offset)
		} else if err != nil {
Michael Yang's avatar
Michael Yang committed
309
			slog.Error("parseToolCalls", "error", err)
Michael Yang's avatar
tools  
Michael Yang committed
310
			return nil, false
Michael Yang's avatar
Michael Yang committed
311
		} else {
312
			offset += int(decoder.InputOffset())
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331

			// collect all nested objects
			var collect func(any) []map[string]any
			collect = func(obj any) (all []map[string]any) {
				switch o := obj.(type) {
				case map[string]any:
					all = append(all, o)
					for _, v := range o {
						all = append(all, collect(v)...)
					}
				case []any:
					for _, v := range o {
						all = append(all, collect(v)...)
					}
				}

				return all
			}
			objs = append(objs, collect(obj)...)
Michael Yang's avatar
tools  
Michael Yang committed
332
333
334
335
		}
	}

	var toolCalls []api.ToolCall
Michael Yang's avatar
Michael Yang committed
336
	for _, kv := range objs {
337
338
339
340
341
342
343
344
345
		n, nok := kv[name].(string)
		a, aok := kv[arguments].(map[string]any)
		if nok && aok {
			toolCalls = append(toolCalls, api.ToolCall{
				Function: api.ToolCallFunction{
					Name:      n,
					Arguments: a,
				},
			})
Michael Yang's avatar
tools  
Michael Yang committed
346
347
348
		}
	}

Michael Yang's avatar
Michael Yang committed
349
	return toolCalls, len(toolCalls) > 0
Michael Yang's avatar
tools  
Michael Yang committed
350
}