File size: 1,547 Bytes
25f9f37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import torch
from torch import nn
from colipri.model.text import TextEncoder
from colipri.pooling import AttentionPool1D
class _StubBackbone(nn.Module):
"""Encode token IDs as deterministic embeddings."""
def __init__(self):
"""Initialize embeddings with a distinctive padding-token vector."""
super().__init__()
weights = torch.tensor(
[
[100.0, -100.0, 50.0, -50.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
]
)
self.embedding = nn.Embedding.from_pretrained(weights, freeze=False)
def forward(self, token_ids, attention_mask=None): # noqa: ARG002
"""Encode token IDs without altering padded positions."""
return {"last_hidden_state": self.embedding(token_ids)}
def test_text_encoder_ignores_masked_padding():
"""Preserve the pooled embedding when masked padding is appended."""
torch.manual_seed(0)
encoder = TextEncoder(
backbone=_StubBackbone(),
pooler=AttentionPool1D(embed_dim=4, num_heads=2),
)
valid_ids = torch.tensor([[1, 2, 3]])
padded_ids = torch.tensor([[1, 2, 3, 0, 0]])
valid_mask = torch.ones_like(valid_ids)
padded_mask = torch.tensor([[1, 1, 1, 0, 0]])
expected = encoder(
valid_ids,
valid_mask,
normalize=False,
)
result = encoder(
padded_ids,
padded_mask,
normalize=False,
)
torch.testing.assert_close(result, expected)
|