ollama

1
package llm
2

3
import (
4
	"bytes"
5
	"encoding/binary"
6
	"fmt"
7
	"io"
8

9
	"github.com/jmorganca/ollama/format"
10
)
11

12
type containerGGUF struct {
13
	bo binary.ByteOrder
14

15
	Version uint32
16

17
	V1 struct {
18
		NumTensor uint32
19
		NumKV     uint32
20
	}
21

22
	V2 struct {
23
		NumTensor uint64
24
		NumKV     uint64
25
	}
26
}
27

28
func (c *containerGGUF) Name() string {
29
	return "gguf"
30
}
31

32
func (c *containerGGUF) Decode(rso *readSeekOffset) (model, error) {
33
	binary.Read(rso, c.bo, &c.Version)
34

35
	switch c.Version {
36
	case 1:
37
		binary.Read(rso, c.bo, &c.V1)
38
	default:
39
		binary.Read(rso, c.bo, &c.V2)
40
	}
41

42
	model := newGGUFModel(c)
43
	if err := model.Decode(rso); err != nil {
44
		return nil, err
45
	}
46

47
	return model, nil
48
}
49

50
const (
51
	ggufTypeUint8 uint32 = iota
52
	ggufTypeInt8
53
	ggufTypeUint16
54
	ggufTypeInt16
55
	ggufTypeUint32
56
	ggufTypeInt32
57
	ggufTypeFloat32
58
	ggufTypeBool
59
	ggufTypeString
60
	ggufTypeArray
61
	ggufTypeUint64
62
	ggufTypeInt64
63
	ggufTypeFloat64
64
)
65

66
type kv map[string]any
67

68
type tensor struct {
69
	name   string
70
	kind   uint32
71
	offset uint64
72

73
	// shape is the number of elements in each dimension
74
	shape [4]uint64
75
}
76

77
func (t tensor) blockSize() uint64 {
78
	switch {
79
	case t.kind < 2:
80
		return 1
81
	case t.kind < 10:
82
		return 32
83
	default:
84
		return 256
85
	}
86
}
87

88
func (t tensor) typeSize() uint64 {
89
	blockSize := t.blockSize()
90

91
	switch t.kind {
92
	case 0: // FP32
93
		return 4
94
	case 1: // FP16
95
		return 2
96
	case 2: // Q4_0
97
		return 2 + blockSize/2
98
	case 3: // Q4_1
99
		return 2 + 2 + blockSize/2
100
	case 6: // Q5_0
101
		return 2 + 4 + blockSize/2
102
	case 7: // Q5_1
103
		return 2 + 2 + 4 + blockSize/2
104
	case 8: // Q8_0
105
		return 2 + blockSize
106
	case 9: // Q8_1
107
		return 4 + 4 + blockSize
108
	case 10: // Q2_K
109
		return blockSize/16 + blockSize/4 + 2 + 2
110
	case 11: // Q3_K
111
		return blockSize/8 + blockSize/4 + 12 + 2
112
	case 12: // Q4_K
113
		return 2 + 2 + 12 + blockSize/2
114
	case 13: // Q5_K
115
		return 2 + 2 + 12 + blockSize/8 + blockSize/2
116
	case 14: // Q6_K
117
		return blockSize/2 + blockSize/4 + blockSize/16 + 2
118
	case 15: // Q8_K
119
		return 2 + blockSize + 2*blockSize/16
120
	case 16: // IQ2_XXS
121
		return 2 + 2*blockSize/8
122
	case 17: // IQ2_XS
123
		return 2 + 2*blockSize/8 + blockSize/32
124
	case 18: // IQ3_XXS
125
		return 2 + 3*blockSize/8
126
	default:
127
		return 0
128
	}
129
}
130

131
func (t tensor) parameters() uint64 {
132
	return t.shape[0] * t.shape[1] * t.shape[2] * t.shape[3]
133
}
134

135
func (t tensor) size() uint64 {
136
	return t.parameters() * t.typeSize() / t.blockSize()
137
}
138

