fix: make tokenizer metadata files optional

fastembed/common/preprocessor_utils.py

@@ -1,6 +1,6 @@
 import json
 import sys
-from typing import Any
+from typing import Any, Iterator
 from pathlib import Path
 
 from tokenizers import AddedToken, Tokenizer
@@ -9,16 +9,29 @@
 
 
 def load_special_tokens(model_dir: Path) -> dict[str, Any]:
+    """Read special_tokens_map.json, treating an absent file as an empty map."""
     tokens_map_path = model_dir / "special_tokens_map.json"
     if not tokens_map_path.exists():
-        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+        return {}
 
     with open(str(tokens_map_path)) as tokens_map_file:
         tokens_map = json.load(tokens_map_file)
 
     return tokens_map
 
 
+def iter_special_tokens(tokens_map: dict[str, Any]) -> Iterator[str | dict[str, Any]]:
+    """Yield the individual tokens declared in a special tokens map.
+
+    Most keys hold one token, but `additional_special_tokens` holds a list of them.
+    """
+    for value in tokens_map.values():
+        if isinstance(value, list):
+            yield from value
+        else:
+            yield value
+
+
 def _valid_context(value: Any) -> int | None:
     """Return `value` if it can be used as a truncation limit, `None` otherwise.
 
@@ -59,10 +72,6 @@ def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> i
 
 
 def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
-    config_path = model_dir / "config.json"
-    if not config_path.exists():
-        raise ValueError(f"Could not find config.json in {model_dir}")
-
     tokenizer_path = model_dir / "tokenizer.json"
     if not tokenizer_path.exists():
         raise ValueError(f"Could not find tokenizer.json in {model_dir}")
@@ -71,8 +80,12 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     if not tokenizer_config_path.exists():
         raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
 
-    with open(str(config_path)) as config_file:
-        config = json.load(config_file)
+    # config.json is optional: transformers v5 no longer writes it for every model.
+    config_path = model_dir / "config.json"
+    config: dict[str, Any] = {}
+    if config_path.exists():
+        with open(str(config_path)) as config_file:
+            config = json.load(config_file)
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
@@ -84,6 +97,14 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=max_context)
 
+    # Registered before the padding is resolved: the map may name a pad token that
+    # tokenizer.json does not carry, and it only gets an id once it is added.
+    for token in iter_special_tokens(tokens_map):
+        if isinstance(token, str):
+            tokenizer.add_special_tokens([token])
+        elif isinstance(token, dict):
+            tokenizer.add_special_tokens([AddedToken(**token)])
+
     # Padding is always normalized to batch-longest. A serialized fixed length shorter than the
     # truncation limit leaves longer encodings untouched, which produces ragged batches, and a
     # fixed length equal to it pads every batch to the maximum. Direction and pad token metadata
@@ -93,29 +114,28 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     if pad_token is None:
         raise ValueError(f"Could not find a pad token for {model_dir}")
 
+    # The vocabulary is the last resort, not a hardcoded 0: that silently disagrees with
+    # `pad_token` for every model whose pad token is not the first entry.
+    pad_id = padding.get("pad_id", config.get("pad_token_id"))
+    if pad_id is None:
+        pad_id = tokenizer.token_to_id(pad_token)
+    if pad_id is None:
+        raise ValueError(f"Could not resolve an id for the pad token {pad_token!r} in {model_dir}")
+
     tokenizer.enable_padding(
         direction=padding.get("direction", "right"),
-        pad_id=padding.get("pad_id", config.get("pad_token_id", 0)),
+        pad_id=pad_id,
         pad_type_id=padding.get("pad_type_id", 0),
         pad_token=pad_token,
         pad_to_multiple_of=padding.get("pad_to_multiple_of"),
         length=None,
     )
 
-    for token in tokens_map.values():
-        if isinstance(token, str):
-            tokenizer.add_special_tokens([token])
-        elif isinstance(token, dict):
-            tokenizer.add_special_tokens([AddedToken(**token)])
-
-    special_token_to_id: dict[str, int] = {}
-
-    for token in tokens_map.values():
-        if isinstance(token, str):
-            special_token_to_id[token] = tokenizer.token_to_id(token)
-        elif isinstance(token, dict):
-            token_str = token.get("content", "")
-            special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
+    special_token_to_id = {
+        token.content: token_id
+        for token_id, token in tokenizer.get_added_tokens_decoder().items()
+        if token.special
+    }
 
     return tokenizer, special_token_to_id
 

tests/test_preprocessor_utils.py

@@ -68,17 +68,23 @@ def factory(
         config: dict[str, Any] | None = None,
         padding: dict[str, Any] | None = None,
         drop_from_tokenizer_config: tuple[str, ...] = (),
+        drop_from_config: tuple[str, ...] = (),
+        drop_files: tuple[str, ...] = (),
     ) -> Path:
         model_dir = tmp_path_factory.mktemp(f"model_dir_{next(counter)}")
         for file_name in TOKENIZER_FILES:
+            if file_name in drop_files:
+                continue
             shutil.copy(source_dir / file_name, model_dir / file_name)
 
-        _patch_json(
-            model_dir / "tokenizer_config.json",
-            tokenizer_config or {},
-            drop_from_tokenizer_config,
-        )
-        _patch_json(model_dir / "config.json", config or {})
+        if "tokenizer_config.json" not in drop_files:
+            _patch_json(
+                model_dir / "tokenizer_config.json",
+                tokenizer_config or {},
+                drop_from_tokenizer_config,
+            )
+        if "config.json" not in drop_files:
+            _patch_json(model_dir / "config.json", config or {}, drop_from_config)
         if padding is not None:
             _set_serialized_padding(model_dir / "tokenizer.json", padding)
 
@@ -230,3 +236,98 @@ def test_absent_max_context_keys_raise(make_model_dir) -> None:
 
     with pytest.raises(ValueError, match="Could not determine the maximum context length"):
         load_tokenizer(model_dir)
+
+
+@pytest.fixture(scope="module")
+def token_id(make_model_dir):
+    """Resolve vocabulary ids by name, so the cases below carry no magic numbers."""
+    return Tokenizer.from_file(str(make_model_dir() / "tokenizer.json")).token_to_id
+
+
+@pytest.mark.parametrize(
+    "dropped",
+    [("config.json",), ("special_tokens_map.json",), ("config.json", "special_tokens_map.json")],
+    ids=["no-config", "no-special-tokens-map", "neither"],
+)
+def test_optional_files_do_not_change_what_is_loaded(make_model_dir, dropped) -> None:
+    """Both files are redundant: everything they carry is already in the tokenizer."""
+    baseline, baseline_specials = load_tokenizer(make_model_dir())
+
+    tokenizer, specials = load_tokenizer(make_model_dir(drop_files=dropped))
+
+    assert specials == baseline_specials
+    assert tokenizer.padding == baseline.padding
+    assert tokenizer.encode("hello world").ids == baseline.encode("hello world").ids
+
+
+@pytest.mark.parametrize("missing", ("tokenizer.json", "tokenizer_config.json"))
+def test_the_remaining_files_are_still_required(make_model_dir, missing) -> None:
+    """Relaxing the optional two must not relax the two that carry irreplaceable data."""
+    model_dir = make_model_dir(drop_files=(missing,))
+
+    with pytest.raises(ValueError, match=f"Could not find {missing}"):
+        load_tokenizer(model_dir)
+
+
+@pytest.mark.parametrize(
+    "model_files",
+    [
+        pytest.param({"drop_from_config": ("pad_token_id",)}, id="config-omits-pad-token-id"),
+        pytest.param({"drop_files": ("config.json",)}, id="config-is-absent"),
+    ],
+)
+def test_pad_id_falls_back_to_the_vocabulary(make_model_dir, token_id, model_files) -> None:
+    """Last link of the chain; a hardcoded 0 would silently disagree with `pad_token`."""
+    expected = token_id("[SEP]")
+    assert expected != 0, "a pad token whose id is 0 would pass even without a lookup"
+    model_dir = make_model_dir(tokenizer_config={"pad_token": "[SEP]"}, **model_files)
+
+    tokenizer, _ = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_id"] == expected
+
+
+def test_pad_token_that_resolves_nowhere_raises(make_model_dir) -> None:
+    """Without config.json a pad token outside the vocabulary has no id left to fall back on."""
+    model_dir = make_model_dir(
+        tokenizer_config={"pad_token": "[NOT_IN_VOCAB]"},
+        drop_files=("config.json",),
+    )
+
+    with pytest.raises(ValueError, match="Could not resolve an id for the pad token"):
+        load_tokenizer(model_dir)
+
+
+def test_pad_token_named_only_in_the_map_resolves(make_model_dir) -> None:
+    """The map is read first, so it can name a pad token tokenizer.json does not carry."""
+    model_dir = make_model_dir(
+        tokenizer_config={"pad_token": "<|mypad|>"},
+        drop_files=("config.json",),
+    )
+    _patch_json(model_dir / "special_tokens_map.json", {"pad_token": "<|mypad|>"})
+
+    tokenizer, specials = load_tokenizer(model_dir)
+
+    assert tokenizer.padding["pad_token"] == "<|mypad|>"
+    assert tokenizer.padding["pad_id"] == specials["<|mypad|>"]
+
+
+@pytest.mark.parametrize(
+    "additional",
+    [
+        pytest.param(["<|list_str|>"], id="list-of-strings"),
+        pytest.param([{"content": "<|list_str|>"}], id="list-of-added-token-dicts"),
+    ],
+)
+def test_list_valued_map_entries_are_registered(make_model_dir, additional) -> None:
+    """`additional_special_tokens` holds a list, which the str/dict dispatch alone drops.
+
+    Real repos ship both spellings, and their tokens are in tokenizer.json already, so
+    only a token living nowhere else shows the drop.
+    """
+    model_dir = make_model_dir()
+    _patch_json(model_dir / "special_tokens_map.json", {"additional_special_tokens": additional})
+
+    _, specials = load_tokenizer(model_dir)
+
+    assert "<|list_str|>" in specials