tokenizer.go 5.87 KB
Newer Older
Patrick Devine's avatar
Patrick Devine committed
1
2
3
package convert

import (
Michael Yang's avatar
Michael Yang committed
4
	"crypto/sha256"
Michael Yang's avatar
Michael Yang committed
5
	"encoding/hex"
Patrick Devine's avatar
Patrick Devine committed
6
	"encoding/json"
Michael Yang's avatar
Michael Yang committed
7
	"errors"
Michael Yang's avatar
Michael Yang committed
8
	"fmt"
9
	"io/fs"
Michael Yang's avatar
Michael Yang committed
10
	"log/slog"
Patrick Devine's avatar
Patrick Devine committed
11
	"os"
Michael Yang's avatar
Michael Yang committed
12
	"slices"
Michael Yang's avatar
bert  
Michael Yang committed
13
14

	"golang.org/x/exp/maps"
Michael Yang's avatar
Michael Yang committed
15
)
Michael Yang's avatar
Michael Yang committed
16

Michael Yang's avatar
Michael Yang committed
17
18
19
20
21
22
23
24
const (
	_ int32 = iota
	tokenTypeNormal
	tokenTypeUnknown
	tokenTypeControl
	tokenTypeUserDefined
	tokenTypeUnused
	tokenTypeByte
Patrick Devine's avatar
Patrick Devine committed
25
26
27
)

type Tokenizer struct {
Michael Yang's avatar
Michael Yang committed
28
29
30
31
32
33
34
35
	*Vocabulary
	SpecialVocabulary []*SpecialVocabulary
	Merges            []string

	Pre      string
	Template string
}

36
37
func parseTokenizer(fsys fs.FS, specialTokenTypes []string) (*Tokenizer, error) {
	v, err := parseVocabulary(fsys)
Michael Yang's avatar
Michael Yang committed
38
39
40
41
42
43
44
45
46
47
	if err != nil {
		return nil, err
	}

	t := &Tokenizer{
		Vocabulary: v,
		Pre:        "default",
	}

	addedTokens := make(map[string]token)
48
	if f, err := fsys.Open("tokenizer.json"); errors.Is(err, os.ErrNotExist) {
Michael Yang's avatar
Michael Yang committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
	} else if err != nil {
		return nil, err
	} else {
		defer f.Close()

		var tt tokenizer
		if err := json.NewDecoder(f).Decode(&tt); err != nil {
			return nil, err
		}

		for _, t := range tt.AddedTokens {
			addedTokens[t.Content] = t
		}

		t.Merges = tt.Model.Merges

		sha256sum := sha256.New()
		for _, pt := range tt.PreTokenizer.PreTokenizers {
			switch pt.Type {
			case "Split":
				if pt.Pattern.Regex != "" {
Michael Yang's avatar
Michael Yang committed
70
71
					// create a checksum of all Split pretokenizers which should be sufficient
					// to identify the pretokenizer
Michael Yang's avatar
Michael Yang committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
					sha256sum.Write([]byte(pt.Pattern.Regex))
				}
			}
		}

		switch digest := hex.EncodeToString(sha256sum.Sum(nil)); digest {
		case "d98f9631be1e9607a9848c26c1f9eac1aa9fc21ac6ba82a2fc0741af9780a48f":
			t.Pre = "llama-bpe"
		case "03df5c5863ad70781dcfdef491ead25140f895fe8010964be0daefe27be32b02":
			t.Pre = "deepseek-llm"
		case "21cde974d587f0d54dc8d56b183cc1e6239600172035c68fbd6d4b9f8da0576e":
			t.Pre = "deepseek-coder"
		case "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855":
			// noop, empty pretokenizer
		default:
			slog.Warn("unknown pretokenizer, using default", "digest", digest)
		}
	}

91
	if f, err := fsys.Open("tokenizer_config.json"); errors.Is(err, os.ErrNotExist) {
Michael Yang's avatar
Michael Yang committed
92
93
94
95
96
97
98
99
100
101
102
	} else if err != nil {
		return nil, err
	} else {
		defer f.Close()

		var p map[string]json.RawMessage
		if err := json.NewDecoder(f).Decode(&p); err != nil {
			return nil, err
		}

		if template, ok := p["chat_template"]; ok {
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
			var s []struct {
				Name     string `json:"name"`
				Template string `json:"template"`
			}
			if err := json.Unmarshal(template, &t.Template); err == nil {
				// noop
			} else if err := json.Unmarshal(template, &s); err == nil {
				for _, e := range s {
					if e.Name == "default" {
						t.Template = e.Template
						break
					}
				}
			} else {
				return nil, fmt.Errorf("invalid chat_template: %w", err)
Michael Yang's avatar
Michael Yang committed
118
119
120
			}
		}

Michael Yang's avatar
Michael Yang committed
121
		for _, st := range specialTokenTypes {
Michael Yang's avatar
Michael Yang committed
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
			sv := SpecialVocabulary{Type: st}
			if bts, ok := p[fmt.Sprintf("add_%s_token", st)]; ok {
				if err := json.Unmarshal(bts, &sv.AddToken); err != nil {
					return nil, err
				}
			}

			if bts, ok := p[fmt.Sprintf("%s_token", st)]; ok {
				var content string
				if err := json.Unmarshal(bts, &content); err != nil {
					var mm map[string]any
					if err := json.Unmarshal(bts, &mm); err != nil {
						continue
					}

					content, ok = mm["content"].(string)
					if !ok {
						continue
					}
				}

				sv.Content = content
			}

			if id, ok := addedTokens[sv.Content]; ok {
				sv.ID = id.ID
				t.SpecialVocabulary = append(t.SpecialVocabulary, &sv)
			}
		}
	}

	return t, nil
}