139
type ggufModel struct {
140
	*containerGGUF
141

142
	kv
143
	tensors []tensor
144

145
	parameters uint64
146
}
147

148
func newGGUFModel(container *containerGGUF) *ggufModel {
149
	return &ggufModel{
150
		containerGGUF: container,
151
		kv:            make(kv),
152
	}
153
}
154

155
func (llm *ggufModel) NumTensor() uint64 {
156
	if llm.Version == 1 {
157
		return uint64(llm.V1.NumTensor)
158
	}
159

160
	return llm.V2.NumTensor
161
}
162

163
func (llm *ggufModel) NumKV() uint64 {
164
	if llm.Version == 1 {
165
		return uint64(llm.V1.NumKV)
166
	}
167

168
	return llm.V2.NumKV
169
}
170

171
func (llm *ggufModel) ModelFamily() string {
172
	if t, ok := llm.kv["general.architecture"].(string); ok {
173
		return t
174
	}
175

176
	return "unknown"
177
}
178

179
func (llm *ggufModel) ModelType() string {
180
	if llm.parameters > 0 {
181
		return format.HumanNumber(llm.parameters)
182
	}
183

184
	return "unknown"
185
}
186

187
func (llm *ggufModel) FileType() string {
188
	if t, ok := llm.kv["general.file_type"].(uint32); ok {
189
		return fileType(t)
190
	}
191

192
	return "unknown"
193
}
194

195
func (llm *ggufModel) Decode(rso *readSeekOffset) error {
196
	// decode key-values
197
	for i := 0; uint64(i) < llm.NumKV(); i++ {
198
		k, err := llm.readString(rso)
199
		if err != nil {
200
			return err
201
		}
202

203
		vtype := llm.readU32(rso)
204

205
		var v any
206
		switch vtype {
207
		case ggufTypeUint8:
208
			v = llm.readU8(rso)
209
		case ggufTypeInt8:
210
			v = llm.readI8(rso)
211
		case ggufTypeUint16:
212
			v = llm.readU16(rso)
213
		case ggufTypeInt16:
214
			v = llm.readI16(rso)
215
		case ggufTypeUint32:
216
			v = llm.readU32(rso)
217
		case ggufTypeInt32:
218
			v = llm.readI32(rso)
219
		case ggufTypeUint64:
220
			v = llm.readU64(rso)
221
		case ggufTypeInt64:
222
			v = llm.readI64(rso)
223
		case ggufTypeFloat32:
224
			v = llm.readF32(rso)
225
		case ggufTypeFloat64:
226
			v = llm.readF64(rso)
227
		case ggufTypeBool:
228
			v = llm.readBool(rso)
229
		case ggufTypeString:
230
			s, err := llm.readString(rso)
231
			if err != nil {
232
				return err
233
			}
234

235
			v = s
236
		case ggufTypeArray:
237
			a, err := llm.readArray(rso)
238
			if err != nil {
239
				return err
240
			}
241

242
			v = a
243
		default:
244
			return fmt.Errorf("invalid type: %d", vtype)
245
		}
246

247
		llm.kv[k] = v
248
	}
249

250
	// decode tensors
251
	for i := 0; uint64(i) < llm.NumTensor(); i++ {
252
		name, err := llm.readString(rso)
253
		if err != nil {
254
			return err
255
		}
256

257
		// dims is the number of dimensions in the tensor
258
		dims := llm.readU32(rso)
259

260
		shape := [4]uint64{1, 1, 1, 1}
261
		for i := 0; uint32(i) < dims; i++ {
262
			shape[i] = llm.readU64(rso)
263
		}
264

265
		tensor := tensor{
266
			name:   name,
267
			kind:   llm.readU32(rso),
268
			offset: llm.readU64(rso),
269
			shape:  shape,
270
		}
271

272
		llm.tensors = append(llm.tensors, tensor)
273
		llm.parameters += tensor.parameters()
274
	}
275

276
	alignment, ok := llm.kv["general.alignment"].(uint32)
277
	if !ok {
278
		alignment = 32
279
	}
280

281
	rso.Seek(int64(alignment)-rso.offset%int64(alignment), io.SeekCurrent)
282
	for _, tensor := range llm.tensors {
283
		padded := (int64(tensor.size()) + int64(alignment) - 1) & ^(int64(alignment) - 1)
284
		rso.Seek(padded, io.SeekCurrent)
285
	}
286

287
	return nil
288
}
289

