gguf.go 13.5 KB
Newer Older
mashun1's avatar
v1  
mashun1 committed
1
2
3
4
package llm

import (
	"bytes"
xuxzh1's avatar
init  
xuxzh1 committed
5
	"cmp"
mashun1's avatar
v1  
mashun1 committed
6
	"encoding/binary"
xuxzh1's avatar
init  
xuxzh1 committed
7
	"encoding/json"
mashun1's avatar
v1  
mashun1 committed
8
9
	"fmt"
	"io"
xuxzh1's avatar
init  
xuxzh1 committed
10
11
	"log/slog"
	"slices"
mashun1's avatar
v1  
mashun1 committed
12
13
	"strings"

xuxzh1's avatar
init  
xuxzh1 committed
14
	"golang.org/x/exp/maps"
mashun1's avatar
v1  
mashun1 committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
)

type containerGGUF struct {
	ByteOrder binary.ByteOrder

	Version uint32

	V1 struct {
		NumTensor uint32
		NumKV     uint32
	}

	V2 struct {
		NumTensor uint64
		NumKV     uint64
	}

	V3 struct {
		NumTensor uint64
		NumKV     uint64
	}
xuxzh1's avatar
init  
xuxzh1 committed
36
37
38
39
40
41

	maxArraySize int
}

func (c *containerGGUF) canCollectArray(size int) bool {
	return c.maxArraySize < 0 || size <= c.maxArraySize
mashun1's avatar
v1  
mashun1 committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
}

func (c *containerGGUF) Name() string {
	return "gguf"
}

func (c *containerGGUF) Decode(rs io.ReadSeeker) (model, error) {
	if err := binary.Read(rs, c.ByteOrder, &c.Version); err != nil {
		return nil, err
	}

	var err error
	switch c.Version {
	case 1:
		err = binary.Read(rs, c.ByteOrder, &c.V1)
	case 2:
		err = binary.Read(rs, c.ByteOrder, &c.V2)
	default:
		err = binary.Read(rs, c.ByteOrder, &c.V3)
	}
	if err != nil {
		return nil, err
	}

	model := newGGUF(c)
	if err := model.Decode(rs); err != nil {
		return nil, err
	}

	return model, nil
}

const (
	ggufTypeUint8 uint32 = iota
	ggufTypeInt8
	ggufTypeUint16
	ggufTypeInt16
	ggufTypeUint32
	ggufTypeInt32
	ggufTypeFloat32
	ggufTypeBool
	ggufTypeString
	ggufTypeArray
	ggufTypeUint64
	ggufTypeInt64
	ggufTypeFloat64
)

type gguf struct {
	*containerGGUF

	kv      KV
	tensors []*Tensor

xuxzh1's avatar
init  
xuxzh1 committed
96
97
98
99
	parameters   uint64
	tensorOffset uint64

	scratch [16 << 10]byte
mashun1's avatar
v1  
mashun1 committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
}

func newGGUF(container *containerGGUF) *gguf {
	return &gguf{
		containerGGUF: container,
		kv:            make(KV),
	}
}

func (llm *gguf) KV() KV {
	return llm.kv
}

func (llm *gguf) Tensors() Tensors {
xuxzh1's avatar
init  
xuxzh1 committed
114
115
116
117
	return Tensors{
		Items:  llm.tensors,
		Offset: llm.tensorOffset,
	}
mashun1's avatar
v1  
mashun1 committed
118
119
120
121
122
123
124
125
126
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
}

func (llm *gguf) numTensor() uint64 {
	switch llm.Version {
	case 1:
		return uint64(llm.V1.NumTensor)
	case 2:
		return llm.V2.NumTensor
	default:
		return llm.V3.NumTensor
	}
}

func (llm *gguf) numKV() uint64 {
	switch llm.Version {
	case 1:
		return uint64(llm.V1.NumKV)
	case 2:
		return llm.V2.NumKV
	default:
		return llm.V3.NumKV
	}
}

