diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index 2fbd9986..c08658db 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/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,9 +9,10 @@ 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) @@ -19,6 +20,18 @@ def load_special_tokens(model_dir: Path) -> dict[str, Any]: 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 diff --git a/tests/test_preprocessor_utils.py b/tests/test_preprocessor_utils.py index 23cfc58c..0b928f50 100644 --- a/tests/test_preprocessor_utils.py +++ b/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