290
func (llm *ggufModel) NumLayers() uint32 {
291
	value, exists := llm.kv[fmt.Sprintf("%s.block_count", llm.ModelFamily())]
292
	if !exists {
293
		return 0
294
	}
295

296
	return value.(uint32)
297
}
298

299
func (llm *ggufModel) NumHead() uint32 {
300
	value, exists := llm.kv[fmt.Sprintf("%s.attention.head_count", llm.ModelFamily())]
301
	if !exists {
302
		return 0
303
	}
304

305
	return value.(uint32)
306
}
307

308
func (llm *ggufModel) NumEmbed() uint32 {
309
	value, exists := llm.kv[fmt.Sprintf("%s.embedding_length", llm.ModelFamily())]
310
	if !exists {
311
		return 0
312
	}
313

314
	return value.(uint32)
315
}
316

317
func (llm *ggufModel) NumHeadKv() uint32 {
318
	value, exists := llm.kv[fmt.Sprintf("%s.attention.head_count_kv", llm.ModelFamily())]
319
	if !exists {
320
		return 0
321
	}
322

323
	return value.(uint32)
324
}
325

326
func (llm *ggufModel) NumCtx() uint32 {
327
	value, exists := llm.kv[fmt.Sprintf("%s.context_length", llm.ModelFamily())]
328
	if !exists {
329
		return 0
330
	}
331

332
	return value.(uint32)
333
}
334

335
func (llm *ggufModel) NumGQA() uint32 {
336
	numHeadKv := llm.NumHeadKv()
337
	if numHeadKv == 0 {
338
		return 0
339
	}
340

341
	return llm.NumHead() / numHeadKv
342
}
343

344
func (llm ggufModel) readU8(r io.Reader) uint8 {
345
	var u8 uint8
346
	binary.Read(r, llm.bo, &u8)
347
	return u8
348
}
349

350
func (llm ggufModel) readI8(r io.Reader) int8 {
351
	var i8 int8
352
	binary.Read(r, llm.bo, &i8)
353
	return i8
354
}
355

356
func (llm ggufModel) readU16(r io.Reader) uint16 {
357
	var u16 uint16
358
	binary.Read(r, llm.bo, &u16)
359
	return u16
360
}
361

362
func (llm ggufModel) readI16(r io.Reader) int16 {
363
	var i16 int16
364
	binary.Read(r, llm.bo, &i16)
365
	return i16
366
}
367

368
func (llm ggufModel) readU32(r io.Reader) uint32 {
369
	var u32 uint32
370
	binary.Read(r, llm.bo, &u32)
371
	return u32
372
}
373

374
func (llm ggufModel) readI32(r io.Reader) int32 {
375
	var i32 int32
376
	binary.Read(r, llm.bo, &i32)
377
	return i32
378
}
379

380
func (llm ggufModel) readU64(r io.Reader) uint64 {
381
	var u64 uint64
382
	binary.Read(r, llm.bo, &u64)
383
	return u64
384
}
385

386
func (llm ggufModel) readI64(r io.Reader) int64 {
387
	var i64 int64
388
	binary.Read(r, llm.bo, &i64)
389
	return i64
390
}
391

392
func (llm ggufModel) readF32(r io.Reader) float32 {
393
	var f32 float32
394
	binary.Read(r, llm.bo, &f32)
395
	return f32
396
}
397

398
func (llm ggufModel) readF64(r io.Reader) float64 {
399
	var f64 float64
400
	binary.Read(r, llm.bo, &f64)
401
	return f64
402
}
403

404
func (llm ggufModel) readBool(r io.Reader) bool {
405
	var b bool
406
	binary.Read(r, llm.bo, &b)
407
	return b
408
}
409