type tokenizer struct {
	Version     string  `json:"version"`
	AddedTokens []token `json:"added_tokens"`
	Model       struct {
		Type   string         `json:"type"`
		Vocab  map[string]int `json:"vocab"`
		Merges []string       `json:"merges"`
	} `json:"model"`
Michael Yang's avatar
Michael Yang committed
164
165

	PreTokenizer struct {
166
		PreTokenizers []struct {
Michael Yang's avatar
Michael Yang committed
167
168
169
170
171
172
			Type    string `json:"type"`
			Pattern struct {
				Regex string `json:"Regex"`
			} `json:"pattern"`
		} `json:"pretokenizers"`
	} `json:"pre_tokenizer"`
Patrick Devine's avatar
Patrick Devine committed
173
174
}

Michael Yang's avatar
Michael Yang committed
175
type token struct {
Patrick Devine's avatar
Patrick Devine committed
176
177
178
179
180
181
	ID          int    `json:"id"`
	Content     string `json:"content"`
	Special     bool   `json:"special"`
	UserDefined bool
}

Michael Yang's avatar
Michael Yang committed
182
183
184
185
186
type Vocabulary struct {
	Model  string
	Tokens []string
	Scores []float32
	Types  []int32
Michael Yang's avatar
Michael Yang committed
187
}
Patrick Devine's avatar
Patrick Devine committed
188

