1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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
|
package llmc
import (
"encoding/binary"
"fmt"
"io"
"os"
"golang.org/x/exp/constraints"
)
type ModelDesc string
const (
GPT2 ModelDesc = "gpt-2"
LLaMA3 ModelDesc = "llama-3"
)
type modelInfo struct {
magic int32
version int32
wide bool // true: uint32 tokens; false: uint16 tokens
}
var headerInfo = map[ModelDesc]modelInfo{
GPT2: {magic: 20240520, version: 1, wide: false},
LLaMA3: {magic: 20240801, version: 7, wide: true},
}
type DataFile[T constraints.Integer] struct {
Model ModelDesc
Tokens []T // token IDs
}
func Deserialize[T constraints.Integer](name string, dst *DataFile[T]) (int, error) {
var file *os.File
if f, err := os.Open(name); err != nil {
return 0, err
} else {
file = f
}
defer file.Close()
var header [headerWords * 4]byte
if _, err := io.ReadFull(file, header[:]); err != nil {
return 0, err
}
magic := int32(binary.LittleEndian.Uint32(header[0:]))
version := int32(binary.LittleEndian.Uint32(header[4:]))
numToks := int32(binary.LittleEndian.Uint32(header[8:]))
var model ModelDesc
var info modelInfo
for m, mi := range headerInfo {
if mi.magic == magic && mi.version == version {
model, info = m, mi
break
}
}
if model == "" {
return 0, fmt.Errorf("unknown magic:version %d:%d", magic, version)
}
tokens := make([]T, numToks)
if info.wide {
tokensU32 := make([]uint32, numToks)
if err := binary.Read(file, binary.LittleEndian, tokensU32); err != nil {
return 0, err
}
for i, t := range tokensU32 {
tokens[i] = T(t)
}
} else {
tokensU16 := make([]uint16, numToks)
if err := binary.Read(file, binary.LittleEndian, tokensU16); err != nil {
return 0, err
}
for i, t := range tokensU16 {
tokens[i] = T(t)
}
}
dst.Model = model
dst.Tokens = tokens
return int(numToks), nil
}
func Serialize[T constraints.Integer](src *DataFile[T], name string) (int, error) {
info, ok := headerInfo[src.Model]
if !ok {
return 0, fmt.Errorf("unknown model descriptor %q", src.Model)
}
h := Header{
Magic: info.magic,
Version: info.version,
NumToks: int32(len(src.Tokens)),
}
var file *os.File
if f, err := os.Create(name); err != nil {
return 0, err
} else {
file = f
}
defer file.Close()
header := h.encode()
if _, err := file.Write(header[:]); err != nil {
return 0, err
}
if info.wide {
tokensU32 := make([]uint32, len(src.Tokens))
for i, t := range src.Tokens {
tokensU32[i] = uint32(t)
}
if err := binary.Write(file, binary.LittleEndian, tokensU32); err != nil {
return 0, err
}
} else {
tokensU16 := make([]uint16, len(src.Tokens))
for i, t := range src.Tokens {
tokensU16[i] = uint16(t)
}
if err := binary.Write(file, binary.LittleEndian, tokensU16); err != nil {
return 0, err
}
}
return len(src.Tokens), nil
}
|