func (llm *gguf) Decode(rs io.ReadSeeker) error {
	// decode key-values
	for i := 0; uint64(i) < llm.numKV(); i++ {
		k, err := readGGUFString(llm, rs)
		if err != nil {
			return err
		}

		t, err := readGGUF[uint32](llm, rs)
		if err != nil {
			return err
		}

		var v any
		switch t {
		case ggufTypeUint8:
			v, err = readGGUF[uint8](llm, rs)
		case ggufTypeInt8:
			v, err = readGGUF[int8](llm, rs)
		case ggufTypeUint16:
			v, err = readGGUF[uint16](llm, rs)
		case ggufTypeInt16:
			v, err = readGGUF[int16](llm, rs)
		case ggufTypeUint32:
			v, err = readGGUF[uint32](llm, rs)
		case ggufTypeInt32:
			v, err = readGGUF[int32](llm, rs)
		case ggufTypeUint64:
			v, err = readGGUF[uint64](llm, rs)
		case ggufTypeInt64:
			v, err = readGGUF[int64](llm, rs)
		case ggufTypeFloat32:
			v, err = readGGUF[float32](llm, rs)
		case ggufTypeFloat64:
			v, err = readGGUF[float64](llm, rs)
		case ggufTypeBool:
			v, err = readGGUF[bool](llm, rs)
		case ggufTypeString:
			v, err = readGGUFString(llm, rs)
		case ggufTypeArray:
			v, err = readGGUFArray(llm, rs)
		default:
			return fmt.Errorf("invalid type: %d", t)
		}

		if err != nil {
			return err
		}

		llm.kv[k] = v
	}

	// decode tensors
xuxzh1's avatar
init  
xuxzh1 committed
195
	for range llm.numTensor() {
mashun1's avatar
v1  
mashun1 committed
196
197
		name, err := readGGUFString(llm, rs)
		if err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
198
			return fmt.Errorf("failed to read tensor name: %w", err)
mashun1's avatar
v1  
mashun1 committed
199
200
201
202
203
		}

		// dims is the number of dimensions in the tensor
		dims, err := readGGUF[uint32](llm, rs)
		if err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
204
			return fmt.Errorf("failed to read tensor dimensions: %w", err)
mashun1's avatar
v1  
mashun1 committed
205
206
		}

xuxzh1's avatar
init  
xuxzh1 committed
207
		shape := make([]uint64, dims)
mashun1's avatar
v1  
mashun1 committed
208
209
210
		for i := 0; uint32(i) < dims; i++ {
			shape[i], err = readGGUF[uint64](llm, rs)
			if err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
211
				return fmt.Errorf("failed to read tensor shape: %w", err)
mashun1's avatar
v1  
mashun1 committed
212
213
214
215
216
			}
		}

		kind, err := readGGUF[uint32](llm, rs)
		if err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
217
			return fmt.Errorf("failed to read tensor kind: %w", err)
mashun1's avatar
v1  
mashun1 committed
218
219
220
221
		}

		offset, err := readGGUF[uint64](llm, rs)
		if err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
222
			return fmt.Errorf("failed to read tensor offset: %w", err)
mashun1's avatar
v1  
mashun1 committed
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
		}

		tensor := Tensor{
			Name:   name,
			Kind:   kind,
			Offset: offset,
			Shape:  shape[:],
		}

		llm.tensors = append(llm.tensors, &tensor)
		llm.parameters += tensor.parameters()
	}

	// patch KV with parameter count
	llm.kv["general.parameter_count"] = llm.parameters

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

	offset, err := rs.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
249
250
	padding := ggufPadding(offset, int64(alignment))
	llm.tensorOffset = uint64(offset + padding)
mashun1's avatar
v1  
mashun1 committed
251
252

	for _, tensor := range llm.tensors {
xuxzh1's avatar
init  
xuxzh1 committed
253
254
255
		offset, err := rs.Seek(0, io.SeekCurrent)
		if err != nil {
			return fmt.Errorf("failed to get current offset: %w", err)
mashun1's avatar
v1  
mashun1 committed
256
257
		}

xuxzh1's avatar
init  
xuxzh1 committed
258
		padding := ggufPadding(offset, int64(alignment))
mashun1's avatar
v1  
mashun1 committed
259
		if _, err := rs.Seek(padding, io.SeekCurrent); err != nil {
xuxzh1's avatar
init  
xuxzh1 committed
260
261
262
263
264
			return fmt.Errorf("failed to seek to init padding: %w", err)
		}

		if _, err := rs.Seek(int64(tensor.Size()), io.SeekCurrent); err != nil {
			return fmt.Errorf("failed to seek to tensor: %w", err)
mashun1's avatar
v1  
mashun1 committed
265
266
267
268
269
270
271
272
273
274
275
276
		}
	}

	return nil
}