189
190
func parseVocabularyFromTokenizer(fsys fs.FS) (*Vocabulary, error) {
	f, err := fsys.Open("tokenizer.json")
Patrick Devine's avatar
Patrick Devine committed
191
	if err != nil {
Michael Yang's avatar
Michael Yang committed
192
		return nil, err
Patrick Devine's avatar
Patrick Devine committed
193
194
195
	}
	defer f.Close()

Michael Yang's avatar
Michael Yang committed
196
	var t tokenizer
Michael Yang's avatar
Michael Yang committed
197
	if err := json.NewDecoder(f).Decode(&t); err != nil {
Michael Yang's avatar
Michael Yang committed
198
		return nil, err
Patrick Devine's avatar
Patrick Devine committed
199
200
	}

Michael Yang's avatar
bert  
Michael Yang committed
201
	tokens := make(map[int]token, len(t.Model.Vocab))
Michael Yang's avatar
Michael Yang committed
202
	for k, v := range t.Model.Vocab {
Michael Yang's avatar
bert  
Michael Yang committed
203
		tokens[v] = token{
Michael Yang's avatar
Michael Yang committed
204
205
			ID:      v,
			Content: k,
Michael Yang's avatar
bert  
Michael Yang committed
206
		}
Patrick Devine's avatar
Patrick Devine committed
207
208
	}

Michael Yang's avatar
bert  
Michael Yang committed
209
210
211
	for _, token := range t.AddedTokens {
		token.UserDefined = true
		tokens[token.ID] = token
Michael Yang's avatar
Michael Yang committed
212
	}
Patrick Devine's avatar
Patrick Devine committed
213

Michael Yang's avatar
bert  
Michael Yang committed
214
215
	keys := maps.Keys(tokens)
	slices.Sort(keys)
Michael Yang's avatar
Michael Yang committed
216
217

	v := Vocabulary{Model: "gpt2"}
Michael Yang's avatar
bert  
Michael Yang committed
218
219
220
221
	for _, k := range keys {
		token := tokens[k]
		v.Tokens = append(v.Tokens, token.Content)
		v.Scores = append(v.Scores, float32(token.ID))
Michael Yang's avatar
Michael Yang committed
222
223

		switch {
Michael Yang's avatar
bert  
Michael Yang committed
224
		case token.Special:
Michael Yang's avatar
Michael Yang committed
225
			v.Types = append(v.Types, tokenTypeControl)
Michael Yang's avatar
bert  
Michael Yang committed
226
		case token.UserDefined:
Michael Yang's avatar
Michael Yang committed
227
228
229
			v.Types = append(v.Types, tokenTypeUserDefined)
		default:
			v.Types = append(v.Types, tokenTypeNormal)
Michael Yang's avatar
Michael Yang committed
230
		}
Patrick Devine's avatar
Patrick Devine committed
231
232
	}

Michael Yang's avatar
Michael Yang committed
233
234
235
	return &v, nil
}

236
func parseVocabulary(fsys fs.FS) (*Vocabulary, error) {
Michael Yang's avatar
Michael Yang committed
237
238
239
240
241
242
	patterns := []struct {
		Pattern string
		Func    func(fs.FS) (*Vocabulary, error)
	}{
		{"tokenizer.model", parseSentencePiece},
		{"tokenizer.json", parseVocabularyFromTokenizer},
Michael Yang's avatar
Michael Yang committed
243
244
	}

Michael Yang's avatar
Michael Yang committed
245
246
	for _, pattern := range patterns {
		if _, err := fs.Stat(fsys, pattern.Pattern); errors.Is(err, os.ErrNotExist) {
Michael Yang's avatar
Michael Yang committed
247
248
			continue
		} else if err != nil {
Michael Yang's avatar
Michael Yang committed
249
250
251
			return nil, err
		}

Michael Yang's avatar
Michael Yang committed
252
		return pattern.Func(fsys)
Michael Yang's avatar
Michael Yang committed
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
	}

	return nil, errors.New("unknown tensor format")
}

type SpecialVocabulary struct {
	Type     string
	ID       int
	Content  string
	AddToken bool
}

func (sv SpecialVocabulary) Key() string {
	switch t := sv.Type; t {
	case "bos", "eos", "cls", "mask":
		return t
	case "unk":
		return "unknown"
	case "sep":
		//nolint:misspell // this is an upstream typo
		return "seperator"
	case "pad":
		return "padding"
Patrick Devine's avatar
Patrick Devine committed
276
277
	}

Michael Yang's avatar
Michael Yang committed
278
	panic("unknown special vocabulary type")
Patrick Devine's avatar
Patrick Devine committed
279
}