diff --git a/x/tokenizer/tokenizer_load.go b/x/tokenizer/tokenizer_load.go index c0b6bc25e8e..38a6500f9c5 100644 --- a/x/tokenizer/tokenizer_load.go +++ b/x/tokenizer/tokenizer_load.go @@ -8,12 +8,41 @@ import ( "strings" ) +const maxTokenizerVocabularySize = 1 << 20 + +type addedToken struct { + ID int32 `json:"id"` + Content string `json:"content"` + Special bool `json:"special"` +} + +func validateTokenizerRecordCount(baseCount, addedCount int) error { + if baseCount > maxTokenizerVocabularySize || addedCount > maxTokenizerVocabularySize-baseCount { + return fmt.Errorf("tokenizer has too many vocabulary records (maximum %d)", maxTokenizerVocabularySize) + } + return nil +} + +func validateTokenizerID(id int32) error { + if id < 0 { + return fmt.Errorf("tokenizer ID %d must not be negative", id) + } + if id >= maxTokenizerVocabularySize { + return fmt.Errorf("tokenizer ID %d exceeds maximum %d", id, maxTokenizerVocabularySize-1) + } + return nil +} + // TokenizerConfig holds optional configuration data that can be passed to LoadFromBytesWithConfig. type TokenizerConfig struct { TokenizerConfigJSON []byte // tokenizer_config.json content GenerationConfigJSON []byte // generation_config.json content SpecialTokensMapJSON []byte // special_tokens_map.json content ConfigJSON []byte // config.json content + // AddedTokenIDLimit is an optional exclusive upper bound for added token IDs. + // When positive, added tokens with negative IDs or IDs at or above the limit + // are ignored. A zero value disables filtering. + AddedTokenIDLimit int32 } // LoadFromBytes loads a tokenizer from tokenizer.json bytes. @@ -21,13 +50,17 @@ type TokenizerConfig struct { // Note: This won't load special token config from companion files. Use LoadFromBytesWithConfig // to provide tokenizer_config.json data for proper PAD/EOS token loading. func LoadFromBytes(data []byte) (*Tokenizer, error) { - return loadFromTokenizerJSON(data) + return loadFromTokenizerJSON(data, 0) } // LoadFromBytesWithConfig loads a tokenizer from tokenizer.json bytes with additional config files. // This is useful when loading from blob storage where companion config files are also blobs. func LoadFromBytesWithConfig(data []byte, config *TokenizerConfig) (*Tokenizer, error) { - t, err := loadFromTokenizerJSON(data) + var addedTokenIDLimit int32 + if config != nil { + addedTokenIDLimit = config.AddedTokenIDLimit + } + t, err := loadFromTokenizerJSON(data, addedTokenIDLimit) if err != nil { return nil, err } @@ -43,7 +76,11 @@ func LoadFromBytesWithConfig(data []byte, config *TokenizerConfig) (*Tokenizer, } // loadFromTokenizerJSON parses tokenizer.json content from bytes. -func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { +func loadFromTokenizerJSON(data []byte, addedTokenIDLimit int32) (*Tokenizer, error) { + if addedTokenIDLimit < 0 { + return nil, fmt.Errorf("added token ID limit must not be negative: %d", addedTokenIDLimit) + } + var raw struct { Model struct { Type string `json:"type"` // "BPE" @@ -52,11 +89,7 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { } `json:"model"` PreTokenizer json.RawMessage `json:"pre_tokenizer"` Decoder json.RawMessage `json:"decoder"` - AddedTokens []struct { - ID int32 `json:"id"` - Content string `json:"content"` - Special bool `json:"special"` - } `json:"added_tokens"` + AddedTokens []addedToken `json:"added_tokens"` } if err := json.Unmarshal(data, &raw); err != nil { @@ -67,6 +100,95 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { if raw.Model.Type != "BPE" { return nil, fmt.Errorf("unsupported tokenizer type: %s", raw.Model.Type) } + if addedTokenIDLimit > 0 { + filtered := raw.AddedTokens[:0] + for _, tok := range raw.AddedTokens { + if tok.ID >= 0 && tok.ID < addedTokenIDLimit { + filtered = append(filtered, tok) + } + } + raw.AddedTokens = filtered + } + + if err := validateTokenizerRecordCount(len(raw.Model.Vocab), len(raw.AddedTokens)); err != nil { + return nil, fmt.Errorf("invalid tokenizer vocabulary: %w", err) + } + + baseByID := make(map[int32]string, len(raw.Model.Vocab)) + maxID := int32(-1) + // Select canonical failures instead of returning from randomized map traversal. + // Range errors take precedence, and the lowest numeric ID wins within each class. + invalidBaseID := int32(0) + hasInvalidBaseID := false + duplicateBaseID := int32(0) + hasDuplicateBaseID := false + for token, id := range raw.Model.Vocab { + if err := validateTokenizerID(id); err != nil { + if !hasInvalidBaseID || id < invalidBaseID { + invalidBaseID = id + hasInvalidBaseID = true + } + continue + } + if _, ok := baseByID[id]; ok { + if !hasDuplicateBaseID || id < duplicateBaseID { + duplicateBaseID = id + hasDuplicateBaseID = true + } + continue + } + baseByID[id] = token + if id > maxID { + maxID = id + } + } + if hasInvalidBaseID { + return nil, fmt.Errorf("invalid base token ID: %w", validateTokenizerID(invalidBaseID)) + } + if hasDuplicateBaseID { + return nil, fmt.Errorf("duplicate base token ID %d", duplicateBaseID) + } + + addedByID := make(map[int32]string, len(raw.AddedTokens)) + addedByContent := make(map[string]int32, len(raw.AddedTokens)) + for _, tok := range raw.AddedTokens { + if err := validateTokenizerID(tok.ID); err != nil { + return nil, fmt.Errorf("invalid added token %q: %w", tok.Content, err) + } + if _, ok := addedByID[tok.ID]; ok { + return nil, fmt.Errorf("duplicate added token ID %d", tok.ID) + } + if previousID, ok := addedByContent[tok.Content]; ok { + first, second := previousID, tok.ID + if first > second { + first, second = second, first + } + return nil, fmt.Errorf("duplicate added token content %q with IDs %d and %d", tok.Content, first, second) + } + addedByID[tok.ID] = tok.Content + addedByContent[tok.Content] = tok.ID + if tok.ID > maxID { + maxID = tok.ID + } + } + + for _, tok := range raw.AddedTokens { + if baseContent, ok := baseByID[tok.ID]; ok && baseContent != tok.Content { + return nil, fmt.Errorf("token ID %d has conflicting base and added content", tok.ID) + } + if baseID, ok := raw.Model.Vocab[tok.Content]; ok && baseID != tok.ID { + first, second := baseID, tok.ID + if first > second { + first, second = second, first + } + return nil, fmt.Errorf("token content %q has conflicting base and added IDs %d and %d", tok.Content, first, second) + } + } + + valuesLen := 0 + if maxID >= 0 { + valuesLen = int(maxID) + 1 + } // Parse merges - can be []string (Llama) or [][]string (GPT-OSS). var mergesStrings []string @@ -91,7 +213,7 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // Build tokenizer t := &Tokenizer{ vocab: &Vocabulary{ - Values: make([]string, len(raw.Model.Vocab)), + Values: make([]string, valuesLen), Reverse: raw.Model.Vocab, Merges: make(map[string]int, len(mergesStrings)), BOS: -1, @@ -102,11 +224,6 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // Build values array for token, id := range raw.Model.Vocab { - if int(id) >= len(t.vocab.Values) { - newValues := make([]string, id+1) - copy(newValues, t.vocab.Values) - t.vocab.Values = newValues - } t.vocab.Values[id] = token } @@ -121,11 +238,6 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // if it's a "truly special" token like BOS/EOS/PAD, but for tokenization we need // to treat all added_tokens as special to match HuggingFace behavior. for _, tok := range raw.AddedTokens { - if int(tok.ID) >= len(t.vocab.Values) { - newValues := make([]string, tok.ID+1) - copy(newValues, t.vocab.Values) - t.vocab.Values = newValues - } t.vocab.Values[tok.ID] = tok.Content t.specialTokens[tok.Content] = tok.ID // Add ALL added_tokens to special tokens } diff --git a/x/tokenizer/tokenizer_load_test.go b/x/tokenizer/tokenizer_load_test.go index c0d52e1a459..f1dbb3a216e 100644 --- a/x/tokenizer/tokenizer_load_test.go +++ b/x/tokenizer/tokenizer_load_test.go @@ -2,10 +2,384 @@ package tokenizer import ( "encoding/json" + "fmt" + "math" "strings" "testing" ) +func tokenizerJSON(vocab, addedTokens string) []byte { + return []byte(fmt.Sprintf(`{"model":{"type":"BPE","vocab":%s,"merges":[]},"added_tokens":%s}`, vocab, addedTokens)) +} + +func TestValidateTokenizerID(t *testing.T) { + tests := []struct { + name string + id int32 + want string + }{ + {name: "negative", id: -1, want: "must not be negative"}, + {name: "minimum int32", id: math.MinInt32, want: "must not be negative"}, + {name: "maximum int32", id: math.MaxInt32, want: "exceeds maximum"}, + {name: "exclusive cap", id: maxTokenizerVocabularySize, want: "exceeds maximum"}, + {name: "inclusive upper ID", id: maxTokenizerVocabularySize - 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateTokenizerID(tt.id) + if tt.want == "" { + if err != nil { + t.Fatalf("validateTokenizerID(%d): %v", tt.id, err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("validateTokenizerID(%d) error = %v, want containing %q", tt.id, err, tt.want) + } + }) + } +} + +func TestValidateTokenizerRecordCount(t *testing.T) { + if err := validateTokenizerRecordCount(maxTokenizerVocabularySize-1, 1); err != nil { + t.Fatalf("maximum record count rejected: %v", err) + } + if err := validateTokenizerRecordCount(maxTokenizerVocabularySize, 1); err == nil { + t.Fatal("expected aggregate record count above cap to fail") + } +} + +func TestLoadFromBytesRejectsInvalidTokenizerIDs(t *testing.T) { + tests := []struct { + name string + data []byte + wants []string + }{ + {name: "negative base", data: tokenizerJSON(`{"a":-1}`, `[]`), wants: []string{"invalid base token", "tokenizer ID -1"}}, + {name: "minimum base", data: tokenizerJSON(fmt.Sprintf(`{"a":%d}`, int64(math.MinInt32)), `[]`), wants: []string{"invalid base token", fmt.Sprint(math.MinInt32)}}, + {name: "negative added", data: tokenizerJSON(`{}`, `[{"id":-1,"content":"a"}]`), wants: []string{"invalid added token", "tokenizer ID -1"}}, + {name: "minimum added", data: tokenizerJSON(`{}`, fmt.Sprintf(`[{"id":%d,"content":"a"}]`, int64(math.MinInt32))), wants: []string{"invalid added token", fmt.Sprint(math.MinInt32)}}, + {name: "base at exclusive cap", data: tokenizerJSON(fmt.Sprintf(`{"a":%d}`, maxTokenizerVocabularySize), `[]`), wants: []string{"invalid base token", "exceeds maximum"}}, + {name: "added at exclusive cap", data: tokenizerJSON(`{}`, fmt.Sprintf(`[{"id":%d,"content":"a"}]`, maxTokenizerVocabularySize)), wants: []string{"invalid added token", "exceeds maximum"}}, + {name: "above int32", data: tokenizerJSON(`{"a":2147483648}`, `[]`), wants: []string{"failed to parse tokenizer", "cannot unmarshal number"}}, + {name: "below int32", data: tokenizerJSON(`{"a":-2147483649}`, `[]`), wants: []string{"failed to parse tokenizer", "cannot unmarshal number"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := LoadFromBytes(tt.data) + if err == nil { + t.Fatal("expected tokenizer load to fail") + } + for _, want := range tt.wants { + if !strings.Contains(err.Error(), want) { + t.Fatalf("error = %v, want containing %q", err, want) + } + } + }) + } +} + +func TestLoadFromBytesWithConfigFiltersAddedTokenIDs(t *testing.T) { + data := tokenizerJSON(`{"base":0}`, `[ + {"id":-1,"content":"negative"}, + {"id":1,"content":"kept"}, + {"id":2,"content":"at-limit"}, + {"id":3,"content":"above-limit"} + ]`) + tok, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + + if got := tok.VocabSize(); got != 2 { + t.Fatalf("VocabSize() = %d, want 2", got) + } + if got := tok.Decode([]int32{0, 1}); got != "basekept" { + t.Fatalf("Decode([0 1]) = %q, want %q", got, "basekept") + } + if id, ok := tok.GetSpecialToken("kept"); !ok || id != 1 { + t.Fatalf("GetSpecialToken(kept) = (%d, %v), want (1, true)", id, ok) + } + for _, content := range []string{"negative", "at-limit", "above-limit"} { + if id, ok := tok.GetSpecialToken(content); ok { + t.Fatalf("GetSpecialToken(%q) = (%d, true), want absent", content, id) + } + } +} + +func TestLoadFromBytesWithConfigAddedTokenIDLimitDoesNotFilterBaseVocabulary(t *testing.T) { + data := tokenizerJSON(`{"base":3}`, `[ + {"id":1,"content":"added"}, + {"id":3,"content":"filtered-added"} + ]`) + tok, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + + if got := tok.VocabSize(); got != 4 { + t.Fatalf("VocabSize() = %d, want 4", got) + } + if got := tok.Decode([]int32{3, 1}); got != "baseadded" { + t.Fatalf("Decode([3 1]) = %q, want %q", got, "baseadded") + } + if _, ok := tok.GetSpecialToken("filtered-added"); ok { + t.Fatal("added token at the limit was retained") + } +} + +func TestLoadFromBytesWithConfigAddedTokenIDLimitDisabled(t *testing.T) { + data := tokenizerJSON(`{}`, `[{"id":-1,"content":"negative"}]`) + for name, config := range map[string]*TokenizerConfig{ + "nil config": nil, + "zero limit": {}, + } { + t.Run(name, func(t *testing.T) { + _, err := LoadFromBytesWithConfig(data, config) + if err == nil || !strings.Contains(err.Error(), "invalid added token") { + t.Fatalf("error = %v, want invalid added token", err) + } + }) + } +} + +func TestLoadFromBytesWithConfigRejectsNegativeAddedTokenIDLimit(t *testing.T) { + _, err := LoadFromBytesWithConfig(tokenizerJSON(`{}`, `[]`), &TokenizerConfig{AddedTokenIDLimit: -1}) + if err == nil || err.Error() != "added token ID limit must not be negative: -1" { + t.Fatalf("error = %v, want negative limit error", err) + } +} + +func TestLoadFromBytesWithConfigValidatesOnlyRetainedAddedTokens(t *testing.T) { + tests := []struct { + name string + data []byte + }{ + { + name: "duplicate IDs are filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"first"}, + {"id":2,"content":"second"}, + {"id":1,"content":"kept"} + ]`), + }, + { + name: "duplicate content is filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"duplicate"}, + {"id":3,"content":"duplicate"}, + {"id":1,"content":"kept"} + ]`), + }, + { + name: "base collision is filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"base"}, + {"id":1,"content":"kept"} + ]`), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tok, err := LoadFromBytesWithConfig(tt.data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + if got := tok.VocabSize(); got != 2 { + t.Fatalf("VocabSize() = %d, want 2", got) + } + }) + } +} + +func TestLoadFromBytesWithConfigRejectsRetainedAddedTokenCollision(t *testing.T) { + data := tokenizerJSON(`{"base":0}`, `[ + {"id":1,"content":"first"}, + {"id":1,"content":"second"}, + {"id":2,"content":"filtered"} + ]`) + _, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err == nil || err.Error() != "duplicate added token ID 1" { + t.Fatalf("error = %v, want retained-token collision", err) + } +} + +func TestLoadFromBytesSupportsBoundedSparseIDs(t *testing.T) { + tests := []struct { + name string + data []byte + wantSize int + decodeID int32 + wantDecoded string + special string + wantSpecial int32 + }{ + { + name: "inclusive upper base ID", + data: tokenizerJSON(fmt.Sprintf(`{"edge":%d}`, maxTokenizerVocabularySize-1), `[]`), + wantSize: maxTokenizerVocabularySize, + decodeID: maxTokenizerVocabularySize - 1, + wantDecoded: "edge", + }, + {name: "sparse base", data: tokenizerJSON(`{"base":1000}`, `[]`), wantSize: 1001, decodeID: 1000, wantDecoded: "base"}, + {name: "sparse added", data: tokenizerJSON(`{}`, `[{"id":1000,"content":"added"}]`), wantSize: 1001, decodeID: 1000, wantDecoded: "added", special: "added", wantSpecial: 1000}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tok, err := LoadFromBytes(tt.data) + if err != nil { + t.Fatal(err) + } + if got := tok.VocabSize(); got != tt.wantSize { + t.Fatalf("VocabSize() = %d, want %d", got, tt.wantSize) + } + if got := tok.Decode([]int32{tt.decodeID}); got != tt.wantDecoded { + t.Fatalf("Decode(%d) = %q, want %q", tt.decodeID, got, tt.wantDecoded) + } + if tt.special != "" { + id, ok := tok.GetSpecialToken(tt.special) + if !ok || id != tt.wantSpecial { + t.Fatalf("GetSpecialToken(%q) = (%d, %v), want (%d, true)", tt.special, id, ok, tt.wantSpecial) + } + } + }) + } +} + +func TestLoadFromBytesRejectsTokenizerCollisions(t *testing.T) { + tests := []struct { + name string + inputs [][]byte + wantErr string + }{ + { + name: "duplicate base ID", + inputs: [][]byte{ + tokenizerJSON(`{"a":1,"b":1}`, `[]`), + tokenizerJSON(`{"b":1,"a":1}`, `[]`), + }, + wantErr: "duplicate base token ID 1", + }, + { + name: "duplicate added ID", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":1,"content":"b"}]`), + tokenizerJSON(`{}`, `[{"id":1,"content":"b"},{"id":1,"content":"a"}]`), + }, + wantErr: "duplicate added token ID 1", + }, + { + name: "repeated added entry", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":1,"content":"a"}]`), + }, + wantErr: "duplicate added token ID 1", + }, + { + name: "duplicate added content", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":2,"content":"a"}]`), + tokenizerJSON(`{}`, `[{"id":2,"content":"a"},{"id":1,"content":"a"}]`), + }, + wantErr: `duplicate added token content "a" with IDs 1 and 2`, + }, + { + name: "cross-source ID conflict", + inputs: [][]byte{ + tokenizerJSON(`{"a":1}`, `[{"id":1,"content":"b"}]`), + }, + wantErr: "token ID 1 has conflicting base and added content", + }, + { + name: "cross-source content conflict", + inputs: [][]byte{ + tokenizerJSON(`{"a":1}`, `[{"id":2,"content":"a"}]`), + }, + wantErr: `token content "a" has conflicting base and added IDs 1 and 2`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for i, data := range tt.inputs { + _, err := LoadFromBytes(data) + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("input %d error = %v, want %q", i, err, tt.wantErr) + } + } + }) + } +} + +func TestLoadFromBytesBaseValidationIsDeterministic(t *testing.T) { + tests := []struct { + name string + inputs [][]byte + wantErr string + }{ + { + name: "range error precedes duplicate ID", + inputs: [][]byte{ + tokenizerJSON(fmt.Sprintf(`{"duplicate-a":5,"negative":-2,"over-cap":%d,"duplicate-b":5}`, maxTokenizerVocabularySize), `[]`), + tokenizerJSON(fmt.Sprintf(`{"duplicate-b":5,"over-cap":%d,"negative":-2,"duplicate-a":5}`, maxTokenizerVocabularySize), `[]`), + }, + wantErr: "invalid base token ID: tokenizer ID -2 must not be negative", + }, + { + name: "lowest duplicate ID wins", + inputs: [][]byte{ + tokenizerJSON(`{"high-a":7,"low-a":2,"high-b":7,"low-b":2}`, `[]`), + tokenizerJSON(`{"low-b":2,"high-b":7,"low-a":2,"high-a":7}`, `[]`), + }, + wantErr: "duplicate base token ID 2", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for i, data := range tt.inputs { + for attempt := range 100 { + _, err := LoadFromBytes(data) + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("input %d attempt %d error = %v, want %q", i, attempt, err, tt.wantErr) + } + } + } + }) + } +} + +func TestLoadFromBytesAcceptsExactAddedTokenPromotion(t *testing.T) { + tok, err := LoadFromBytes(tokenizerJSON(`{"a":1}`, `[{"id":1,"content":"a"}]`)) + if err != nil { + t.Fatal(err) + } + if got := tok.Decode([]int32{1}); got != "a" { + t.Fatalf("Decode(1) = %q, want a", got) + } + id, ok := tok.GetSpecialToken("a") + if !ok || id != 1 { + t.Fatalf("GetSpecialToken(a) = (%d, %v), want (1, true)", id, ok) + } +} + +func TestLoadFromBytesAcceptsEmptyAndContiguousVocabulary(t *testing.T) { + for name, data := range map[string][]byte{ + "empty": tokenizerJSON(`{}`, `[]`), + "contiguous": tokenizerJSON(`{"a":0,"b":1}`, `[]`), + } { + t.Run(name, func(t *testing.T) { + if _, err := LoadFromBytes(data); err != nil { + t.Fatal(err) + } + }) + } +} + func TestLoadFromBytesRejectsWordPiece(t *testing.T) { data := []byte(`{ "model": {