func readGGUF[T any](llm *gguf, r io.Reader) (T, error) {
	var t T
	err := binary.Read(r, llm.ByteOrder, &t)
	return t, err
}

xuxzh1's avatar
init  
xuxzh1 committed
277
278
func writeGGUF[V any](w io.Writer, t uint32, v V) error {
	if err := binary.Write(w, binary.LittleEndian, t); err != nil {
mashun1's avatar
v1  
mashun1 committed
279
280
281
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
282
	return binary.Write(w, binary.LittleEndian, v)
mashun1's avatar
v1  
mashun1 committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
}

func readGGUFV1String(llm *gguf, r io.Reader) (string, error) {
	var length uint64
	if err := binary.Read(r, llm.ByteOrder, &length); err != nil {
		return "", err
	}

	var b bytes.Buffer
	if _, err := io.CopyN(&b, r, int64(length)); err != nil {
		return "", err
	}

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

	return b.String(), nil
}

xuxzh1's avatar
init  
xuxzh1 committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
func discardGGUFString(llm *gguf, r io.Reader) error {
	buf := llm.scratch[:8]
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return err
	}

	size := int(llm.ByteOrder.Uint64(buf))
	for size > 0 {
		n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))])
		if err != nil {
			return err
		}
		size -= n
	}
	return nil
}

mashun1's avatar
v1  
mashun1 committed
320
321
322
323
324
func readGGUFString(llm *gguf, r io.Reader) (string, error) {
	if llm.Version == 1 {
		return readGGUFV1String(llm, r)
	}

xuxzh1's avatar
init  
xuxzh1 committed
325
326
327
	buf := llm.scratch[:8]
	_, err := io.ReadFull(r, buf)
	if err != nil {
mashun1's avatar
v1  
mashun1 committed
328
329
330
		return "", err
	}

xuxzh1's avatar
init  
xuxzh1 committed
331
332
333
334
335
	length := int(llm.ByteOrder.Uint64(buf))
	if length > len(llm.scratch) {
		buf = make([]byte, length)
	} else {
		buf = llm.scratch[:length]
mashun1's avatar
v1  
mashun1 committed
336
	}
xuxzh1's avatar
init  
xuxzh1 committed
337
	clear(buf)
mashun1's avatar
v1  
mashun1 committed
338

xuxzh1's avatar
init  
xuxzh1 committed
339
340
341
342
343
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return "", err
	}
	return string(buf), nil
mashun1's avatar
v1  
mashun1 committed
344
345
}

xuxzh1's avatar
init  
xuxzh1 committed
346
347
func writeGGUFString(w io.Writer, s string) error {
	if err := binary.Write(w, binary.LittleEndian, ggufTypeString); err != nil {
mashun1's avatar
v1  
mashun1 committed
348
349
350
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
351
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
mashun1's avatar
v1  
mashun1 committed
352
353
354
355
356
357
358
		return err
	}

	_, err := io.Copy(w, strings.NewReader(s))
	return err
}

xuxzh1's avatar
init  
xuxzh1 committed
359
360
361
362
363
364
365
366
367
368
type array struct {
	size   int
	values []any
}

func (a *array) MarshalJSON() ([]byte, error) {
	return json.Marshal(a.values)
}

func readGGUFV1Array(llm *gguf, r io.Reader) (*array, error) {
mashun1's avatar
v1  
mashun1 committed
369
370
371
372
373
374
375
376
377
378
	t, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}

	n, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}

