diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 2db3c8822..f2dbc463d 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/tests/unit/model_bridge/test_boot_native.py @@ -7,6 +7,8 @@ import pytest import torch import torch.nn as nn +from tokenizers import Tokenizer, models +from transformers import TokenizersBackend from transformer_lens.config import TransformerBridgeConfig from transformer_lens.model_bridge import TransformerBridge @@ -120,6 +122,30 @@ def test_boot_native_input_to_embed_round_trip(): torch.testing.assert_close(actual, expected) +def test_input_to_embed_without_pad_id_attends_every_token(): + bridge = TransformerBridge.boot_native(_cfg(device="cpu")) + bridge.eval() + bridge.tokenizer = TokenizersBackend( + tokenizer_object=Tokenizer(models.WordLevel({str(i): i for i in range(16)})), + padding_side="left", + ) + assert bridge.tokenizer.pad_token_id is None + tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) + + with torch.no_grad(): + residual, returned_tokens, _, attention_mask = bridge.input_to_embed(tokens) + expected_residual, _, _, _ = bridge.input_to_embed( + tokens, attention_mask=torch.ones_like(tokens) + ) + logits = bridge(tokens) + expected_logits = bridge(tokens, attention_mask=torch.ones_like(tokens)) + + assert torch.equal(returned_tokens, tokens) + torch.testing.assert_close(attention_mask, torch.ones_like(tokens)) + torch.testing.assert_close(residual, expected_residual) + torch.testing.assert_close(logits, expected_logits) + + def test_native_state_dict_round_trip_restores_parameters(): bridge = TransformerBridge.boot_native(_cfg()) diff --git a/tests/unit/utilities/test_tokenize_utils.py b/tests/unit/utilities/test_tokenize_utils.py index 63596e07f..29d7ea634 100644 --- a/tests/unit/utilities/test_tokenize_utils.py +++ b/tests/unit/utilities/test_tokenize_utils.py @@ -35,6 +35,35 @@ def test_attention_mask_uses_explicit_padding_side(gpt2_tokenizer) -> None: assert torch.equal(mask, torch.tensor([[0, 1, 1, 1], [1, 1, 1, 1]])) +@pytest.mark.parametrize("padding_side", ["left", "right"]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_attention_mask_without_pad_id_attends_every_token(padding_side, dtype) -> None: + tokenizer = _tokenizer(bos_token_id=7, pad_token_id=None, padding_side=padding_side) + tokens = torch.tensor([[7, 8, 9], [3, 4, 5]], dtype=dtype) + + mask = utils.get_attention_mask(tokenizer, tokens, prepend_bos=True) + + torch.testing.assert_close(mask, torch.ones_like(tokens)) + + +def test_attention_mask_preserves_bos_when_bos_is_pad() -> None: + tokenizer = _tokenizer(bos_token_id=0, pad_token_id=0, padding_side="left") + tokens = torch.tensor([[0, 0, 10, 11], [0, 20, 21, 22]]) + + mask = utils.get_attention_mask(tokenizer, tokens, prepend_bos=True) + + assert torch.equal(mask, torch.tensor([[0, 1, 1, 1], [1, 1, 1, 1]])) + + +def test_attention_mask_masks_trailing_pads_on_right() -> None: + tokenizer = _tokenizer(bos_token_id=1, pad_token_id=0, padding_side="right") + tokens = torch.tensor([[1, 10, 0, 0], [1, 20, 21, 22]]) + + mask = utils.get_attention_mask(tokenizer, tokens, prepend_bos=True) + + assert torch.equal(mask, torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]])) + + def test_bos_removal_uses_explicit_padding_side(gpt2_tokenizer) -> None: tokenizer = deepcopy(gpt2_tokenizer) tokenizer.pad_token = tokenizer.eos_token diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py index bcca9df56..f99b08d61 100644 --- a/transformer_lens/utilities/tokenize_utils.py +++ b/transformer_lens/utilities/tokenize_utils.py @@ -259,7 +259,7 @@ def get_tokens_with_bos_removed( def get_attention_mask( - tokenizer: PreTrainedTokenizerBase, + tokenizer: PreTrainedTokenizerBase | None, tokens: torch.Tensor, prepend_bos: bool, padding_side: str | None = None, @@ -282,7 +282,7 @@ def get_attention_mask( # Default to attending every token; pad positions are zeroed below. attention_mask = torch.ones_like(tokens) - if tokenizer is None: + if tokenizer is None or tokenizer.pad_token_id is None: return attention_mask if padding_side is None: padding_side = tokenizer.padding_side