diff --git a/gliner/model.py b/gliner/model.py index 3a21ee7..af9b8cc 100644 --- a/gliner/model.py +++ b/gliner/model.py @@ -899,6 +899,91 @@ def _download_model( return model_dir + @staticmethod + def _token_missing_from_tokenizer(tokenizer, token: str) -> bool: + """Check whether a special token is absent from the tokenizer vocabulary. + + A token is considered missing when it cannot be converted to an id, or + when it silently maps to the unknown-token id. + """ + token_id = tokenizer.convert_tokens_to_ids(token) + if token_id is None: + return True + unk_token_id = getattr(tokenizer, "unk_token_id", None) + return ( + unk_token_id is not None + and token_id == unk_token_id + and token != getattr(tokenizer, "unk_token", None) + ) + + @classmethod + def validate_special_token_config(cls, config_instance, tokenizer) -> None: + """Validate explicit ``class_token_index``/``vocab_size`` against the tokenizer. + + When these values are hardcoded in the config, GLiNER trusts them and does + not add its special tokens (``ent_token``, ``sep_token``). If the tokenizer + does not actually contain those tokens, the class-token mask never matches + and training proceeds with ``loss=0`` without any error (see issue #332). + This check fails fast with actionable guidance instead. + + Args: + config_instance: Model configuration with explicit special-token indices. + tokenizer: The transformer tokenizer that will be used by the model. + + Raises: + ValueError: If ``class_token_index`` is out of range for the tokenizer, + or if ``ent_token`` is missing from the tokenizer vocabulary. + """ + tokenizer_size = len(tokenizer) + + if not (0 <= config_instance.class_token_index < tokenizer_size): + raise ValueError( + f"class_token_index={config_instance.class_token_index} is out of range for the " + f"tokenizer (len={tokenizer_size}). The configured index assumes GLiNER special " + f"tokens that are not present in this tokenizer. Set class_token_index: -1 and " + f"vocab_size: -1 in your config to let GLiNER add the special tokens and detect " + f"the indices automatically." + ) + + ent_token = getattr(config_instance, "ent_token", None) + if ent_token is not None and cls._token_missing_from_tokenizer(tokenizer, ent_token): + raise ValueError( + f"The entity marker token {ent_token!r} is not present in the tokenizer " + f"vocabulary, but class_token_index={config_instance.class_token_index} and " + f"vocab_size={config_instance.vocab_size} are set explicitly. Training in this " + f"state produces an empty entity prompt and a constant zero loss. Set " + f"class_token_index: -1 and vocab_size: -1 in your config so GLiNER adds " + f"{ent_token!r} and detects its index automatically." + ) + + for optional_token_attr in ("sep_token", "rel_token"): + optional_token = getattr(config_instance, optional_token_attr, None) + if optional_token is not None and cls._token_missing_from_tokenizer(tokenizer, optional_token): + warnings.warn( + f"The special token {optional_token!r} ({optional_token_attr}) is not present " + f"in the tokenizer vocabulary. Prompts will be built from its subword pieces, " + f"which can degrade quality. Consider setting class_token_index: -1 and " + f"vocab_size: -1 so GLiNER adds all special tokens automatically.", + UserWarning, + ) + + if ent_token is not None: + ent_token_id = tokenizer.convert_tokens_to_ids(ent_token) + if ent_token_id is not None and ent_token_id != config_instance.class_token_index: + warnings.warn( + f"class_token_index={config_instance.class_token_index} does not match the " + f"tokenizer id of {ent_token!r} ({ent_token_id}). The class-token mask will " + f"match a different token, which can silently disable learning.", + UserWarning, + ) + + if config_instance.vocab_size != tokenizer_size: + warnings.warn( + f"config.vocab_size={config_instance.vocab_size} differs from the tokenizer size " + f"({tokenizer_size}). Embedding lookups may be misaligned or fail at runtime.", + UserWarning, + ) + @staticmethod def _resize_token_embeddings(instance, config_instance, tokenizer, resize_token_embeddings=True): add_tokens = instance._get_special_tokens() @@ -907,6 +992,14 @@ def _resize_token_embeddings(instance, config_instance, tokenizer, resize_token_ if tokenizer is not None: tokenizer.add_tokens(add_tokens, special_tokens=True) instance.resize_embeddings() + elif ( + tokenizer is not None + and config_instance.class_token_index != -1 + and config_instance.vocab_size != -1 + ): + # Explicit indices: verify they are consistent with the tokenizer to + # avoid silent zero-loss training (issue #332). + instance.validate_special_token_config(config_instance, tokenizer) @classmethod def load_from_config( diff --git a/tests/test_special_token_validation.py b/tests/test_special_token_validation.py new file mode 100644 index 0000000..56f6a0b --- /dev/null +++ b/tests/test_special_token_validation.py @@ -0,0 +1,120 @@ +"""Tests for special-token/config consistency validation (issue #332). + +When ``class_token_index`` and ``vocab_size`` are hardcoded in the config but the +tokenizer does not actually contain the GLiNER special tokens, training silently +runs with ``loss=0``. These tests cover the fail-fast validation that replaces +that silent failure. +""" + +import pytest + +from gliner.model import BaseGLiNER + + +class _FakeTokenizer: + """Minimal tokenizer stub: a vocab dict plus unk handling.""" + + def __init__(self, vocab, unk_token="[UNK]"): + self._vocab = dict(vocab) + self.unk_token = unk_token + self.unk_token_id = self._vocab.get(unk_token) + + def __len__(self): + return len(self._vocab) + + def convert_tokens_to_ids(self, token): + if token in self._vocab: + return self._vocab[token] + return self.unk_token_id + + def add_tokens(self, tokens, special_tokens=False): + added = 0 + for token in tokens: + if token not in self._vocab: + self._vocab[token] = len(self._vocab) + added += 1 + return added + + +class _FakeConfig: + def __init__(self, class_token_index, vocab_size, ent_token="<>", sep_token="<>"): + self.class_token_index = class_token_index + self.vocab_size = vocab_size + self.ent_token = ent_token + self.sep_token = sep_token + + +def _base_vocab(size=10): + vocab = {f"tok{i}": i for i in range(size - 1)} + vocab["[UNK]"] = size - 1 + return vocab + + +def test_out_of_range_class_token_index_raises(): + # Mirrors issue #332: mmBERT-base has base vocab 256000 but the config + # hardcoded class_token_index=256001 assuming special tokens existed. + tokenizer = _FakeTokenizer(_base_vocab(10)) + config = _FakeConfig(class_token_index=11, vocab_size=13) + + with pytest.raises(ValueError, match="out of range"): + BaseGLiNER.validate_special_token_config(config, tokenizer) + + +def test_missing_ent_token_raises(): + # Index is in range but points at an ordinary token; <> is absent, so + # the class-token mask would never match and loss would stay at zero. + tokenizer = _FakeTokenizer(_base_vocab(10)) + config = _FakeConfig(class_token_index=3, vocab_size=10) + + with pytest.raises(ValueError, match="zero loss"): + BaseGLiNER.validate_special_token_config(config, tokenizer) + + +def test_consistent_config_passes_silently(): + vocab = _base_vocab(10) + vocab["[FLERT]"] = 10 + vocab["<>"] = 11 + vocab["<>"] = 12 + tokenizer = _FakeTokenizer(vocab) + config = _FakeConfig(class_token_index=11, vocab_size=13) + + import warnings as _warnings + + with _warnings.catch_warnings(): + _warnings.simplefilter("error") + BaseGLiNER.validate_special_token_config(config, tokenizer) + + +def test_mismatched_class_token_index_warns(): + vocab = _base_vocab(10) + vocab["[FLERT]"] = 10 + vocab["<>"] = 11 + vocab["<>"] = 12 + tokenizer = _FakeTokenizer(vocab) + # <> exists at 11 but the config points the mask at 10. + config = _FakeConfig(class_token_index=10, vocab_size=13) + + with pytest.warns(UserWarning, match="does not match the tokenizer id"): + BaseGLiNER.validate_special_token_config(config, tokenizer) + + +def test_missing_sep_token_warns_but_does_not_raise(): + vocab = _base_vocab(10) + vocab["<>"] = 10 + tokenizer = _FakeTokenizer(vocab) + config = _FakeConfig(class_token_index=10, vocab_size=11) + + with pytest.warns(UserWarning, match="sep_token"): + BaseGLiNER.validate_special_token_config(config, tokenizer) + + +def test_vocab_size_mismatch_warns(): + vocab = _base_vocab(10) + vocab["[FLERT]"] = 10 + vocab["<>"] = 11 + vocab["<>"] = 12 + tokenizer = _FakeTokenizer(vocab) + config = _FakeConfig(class_token_index=11, vocab_size=999) + + with pytest.warns(UserWarning, match="vocab_size"): + BaseGLiNER.validate_special_token_config(config, tokenizer)