fix: handle attention masks without a padding token

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())
 

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

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