From 9a181c10094d973e9cdaed3d94f30576176b4ba5 Mon Sep 17 00:00:00 2001 From: Sohan Venkatesh Date: Sat, 8 Aug 2026 17:08:43 +0100 Subject: [PATCH] fix(tokenizer): do not strip a BOS token the tokenizer does not have get_tokens_with_bos_removed assumed a bos_token_id exists. Callers gate it on cfg.tokenizer_prepends_bos, which detect_tokenizer_bos_eos only sets when the tokenizer has one, but the flag goes stale on a bridge built via build_bridge_from_module(tokenizer=None) and given a tokenizer afterwards: the setter re-runs configure_tokenizer only on reassignment, so the config default of True survives. Trusting it then does damage in both directions. Under right padding, the default, the helper drops the first token unconditionally, silently removing [CLS] from a BERT tokenizer's output and returning a plausible-looking wrong result. Under left padding it evaluates (tokens == None).int() and raises AttributeError: 'bool' object has no attribute 'int', which names neither the tokenizer nor the flag. Return the tokens unchanged when there is no bos_token_id: with no BOS there is nothing to remove, which is correct however the config got out of sync. Reproduced with two off-the-shelf tokenizers, bert-base-cased and t5-small, both of which have bos_token_id None. The normal boot_transformers path is unaffected, since detection runs there. The root cause is the reassignment test in bridge_core.py, left alone deliberately. Correcting the flag would route to_tokens(prepend_bos=True) into the manual-prepend branch at transformer_bridge.py:716, which calls get_input_with_manually_prepended_bos(tokenizer.bos_token, ...) and raises TypeError on a None bos_token. The stale flag currently masks that, so the root-cause fix needs the prepend path hardened first. Fixes #1628 Co-Authored-By: Claude Opus 5 --- .../test_get_tokens_with_bos_removed.py | 82 +++++++++++++++++++ transformer_lens/utilities/tokenize_utils.py | 8 ++ 2 files changed, 90 insertions(+) create mode 100644 tests/unit/utilities/test_get_tokens_with_bos_removed.py diff --git a/tests/unit/utilities/test_get_tokens_with_bos_removed.py b/tests/unit/utilities/test_get_tokens_with_bos_removed.py new file mode 100644 index 000000000..75e51ad74 --- /dev/null +++ b/tests/unit/utilities/test_get_tokens_with_bos_removed.py @@ -0,0 +1,82 @@ +"""Tests for get_tokens_with_bos_removed when the tokenizer has no BOS token. + +Callers gate this helper on ``cfg.tokenizer_prepends_bos``. That flag is set by +``detect_tokenizer_bos_eos``, which requires a ``bos_token_id`` — so a tokenizer +with none should never reach here. It does when the flag is stale: a bridge built +via ``build_bridge_from_module(tokenizer=None)`` keeps the config default of True, +and the tokenizer setter only re-runs detection on *re*-assignment. + +Trusting a stale flag is not harmless. Under right padding the helper drops the +first token unconditionally, which silently removes ``[CLS]`` from a BERT +tokenizer's output; under left padding it compares tokens against ``None`` and +raises an ``AttributeError`` naming neither the tokenizer nor the flag. +""" + +from __future__ import annotations + +import pytest +import torch +from transformers import AutoTokenizer + +from transformer_lens.utilities.tokenize_utils import get_tokens_with_bos_removed + + +@pytest.fixture(scope="module") +def no_bos_tokenizer(): + """BERT uses [CLS] rather than a BOS token, so bos_token_id is None.""" + tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased") + assert tokenizer.bos_token_id is None + return tokenizer + + +@pytest.fixture(scope="module") +def bos_tokenizer(): + tokenizer = AutoTokenizer.from_pretrained("distilgpt2") + assert tokenizer.bos_token_id is not None + return tokenizer + + +@pytest.mark.parametrize("padding_side", ["left", "right"]) +def test_no_bos_token_returns_tokens_unchanged(no_bos_tokenizer, padding_side) -> None: + """There is no BOS to remove, so the tokens must come back untouched.""" + no_bos_tokenizer.padding_side = padding_side + tokens = torch.tensor([[101, 19082, 1362, 102]]) + + result = get_tokens_with_bos_removed(no_bos_tokenizer, tokens) + + torch.testing.assert_close(result, tokens) + + +def test_no_bos_token_does_not_drop_cls_under_right_padding(no_bos_tokenizer) -> None: + """The damaging case: [CLS] is not a BOS token, and dropping it changes what + the model is asked to encode.""" + no_bos_tokenizer.padding_side = "right" + tokens = no_bos_tokenizer("hello world", return_tensors="pt")["input_ids"] + + result = get_tokens_with_bos_removed(no_bos_tokenizer, tokens) + + assert result.shape == tokens.shape + assert result[0, 0].item() == no_bos_tokenizer.cls_token_id + + +def test_no_bos_token_does_not_raise_under_left_padding(no_bos_tokenizer) -> None: + """Previously `(tokens == None).int()` — a Python bool, not a tensor.""" + no_bos_tokenizer.padding_side = "left" + tokens = torch.tensor([[101, 19082, 1362, 102]]) + + result = get_tokens_with_bos_removed(no_bos_tokenizer, tokens) + + assert result.shape == tokens.shape + + +@pytest.mark.parametrize("padding_side", ["left", "right"]) +def test_a_real_bos_is_still_removed(bos_tokenizer, padding_side) -> None: + """The guard must not disturb the case the helper exists for.""" + bos_tokenizer.padding_side = padding_side + bos = bos_tokenizer.bos_token_id + tokens = torch.tensor([[bos, 15496, 995]]) + + result = get_tokens_with_bos_removed(bos_tokenizer, tokens) + + assert result.shape[-1] == tokens.shape[-1] - 1 + assert bos not in result[0].tolist() diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py index 418694798..3abe27af9 100644 --- a/transformer_lens/utilities/tokenize_utils.py +++ b/transformer_lens/utilities/tokenize_utils.py @@ -216,6 +216,14 @@ def get_tokens_with_bos_removed( Returns: torch.Tensor: The tokenized input with the bos token removed. """ + if tokenizer.bos_token_id is None: + # Nothing to remove (#1628). Callers reach this when cfg.tokenizer_prepends_bos + # says the tokenizer prepends a BOS but the tokenizer has none — a stale + # flag, since detect_tokenizer_bos_eos() requires a bos_token_id. Trusting + # it here would drop a real first token under right padding ([CLS] for a + # BERT tokenizer), and compare tokens against None under left padding. + return tokens + if tokenizer.padding_side == "right": return tokens[..., 1:]