410
func (llm ggufModel) readStringV1(r io.Reader) (string, error) {
411
	var nameLength uint32
412
	binary.Read(r, llm.bo, &nameLength)
413

414
	var b bytes.Buffer
415
	if _, err := io.CopyN(&b, r, int64(nameLength)); err != nil {
416
		return "", err
417
	}
418

419
	// gguf v1 strings are null-terminated
420
	b.Truncate(b.Len() - 1)
421

422
	return b.String(), nil
423
}
424

425
func (llm ggufModel) readString(r io.Reader) (string, error) {
426
	if llm.Version == 1 {
427
		return llm.readStringV1(r)
428
	}
429

430
	var nameLength uint64
431
	binary.Read(r, llm.bo, &nameLength)
432

433
	var b bytes.Buffer
434
	if _, err := io.CopyN(&b, r, int64(nameLength)); err != nil {
435
		return "", err
436
	}
437

438
	return b.String(), nil
439
}
440

441
func (llm *ggufModel) readArrayV1(r io.Reader) (arr []any, err error) {
442
	atype := llm.readU32(r)
443
	n := llm.readU32(r)
444

445
	for i := 0; uint32(i) < n; i++ {
446
		switch atype {
447
		case ggufTypeUint8:
448
			arr = append(arr, llm.readU8(r))
449
		case ggufTypeInt8:
450
			arr = append(arr, llm.readI8(r))
451
		case ggufTypeUint16:
452
			arr = append(arr, llm.readU16(r))
453
		case ggufTypeInt16:
454
			arr = append(arr, llm.readI16(r))
455
		case ggufTypeUint32:
456
			arr = append(arr, llm.readU32(r))
457
		case ggufTypeInt32:
458
			arr = append(arr, llm.readI32(r))
459
		case ggufTypeFloat32:
460
			arr = append(arr, llm.readF32(r))
461
		case ggufTypeBool:
462
			arr = append(arr, llm.readBool(r))
463
		case ggufTypeString:
464
			s, err := llm.readStringV1(r)
465
			if err != nil {
466
				return nil, err
467
			}
468

469
			arr = append(arr, s)
470
		default:
471
			return nil, fmt.Errorf("invalid array type: %d", atype)
472
		}
473
	}
474

475
	return
476
}
477

478
func (llm *ggufModel) readArray(r io.Reader) (arr []any, err error) {
479
	if llm.Version == 1 {
480
		return llm.readArrayV1(r)
481
	}
482

483
	atype := llm.readU32(r)
484
	n := llm.readU64(r)
485

486
	for i := 0; uint64(i) < n; i++ {
487
		switch atype {
488
		case ggufTypeUint8:
489
			arr = append(arr, llm.readU8(r))
490
		case ggufTypeInt8:
491
			arr = append(arr, llm.readI8(r))
492
		case ggufTypeUint16:
493
			arr = append(arr, llm.readU16(r))
494
		case ggufTypeInt16:
495
			arr = append(arr, llm.readI16(r))
496
		case ggufTypeUint32:
497
			arr = append(arr, llm.readU32(r))
498
		case ggufTypeInt32:
499
			arr = append(arr, llm.readI32(r))
500
		case ggufTypeUint64:
501
			arr = append(arr, llm.readU64(r))
502
		case ggufTypeInt64:
503
			arr = append(arr, llm.readI64(r))
504
		case ggufTypeFloat32:
505
			arr = append(arr, llm.readF32(r))
506
		case ggufTypeFloat64:
507
			arr = append(arr, llm.readF64(r))
508
		case ggufTypeBool:
509
			arr = append(arr, llm.readBool(r))
510
		case ggufTypeString:
511
			s, err := llm.readString(r)
512
			if err != nil {
513
				return nil, err
514
			}
515

516
			arr = append(arr, s)
517
		default:
518
			return nil, fmt.Errorf("invalid array type: %d", atype)
519
		}
520
	}
521

522
	return
523
}
524