ggml.go 3.47 KB
Newer Older
1
2
3
4
5
6
7
8
package llm

import (
	"encoding/binary"
	"errors"
	"io"
)

Michael Yang's avatar
Michael Yang committed
9
10
11
type GGML struct {
	container
	model
12
13

	Size int64
Michael Yang's avatar
Michael Yang committed
14
}
15
16

const (
Michael Yang's avatar
Michael Yang committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
	fileTypeF32 uint32 = iota
	fileTypeF16
	fileTypeQ4_0
	fileTypeQ4_1
	fileTypeQ4_1_F16
	fileTypeQ8_0 uint32 = iota + 2
	fileTypeQ5_0
	fileTypeQ5_1
	fileTypeQ2_K
	fileTypeQ3_K_S
	fileTypeQ3_K_M
	fileTypeQ3_K_L
	fileTypeQ4_K_S
	fileTypeQ4_K_M
	fileTypeQ5_K_S
	fileTypeQ5_K_M
	fileTypeQ6_K
34
35
)

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
66
67
68
69
70
71
func fileType(fileType uint32) string {
	switch fileType {
	case fileTypeF32:
		return "F32"
	case fileTypeF16:
		return "F16"
	case fileTypeQ4_0:
		return "Q4_0"
	case fileTypeQ4_1:
		return "Q4_1"
	case fileTypeQ4_1_F16:
		return "Q4_1_F16"
	case fileTypeQ8_0:
		return "Q8_0"
	case fileTypeQ5_0:
		return "Q5_0"
	case fileTypeQ5_1:
		return "Q5_1"
	case fileTypeQ2_K:
		return "Q2_K"
	case fileTypeQ3_K_S:
		return "Q3_K_S"
	case fileTypeQ3_K_M:
		return "Q3_K_M"
	case fileTypeQ3_K_L:
		return "Q3_K_L"
	case fileTypeQ4_K_S:
		return "Q4_K_S"
	case fileTypeQ4_K_M:
		return "Q4_K_M"
	case fileTypeQ5_K_S:
		return "Q5_K_S"
	case fileTypeQ5_K_M:
		return "Q5_K_M"
	case fileTypeQ6_K:
		return "Q6_K"
Michael Yang's avatar
Michael Yang committed
72
	default:
Michael Yang's avatar
Michael Yang committed
73
		return "unknown"
Michael Yang's avatar
Michael Yang committed
74
75
	}
}
76

Michael Yang's avatar
Michael Yang committed
77
type model interface {
Michael Yang's avatar
Michael Yang committed
78
79
80
	ModelFamily() string
	ModelType() string
	FileType() string
81
82
83
84
85
	NumLayers() uint32
	NumGQA() uint32
	NumEmbed() uint32
	NumHead() uint32
	NumHeadKv() uint32
Michael Yang's avatar
Michael Yang committed
86
	NumCtx() uint32
87
88
89
90
}

type container interface {
	Name() string
Michael Yang's avatar
Michael Yang committed
91
	Decode(*readSeekOffset) (model, error)
92
93
94
95
96
97
98
99
100
101
}

type containerLORA struct {
	version uint32
}

func (c *containerLORA) Name() string {
	return "ggla"
}

Michael Yang's avatar
Michael Yang committed
102
func (c *containerLORA) Decode(rso *readSeekOffset) (model, error) {
103
	var version uint32
Michael Yang's avatar
Michael Yang committed
104
	binary.Read(rso, binary.LittleEndian, &version)
105
106
107
108

	switch version {
	case 1:
	default:
Bruce MacDonald's avatar
Bruce MacDonald committed
109
		return nil, errors.New("invalid version")
110
111
112
	}

	c.version = version
113
114

	// remaining file contents aren't decoded
Michael Yang's avatar
Michael Yang committed
115
	rso.Seek(0, io.SeekEnd)
116

Bruce MacDonald's avatar
Bruce MacDonald committed
117
118
119
	return nil, nil
}

120
const (
Bruce MacDonald's avatar
Bruce MacDonald committed
121
	// Magic constant for `ggml` files (unversioned).
122
	FILE_MAGIC_GGML = 0x67676d6c
Bruce MacDonald's avatar
Bruce MacDonald committed
123
	// Magic constant for `ggml` files (versioned, ggmf).
124
	FILE_MAGIC_GGMF = 0x67676d66
Bruce MacDonald's avatar
Bruce MacDonald committed
125
	// Magic constant for `ggml` files (versioned, ggjt).
126
	FILE_MAGIC_GGJT = 0x67676a74
Bruce MacDonald's avatar
Bruce MacDonald committed
127
	// Magic constant for `ggla` files (LoRA adapter).
128
	FILE_MAGIC_GGLA = 0x67676C61
Bruce MacDonald's avatar
Bruce MacDonald committed
129
	// Magic constant for `gguf` files (versioned, gguf)
Michael Yang's avatar
ggufv3  
Michael Yang committed
130
131
	FILE_MAGIC_GGUF_LE = 0x46554747
	FILE_MAGIC_GGUF_BE = 0x47475546
132
133
)

Bruce MacDonald's avatar
Bruce MacDonald committed
134
135
var ErrUnsupportedFormat = errors.New("unsupported model format")

Michael Yang's avatar
Michael Yang committed
136
137
func DecodeGGML(r io.ReadSeeker) (*GGML, error) {
	ro := readSeekOffset{ReadSeeker: r}
138

139
140
141
142
143
144
145
	var magic uint32
	if err := binary.Read(&ro, binary.LittleEndian, &magic); err != nil {
		return nil, err
	}

	var c container
	switch magic {
Bruce MacDonald's avatar
Bruce MacDonald committed
146
147
	case FILE_MAGIC_GGML, FILE_MAGIC_GGMF, FILE_MAGIC_GGJT:
		return nil, ErrUnsupportedFormat
148
	case FILE_MAGIC_GGLA:
149
		c = &containerLORA{}
Michael Yang's avatar
ggufv3  
Michael Yang committed
150
	case FILE_MAGIC_GGUF_LE:
151
		c = &containerGGUF{bo: binary.LittleEndian}
Michael Yang's avatar
ggufv3  
Michael Yang committed
152
	case FILE_MAGIC_GGUF_BE:
153
		c = &containerGGUF{bo: binary.BigEndian}
154
155
156
157
	default:
		return nil, errors.New("invalid file magic")
	}

158
	model, err := c.Decode(&ro)
Bruce MacDonald's avatar
Bruce MacDonald committed
159
	if err != nil {
160
161
162
163
		return nil, err
	}

	// final model type
164
165
166
167
168
169
170
	return &GGML{
		container: c,
		model:     model,
		Size:      ro.offset,
	}, nil
}

Michael Yang's avatar
Michael Yang committed
171
172
type readSeekOffset struct {
	io.ReadSeeker
173
174
175
	offset int64
}

Michael Yang's avatar
Michael Yang committed
176
177
178
179
180
181
182
183
184
185
186
187
188
func (rso *readSeekOffset) Seek(offset int64, whence int) (int64, error) {
	offset, err := rso.ReadSeeker.Seek(offset, whence)
	if err != nil {
		return 0, err
	}

	rso.offset = offset
	return offset, nil
}

func (rso *readSeekOffset) Read(p []byte) (int, error) {
	n, err := rso.ReadSeeker.Read(p)
	rso.offset += int64(n)
189
	return n, err
190
}