xuxzh1's avatar
init  
xuxzh1 committed
379
380
381
382
383
384
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, 0, int(n))
	}

	for i := range n {
mashun1's avatar
v1  
mashun1 committed
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
		var e any
		switch t {
		case ggufTypeUint8:
			e, err = readGGUF[uint8](llm, r)
		case ggufTypeInt8:
			e, err = readGGUF[int8](llm, r)
		case ggufTypeUint16:
			e, err = readGGUF[uint16](llm, r)
		case ggufTypeInt16:
			e, err = readGGUF[int16](llm, r)
		case ggufTypeUint32:
			e, err = readGGUF[uint32](llm, r)
		case ggufTypeInt32:
			e, err = readGGUF[int32](llm, r)
		case ggufTypeUint64:
			e, err = readGGUF[uint64](llm, r)
		case ggufTypeInt64:
			e, err = readGGUF[int64](llm, r)
		case ggufTypeFloat32:
			e, err = readGGUF[float32](llm, r)
		case ggufTypeFloat64:
			e, err = readGGUF[float64](llm, r)
		case ggufTypeBool:
			e, err = readGGUF[bool](llm, r)
		case ggufTypeString:
			e, err = readGGUFV1String(llm, r)
		default:
			return nil, fmt.Errorf("invalid array type: %d", t)
		}
		if err != nil {
			return nil, err
		}

xuxzh1's avatar
init  
xuxzh1 committed
418
419
420
		if a.values != nil {
			a.values[i] = e
		}
mashun1's avatar
v1  
mashun1 committed
421
422
	}

xuxzh1's avatar
init  
xuxzh1 committed
423
	return a, nil
mashun1's avatar
v1  
mashun1 committed
424
425
}

xuxzh1's avatar
init  
xuxzh1 committed
426
func readGGUFArray(llm *gguf, r io.Reader) (*array, error) {
mashun1's avatar
v1  
mashun1 committed
427
428
429
430
431
432
433
434
435
436
437
438
439
440
	if llm.Version == 1 {
		return readGGUFV1Array(llm, r)
	}

	t, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}

	n, err := readGGUF[uint64](llm, r)
	if err != nil {
		return nil, err
	}

xuxzh1's avatar
init  
xuxzh1 committed
441
442
443
444
445
446
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, int(n))
	}

	for i := range n {
mashun1's avatar
v1  
mashun1 committed
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
		var e any
		switch t {
		case ggufTypeUint8:
			e, err = readGGUF[uint8](llm, r)
		case ggufTypeInt8:
			e, err = readGGUF[int8](llm, r)
		case ggufTypeUint16:
			e, err = readGGUF[uint16](llm, r)
		case ggufTypeInt16:
			e, err = readGGUF[int16](llm, r)
		case ggufTypeUint32:
			e, err = readGGUF[uint32](llm, r)
		case ggufTypeInt32:
			e, err = readGGUF[int32](llm, r)
		case ggufTypeUint64:
			e, err = readGGUF[uint64](llm, r)
		case ggufTypeInt64:
			e, err = readGGUF[int64](llm, r)
		case ggufTypeFloat32:
			e, err = readGGUF[float32](llm, r)
		case ggufTypeFloat64:
			e, err = readGGUF[float64](llm, r)
		case ggufTypeBool:
			e, err = readGGUF[bool](llm, r)
		case ggufTypeString:
xuxzh1's avatar
init  
xuxzh1 committed
472
473
474
475
476
			if a.values != nil {
				e, err = readGGUFString(llm, r)
			} else {
				err = discardGGUFString(llm, r)
			}
mashun1's avatar
v1  
mashun1 committed
477
478
479
480
481
482
483
		default:
			return nil, fmt.Errorf("invalid array type: %d", t)
		}
		if err != nil {
			return nil, err
		}

xuxzh1's avatar
init  
xuxzh1 committed
484
485
486
		if a.values != nil {
			a.values[i] = e
		}
mashun1's avatar
v1  
mashun1 committed
487
488
	}

xuxzh1's avatar
init  
xuxzh1 committed
489
	return a, nil
mashun1's avatar
v1  
mashun1 committed
490
491
}

xuxzh1's avatar
init  
xuxzh1 committed
492
493
494
// writeGGUFArray writes a slice s of type E to the write with a gguf type of t
func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
	if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
