Gate byte token types on byte fallback

convert/tokenizer.go

@@ -210,10 +210,12 @@ func parseTokenizer(fsys fs.FS, specialTokenTypes []string) (*Tokenizer, error)
 type tokenizer struct {
 	AddedTokens []token `json:"added_tokens"`
 	Model       struct {
-		Type   string          `json:"type"`
-		Vocab  map[string]int  `json:"vocab"`
-		Merges json.RawMessage `json:"merges"`
+		Type         string          `json:"type"`
+		Vocab        map[string]int  `json:"vocab"`
+		Merges       json.RawMessage `json:"merges"`
+		ByteFallback bool            `json:"byte_fallback"`
 	} `json:"model"`
+	Decoder json.RawMessage `json:"decoder"`
 
 	PreTokenizer struct {
 		PreTokenizers []struct {
@@ -269,6 +271,8 @@ func parseVocabularyFromTokenizer(fsys fs.FS) (*Vocabulary, error) {
 		tokens[token.ID] = token
 	}
 
+	hasByteFallback := t.Model.ByteFallback || hasDecoderType(t.Decoder, "ByteFallback")
+
 	v := Vocabulary{Model: "gpt2"}
 	for _, k := range slices.Sorted(maps.Keys(tokens)) {
 		token := tokens[k]
@@ -280,7 +284,7 @@ func parseVocabularyFromTokenizer(fsys fs.FS) (*Vocabulary, error) {
 			v.Types = append(v.Types, tokenTypeControl)
 		case token.UserDefined:
 			v.Types = append(v.Types, tokenTypeUserDefined)
-		case isByteFallbackToken(token.Content):
+		case hasByteFallback && isByteFallbackToken(token.Content):
 			v.Types = append(v.Types, tokenTypeByte)
 		default:
 			v.Types = append(v.Types, tokenTypeNormal)
@@ -326,6 +330,41 @@ func isByteFallbackToken(s string) bool {
 	return true
 }
 
+func hasDecoderType(data json.RawMessage, typ string) bool {
+	if len(data) == 0 {
+		return false
+	}
+
+	var v any
+	if err := json.Unmarshal(data, &v); err != nil {
+		return false
+	}
+
+	return containsDecoderType(v, typ)
+}
+
+func containsDecoderType(v any, typ string) bool {
+	switch v := v.(type) {
+	case map[string]any:
+		if s, ok := v["type"].(string); ok && s == typ {
+			return true
+		}
+		for _, vv := range v {
+			if containsDecoderType(vv, typ) {
+				return true
+			}
+		}
+	case []any:
+		for _, vv := range v {
+			if containsDecoderType(vv, typ) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
 type SpecialVocabulary struct {
 	Type     string
 	ID       int

convert/tokenizer_test.go

@@ -141,7 +141,8 @@ func TestParseTokenizer(t *testing.T) {
 							"<0x82>": 2,
 							"<0x87>": 3,
 							"plain": 4
-						}
+						},
+						"byte_fallback": true
 					}
 				}`),
 			}),
@@ -155,6 +156,58 @@ func TestParseTokenizer(t *testing.T) {
 				Pre: "default",
 			},
 		},
+		{
+			name: "byte-shaped tokens without byte fallback",
+			fsys: createTokenizerFS(t, t.TempDir(), map[string]io.Reader{
+				"tokenizer.json": strings.NewReader(`{
+					"model": {
+						"vocab": {
+							"<0x41>": 0,
+							"plain": 1
+						}
+					}
+				}`),
+			}),
+			want: &Tokenizer{
+				Vocabulary: &Vocabulary{
+					Model:  "gpt2",
+					Tokens: []string{"<0x41>", "plain"},
+					Scores: []float32{0, 1},
+					Types:  []int32{1, 1},
+				},
+				Pre: "default",
+			},
+		},
+		{
+			name: "byte fallback decoder",
+			fsys: createTokenizerFS(t, t.TempDir(), map[string]io.Reader{
+				"tokenizer.json": strings.NewReader(`{
+					"decoder": {
+						"type": "Sequence",
+						"decoders": [
+							{"type": "Replace"},
+							{"type": "ByteFallback"},
+							{"type": "Fuse"}
+						]
+					},
+					"model": {
+						"vocab": {
+							"<0x41>": 0,
+							"plain": 1
+						}
+					}
+				}`),
+			}),
+			want: &Tokenizer{
+				Vocabulary: &Vocabulary{
+					Model:  "gpt2",
+					Tokens: []string{"<0x41>", "plain"},
+					Scores: []float32{0, 1},
+					Types:  []int32{6, 1},
+				},
+				Pre: "default",
+			},
+		},
 		{
 			name: "special token types",
 			fsys: createTokenizerFS(t, t.TempDir(), map[string]io.Reader{