core/cli
8
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions 1

fix: handle both string and array merge formats in tokenizer

Browse source 
Gemma 3 tokenizer.json uses [["a","b"],...] format for merges
instead of the ["a b",...] format. Support both.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude 2026-02-16 02:02:55 +00:00 • committed by Snider
parent 4d3e54c81a
commit 7fc1571f93
1 changed files with 25 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
pkg/mlx/tokenizer/tokenizer.go
View file

@ -29,10 +29,10 @@ type mergePair struct {
// tokenizerJSON is the HuggingFace tokenizer.json format.
type tokenizerJSON struct {
Model struct {
Type string `json:"type"`
Vocab json.RawMessage `json:"vocab"`
Merges []string `json:"merges"`
ByteFallback bool `json:"byte_fallback"`
Type string `json:"type"`
Vocab json.RawMessage `json:"vocab"`
Merges json.RawMessage `json:"merges"`
ByteFallback bool `json:"byte_fallback"`
} `json:"model"`
AddedTokens []struct {
ID int32 `json:"id"`
@ -69,11 +69,27 @@ func Load(path string) (*Tokenizer, error) {
t.invVocab[v] = k
}
// Parse merges
for rank, merge := range tj.Model.Merges {
parts := strings.SplitN(merge, " ", 2)
if len(parts) == 2 {
t.merges = append(t.merges, mergePair{a: parts[0], b: parts[1], rank: rank})
// Parse merges — supports both ["a b", ...] and [["a","b"], ...] formats
if len(tj.Model.Merges) > 0 {
// Try array-of-strings first
var stringMerges []string
if err := json.Unmarshal(tj.Model.Merges, &stringMerges); err == nil {
for rank, merge := range stringMerges {
parts := strings.SplitN(merge, " ", 2)
if len(parts) == 2 {
t.merges = append(t.merges, mergePair{a: parts[0], b: parts[1], rank: rank})
}
}
} else {
// Try array-of-arrays: [["a","b"], ...]
var arrayMerges [][]string
if err := json.Unmarshal(tj.Model.Merges, &arrayMerges); err == nil {
for rank, pair := range arrayMerges {
if len(pair) == 2 {
t.merges = append(t.merges, mergePair{a: pair[0], b: pair[1], rank: rank})
}
}
}
}
}