mashun1's avatar
v1  
mashun1 committed
495
496
497
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
498
	if err := binary.Write(w, binary.LittleEndian, t); err != nil {
mashun1's avatar
v1  
mashun1 committed
499
500
501
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
502
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
mashun1's avatar
v1  
mashun1 committed
503
504
505
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
506
	return binary.Write(w, binary.LittleEndian, s)
mashun1's avatar
v1  
mashun1 committed
507
508
}

xuxzh1's avatar
init  
xuxzh1 committed
509
510
func WriteGGUF(ws io.WriteSeeker, kv KV, ts []Tensor) error {
	if err := binary.Write(ws, binary.LittleEndian, []byte("GGUF")); err != nil {
mashun1's avatar
v1  
mashun1 committed
511
512
513
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
514
	if err := binary.Write(ws, binary.LittleEndian, uint32(3)); err != nil {
mashun1's avatar
v1  
mashun1 committed
515
516
517
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
518
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(ts))); err != nil {
mashun1's avatar
v1  
mashun1 committed
519
520
521
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
522
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(kv))); err != nil {
mashun1's avatar
v1  
mashun1 committed
523
524
525
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
526
527
528
529
530
531
532
	keys := maps.Keys(kv)
	slices.Sort(keys)

	for _, key := range keys {
		if err := ggufWriteKV(ws, key, kv[key]); err != nil {
			return err
		}
mashun1's avatar
v1  
mashun1 committed
533
534
	}

xuxzh1's avatar
init  
xuxzh1 committed
535
536
537
538
539
540
	slices.SortFunc(ts, func(a, b Tensor) int {
		var i, j int
		if n, err := fmt.Sscanf(a.Name, "blk.%d", &i); err != nil || n != 1 {
			return cmp.Compare(a.Name, b.Name)
		} else if n, err := fmt.Sscanf(b.Name, "blk.%d", &j); err != nil || n != 1 {
			return cmp.Compare(a.Name, b.Name)
mashun1's avatar
v1  
mashun1 committed
541
542
		}

xuxzh1's avatar
init  
xuxzh1 committed
543
544
545
546
547
548
549
		return cmp.Compare(i, j)
	})

	var s uint64
	for _, t := range ts {
		t.Offset = s
		if err := ggufWriteTensorInfo(ws, t); err != nil {
mashun1's avatar
v1  
mashun1 committed
550
551
			return err
		}
xuxzh1's avatar
init  
xuxzh1 committed
552
553
		s += t.Size()
	}
mashun1's avatar
v1  
mashun1 committed
554

xuxzh1's avatar
init  
xuxzh1 committed
555
556
557
	var alignment int64 = 32
	for _, t := range ts {
		if err := ggufWriteTensor(ws, t, alignment); err != nil {
mashun1's avatar
v1  
mashun1 committed
558
559
			return err
		}
xuxzh1's avatar
init  
xuxzh1 committed
560
	}
mashun1's avatar
v1  
mashun1 committed
561

xuxzh1's avatar
init  
xuxzh1 committed
562
563
	return nil
}
mashun1's avatar
v1  
mashun1 committed
564

xuxzh1's avatar
init  
xuxzh1 committed
565
566
567
568
func ggufWriteKV(ws io.WriteSeeker, k string, v any) error {
	slog.Debug(k, "type", fmt.Sprintf("%T", v))
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(k))); err != nil {
		return err
mashun1's avatar
v1  
mashun1 committed
569
570
	}

xuxzh1's avatar
init  
xuxzh1 committed
571
572
	if err := binary.Write(ws, binary.LittleEndian, []byte(k)); err != nil {
		return err
mashun1's avatar
v1  
mashun1 committed
573
574
	}

