ggml.go 4.6 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
	NumLayers() int64
82
83
84
85
}

type container interface {
	Name() string
Michael Yang's avatar
Michael Yang committed
86
	Decode(*readSeekOffset) (model, error)
87
88
}

Bruce MacDonald's avatar
Bruce MacDonald committed
89
type containerGGML struct{}
90
91
92
93
94

func (c *containerGGML) Name() string {
	return "ggml"
}

Michael Yang's avatar
Michael Yang committed
95
func (c *containerGGML) Decode(ro *readSeekOffset) (model, error) {
96
97
	// file contents aren't decoded
	ro.Seek(0, io.SeekEnd)
Bruce MacDonald's avatar
Bruce MacDonald committed
98
	return nil, nil
99
100
101
102
103
104
105
106
107
108
}

type containerGGMF struct {
	version uint32
}

func (c *containerGGMF) Name() string {
	return "ggmf"
}

Michael Yang's avatar
Michael Yang committed
109
func (c *containerGGMF) Decode(ro *readSeekOffset) (model, error) {
110
	var version uint32
111
	binary.Read(ro, binary.LittleEndian, &version)
112
113
114
115

	switch version {
	case 1:
	default:
Bruce MacDonald's avatar
Bruce MacDonald committed
116
		return nil, errors.New("invalid version")
117
118
119
	}

	c.version = version
120
121
122
123

	// remaining file contents aren't decoded
	ro.Seek(0, io.SeekEnd)

Bruce MacDonald's avatar
Bruce MacDonald committed
124
	return nil, nil
125
126
127
128
129
130
131
132
133
134
}

type containerGGJT struct {
	version uint32
}

func (c *containerGGJT) Name() string {
	return "ggjt"
}

Michael Yang's avatar
Michael Yang committed
135
func (c *containerGGJT) Decode(ro *readSeekOffset) (model, error) {
136
	var version uint32
137
	binary.Read(ro, binary.LittleEndian, &version)
138
139
140
141

	switch version {
	case 1, 2, 3:
	default:
Bruce MacDonald's avatar
Bruce MacDonald committed
142
		return nil, errors.New("invalid version")
143
144
145
	}

	c.version = version
Bruce MacDonald's avatar
Bruce MacDonald committed
146
147
148

	// different model types may have different layouts for hyperparameters
	var llama llamaModel
149
	binary.Read(ro, binary.LittleEndian, &llama.hyperparameters)
150
151
152
153

	// remaining file contents aren't decoded
	ro.Seek(0, io.SeekEnd)

Bruce MacDonald's avatar
Bruce MacDonald committed
154
	return &llama, nil
155
156
157
158
159
160
161
162
163
164
}

type containerLORA struct {
	version uint32
}

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

Michael Yang's avatar
Michael Yang committed
165
func (c *containerLORA) Decode(ro *readSeekOffset) (model, error) {
166
	var version uint32
167
	binary.Read(ro, binary.LittleEndian, &version)
168
169
170
171

	switch version {
	case 1:
	default:
Bruce MacDonald's avatar
Bruce MacDonald committed
172
		return nil, errors.New("invalid version")
173
174
175
	}

	c.version = version
176
177
178
179

	// remaining file contents aren't decoded
	ro.Seek(0, io.SeekEnd)

Bruce MacDonald's avatar
Bruce MacDonald committed
180
181
182
	return nil, nil
}

183
const (
Bruce MacDonald's avatar
Bruce MacDonald committed
184
	// Magic constant for `ggml` files (unversioned).
185
	FILE_MAGIC_GGML = 0x67676d6c
Bruce MacDonald's avatar
Bruce MacDonald committed
186
	// Magic constant for `ggml` files (versioned, ggmf).
187
	FILE_MAGIC_GGMF = 0x67676d66
Bruce MacDonald's avatar
Bruce MacDonald committed
188
	// Magic constant for `ggml` files (versioned, ggjt).
189
	FILE_MAGIC_GGJT = 0x67676a74
Bruce MacDonald's avatar
Bruce MacDonald committed
190
	// Magic constant for `ggla` files (LoRA adapter).
191
	FILE_MAGIC_GGLA = 0x67676C61
Bruce MacDonald's avatar
Bruce MacDonald committed
192
	// Magic constant for `gguf` files (versioned, gguf)
Michael Yang's avatar
ggufv3  
Michael Yang committed
193
194
	FILE_MAGIC_GGUF_LE = 0x46554747
	FILE_MAGIC_GGUF_BE = 0x47475546
195
196
)

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

200
201
202
203
204
205
206
	var magic uint32
	if err := binary.Read(&ro, binary.LittleEndian, &magic); err != nil {
		return nil, err
	}

	var c container
	switch magic {
207
	case FILE_MAGIC_GGML:
208
		c = &containerGGML{}
209
	case FILE_MAGIC_GGMF:
210
		c = &containerGGMF{}
211
	case FILE_MAGIC_GGJT:
212
		c = &containerGGJT{}
213
	case FILE_MAGIC_GGLA:
214
		c = &containerLORA{}
Michael Yang's avatar
ggufv3  
Michael Yang committed
215
	case FILE_MAGIC_GGUF_LE:
216
		c = &containerGGUF{bo: binary.LittleEndian}
Michael Yang's avatar
ggufv3  
Michael Yang committed
217
	case FILE_MAGIC_GGUF_BE:
218
		c = &containerGGUF{bo: binary.BigEndian}
219
220
221
222
	default:
		return nil, errors.New("invalid file magic")
	}

223
	model, err := c.Decode(&ro)
Bruce MacDonald's avatar
Bruce MacDonald committed
224
	if err != nil {
225
226
227
228
		return nil, err
	}

	// final model type
229
230
231
232
233
234
235
	return &GGML{
		container: c,
		model:     model,
		Size:      ro.offset,
	}, nil
}

Michael Yang's avatar
Michael Yang committed
236
237
type readSeekOffset struct {
	io.ReadSeeker
238
239
240
	offset int64
}

Michael Yang's avatar
Michael Yang committed
241
242
243
244
245
246
247
248
249
250
251
252
253
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)
254
	return n, err
255
}