model.go 6.21 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
8
9
package server

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
10
	"log/slog"
Michael Yang's avatar
Michael Yang committed
11
12
13
14
15
16
17
	"net/http"
	"os"
	"path/filepath"

	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/convert"
	"github.com/ollama/ollama/llm"
Michael Yang's avatar
Michael Yang committed
18
	"github.com/ollama/ollama/template"
Michael Yang's avatar
Michael Yang committed
19
20
21
	"github.com/ollama/ollama/types/model"
)

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

24
type layerGGML struct {
Michael Yang's avatar
Michael Yang committed
25
26
27
28
	*Layer
	*llm.GGML
}

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

37
		m, err = ParseNamedManifest(name)
Michael Yang's avatar
Michael Yang committed
38
39
40
		if err != nil {
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
41
42
43
44
	case err != nil:
		return nil, err
	}

45
46
	for _, layer := range m.Layers {
		layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, name.DisplayShortest())
Michael Yang's avatar
Michael Yang committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
		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()

66
			ggml, _, err := llm.DecodeGGML(blob, 0)
Michael Yang's avatar
Michael Yang committed
67
68
69
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
70

71
			layers = append(layers, &layerGGML{layer, ggml})
Michael Yang's avatar
Michael Yang committed
72
		default:
73
			layers = append(layers, &layerGGML{layer, nil})
Michael Yang's avatar
Michael Yang committed
74
75
76
77
78
79
		}
	}

	return layers, nil
}

80
func extractFromZipFile(p string, file *os.File, fn func(api.ProgressResponse)) error {
Michael Yang's avatar
Michael Yang committed
81
82
	stat, err := file.Stat()
	if err != nil {
83
		return err
Michael Yang's avatar
Michael Yang committed
84
85
86
87
	}

	r, err := zip.NewReader(file, stat.Size())
	if err != nil {
88
		return err
Michael Yang's avatar
Michael Yang committed
89
90
91
92
	}

	fn(api.ProgressResponse{Status: "unpacking model metadata"})
	for _, f := range r.File {
Michael Yang's avatar
Michael Yang committed
93
94
		if !filepath.IsLocal(f.Name) {
			return fmt.Errorf("%w: %s", zip.ErrInsecurePath, f.Name)
95
96
		}

Michael Yang's avatar
Michael Yang committed
97
		n := filepath.Join(p, f.Name)
98
99
100
101
		if err := os.MkdirAll(filepath.Dir(n), 0o750); err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
102
		// TODO(mxyng): this should not write out all files to disk
103
		outfile, err := os.Create(n)
Michael Yang's avatar
Michael Yang committed
104
		if err != nil {
105
			return err
Michael Yang's avatar
Michael Yang committed
106
		}
Michael Yang's avatar
Michael Yang committed
107
		defer outfile.Close()
Michael Yang's avatar
Michael Yang committed
108
109
110

		infile, err := f.Open()
		if err != nil {
111
			return err
Michael Yang's avatar
Michael Yang committed
112
		}
Michael Yang's avatar
Michael Yang committed
113
		defer infile.Close()
Michael Yang's avatar
Michael Yang committed
114
115

		if _, err = io.Copy(outfile, infile); err != nil {
116
			return err
Michael Yang's avatar
Michael Yang committed
117
118
119
		}

		if err := outfile.Close(); err != nil {
120
			return err
Michael Yang's avatar
Michael Yang committed
121
122
123
		}

		if err := infile.Close(); err != nil {
124
			return err
Michael Yang's avatar
Michael Yang committed
125
126
127
		}
	}

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
	return nil
}

func parseFromZipFile(_ context.Context, file *os.File, digest string, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
	tempDir, err := os.MkdirTemp(filepath.Dir(file.Name()), "")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tempDir)

	if err := extractFromZipFile(tempDir, file, fn); err != nil {
		return nil, err
	}

	mf, err := convert.GetModelFormat(tempDir)
Michael Yang's avatar
Michael Yang committed
143
144
145
146
	if err != nil {
		return nil, err
	}

147
	params, err := mf.GetParams(tempDir)
Michael Yang's avatar
Michael Yang committed
148
149
150
151
	if err != nil {
		return nil, err
	}

152
	mArch, err := mf.GetModelArch("", tempDir, params)
Michael Yang's avatar
Michael Yang committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
	if err != nil {
		return nil, err
	}

	fn(api.ProgressResponse{Status: "processing tensors"})
	if err := mArch.GetTensors(); err != nil {
		return nil, err
	}

	if err := mArch.LoadVocab(); err != nil {
		return nil, err
	}

	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")
170
	temp, err := os.CreateTemp(tempDir, "fp16")
Michael Yang's avatar
Michael Yang committed
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
	if err != nil {
		return nil, err
	}
	defer temp.Close()
	defer os.Remove(temp.Name())

	if err = mArch.WriteGGUF(temp); err != nil {
		return nil, err
	}

	if _, err := temp.Seek(0, io.SeekStart); err != nil {
		return nil, err
	}

	layer, err := NewLayer(temp, "application/vnd.ollama.image.model")
	if err != nil {
Michael Yang's avatar
Michael Yang committed
187
		return nil, err
Michael Yang's avatar
Michael Yang committed
188
189
	}

190
	bin, err := layer.Open()
Michael Yang's avatar
Michael Yang committed
191
192
193
194
195
	if err != nil {
		return nil, err
	}
	defer bin.Close()

196
	ggml, _, err := llm.DecodeGGML(bin, 0)
Michael Yang's avatar
Michael Yang committed
197
198
199
200
	if err != nil {
		return nil, err
	}

201
	layers = append(layers, &layerGGML{layer, ggml})
202

Michael Yang's avatar
Michael Yang committed
203
	intermediateBlobs[digest] = layer.Digest
204
	return detectChatTemplate(layers)
Michael Yang's avatar
Michael Yang committed
205
206
}

207
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
208
209
210
211
212
213
214
215
216
217
	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":
218
		return parseFromZipFile(ctx, file, digest, fn)
Michael Yang's avatar
Michael Yang committed
219
220
221
222
223
224
225
226
227
228
229
	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() {
230
		ggml, n, err := llm.DecodeGGML(file, 0)
Michael Yang's avatar
Michael Yang committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
		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
		}

249
		layers = append(layers, &layerGGML{layer, ggml})
Michael Yang's avatar
Michael Yang committed
250
251
252
		offset = n
	}

253
254
255
256
257
258
	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
259
			if t, err := template.Named(s); err != nil {
260
261
262
263
264
265
266
267
268
269
270
271
272
				slog.Debug("template detection", "error", err)
			} else {
				tmpl, err := NewLayer(t.Reader(), "application/vnd.ollama.image.template")
				if err != nil {
					return nil, err
				}

				tmpl.status = fmt.Sprintf("using autodetected template %s", t.Name)
				layers = append(layers, &layerGGML{tmpl, nil})
			}
		}
	}

Michael Yang's avatar
Michael Yang committed
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
	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
}