xuxzh1's avatar
init  
xuxzh1 committed
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
	var err error
	switch v := v.(type) {
	case uint32:
		err = writeGGUF(ws, ggufTypeUint32, v)
	case float32:
		err = writeGGUF(ws, ggufTypeFloat32, v)
	case bool:
		err = writeGGUF(ws, ggufTypeBool, v)
	case string:
		err = writeGGUFString(ws, v)
	case []int32:
		err = writeGGUFArray(ws, ggufTypeInt32, v)
	case []uint32:
		err = writeGGUFArray(ws, ggufTypeUint32, v)
	case []float32:
		err = writeGGUFArray(ws, ggufTypeFloat32, v)
	case []string:
		if err := binary.Write(ws, binary.LittleEndian, ggufTypeArray); err != nil {
mashun1's avatar
v1  
mashun1 committed
593
594
595
			return err
		}

xuxzh1's avatar
init  
xuxzh1 committed
596
		if err := binary.Write(ws, binary.LittleEndian, ggufTypeString); err != nil {
mashun1's avatar
v1  
mashun1 committed
597
598
599
			return err
		}

xuxzh1's avatar
init  
xuxzh1 committed
600
		if err := binary.Write(ws, binary.LittleEndian, uint64(len(v))); err != nil {
mashun1's avatar
v1  
mashun1 committed
601
602
603
			return err
		}

xuxzh1's avatar
init  
xuxzh1 committed
604
605
		for _, e := range v {
			if err := binary.Write(ws, binary.LittleEndian, uint64(len(e))); err != nil {
mashun1's avatar
v1  
mashun1 committed
606
607
608
				return err
			}

xuxzh1's avatar
init  
xuxzh1 committed
609
610
611
			if err := binary.Write(ws, binary.LittleEndian, []byte(e)); err != nil {
				return err
			}
mashun1's avatar
v1  
mashun1 committed
612
		}
xuxzh1's avatar
init  
xuxzh1 committed
613
614
615
	default:
		return fmt.Errorf("improper type for '%s'", k)
	}
mashun1's avatar
v1  
mashun1 committed
616

xuxzh1's avatar
init  
xuxzh1 committed
617
618
619
620
621
622
623
	return err
}

func ggufWriteTensorInfo(ws io.WriteSeeker, t Tensor) error {
	slog.Debug(t.Name, "kind", t.Kind, "shape", t.Shape, "offset", t.Offset)
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(t.Name))); err != nil {
		return err
mashun1's avatar
v1  
mashun1 committed
624
625
	}

xuxzh1's avatar
init  
xuxzh1 committed
626
	if err := binary.Write(ws, binary.LittleEndian, []byte(t.Name)); err != nil {
mashun1's avatar
v1  
mashun1 committed
627
628
629
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
630
	if err := binary.Write(ws, binary.LittleEndian, uint32(len(t.Shape))); err != nil {
mashun1's avatar
v1  
mashun1 committed
631
632
633
		return err
	}

xuxzh1's avatar
init  
xuxzh1 committed
634
635
	for i := range len(t.Shape) {
		if err := binary.Write(ws, binary.LittleEndian, t.Shape[len(t.Shape)-i-1]); err != nil {
mashun1's avatar
v1  
mashun1 committed
636
637
			return err
		}
xuxzh1's avatar
init  
xuxzh1 committed
638
	}
mashun1's avatar
v1  
mashun1 committed
639

xuxzh1's avatar
init  
xuxzh1 committed
640
641
642
	if err := binary.Write(ws, binary.LittleEndian, t.Kind); err != nil {
		return err
	}
mashun1's avatar
v1  
mashun1 committed
643

xuxzh1's avatar
init  
xuxzh1 committed
644
645
646
647
648
649
650
	return binary.Write(ws, binary.LittleEndian, t.Offset)
}

func ggufWriteTensor(ws io.WriteSeeker, t Tensor, alignment int64) error {
	offset, err := ws.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
mashun1's avatar
v1  
mashun1 committed
651
652
	}

xuxzh1's avatar
init  
xuxzh1 committed
653
654
655
656
657
658
	if err := binary.Write(ws, binary.LittleEndian, bytes.Repeat([]byte{0}, int(ggufPadding(offset, alignment)))); err != nil {
		return err
	}

	_, err = t.WriteTo(ws)
	return err
mashun1's avatar
v1  
mashun1 committed
659
660
}

xuxzh1's avatar
init  
xuxzh1 committed
661
func ggufPadding(offset, align int64) int64 {
mashun1's avatar
v1  
mashun1 committed
662
663
	return (align - offset%align) % align
}