File size: 2,108 Bytes
bc5e3b9 | 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 54 55 56 57 58 59 60 61 | """Tiny Encoder module for CAT V3."""
from __future__ import annotations
import torch
import torch.nn as nn
class TinyEncoder(nn.Module):
"""A self-contained lightweight transformer encoder for processing text queries."""
def __init__(
self,
vocab_size: int,
hidden_size: int = 128,
nhead: int = 4,
num_layers: int = 2,
max_seq_len: int = 128,
) -> None:
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_size)
self.pos_embedding = nn.Parameter(torch.randn(1, max_seq_len, hidden_size))
encoder_layer = nn.TransformerEncoderLayer(
d_model=hidden_size,
nhead=nhead,
dim_feedforward=hidden_size * 2,
batch_first=True,
norm_first=True
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.hidden_size = hidden_size
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
"""Encode token sequence to a dense vector.
Args:
input_ids: [batch_size, seq_len] tensor of token IDs
attention_mask: [batch_size, seq_len] mask tensor (1 for active, 0 for pad)
Returns:
[batch_size, hidden_size] dense sentence representation
"""
seq_len = input_ids.size(1)
# Add word embeddings and positional encodings
x = self.embedding(input_ids) + self.pos_embedding[:, :seq_len]
# PyTorch src_key_padding_mask requires True on padding positions
padding_mask = (attention_mask == 0)
# Run Transformer
out = self.transformer(x, src_key_padding_mask=padding_mask)
# Perform mean pooling over non-padded tokens
mask_expanded = attention_mask.unsqueeze(-1).float()
sum_embeddings = torch.sum(out * mask_expanded, dim=1)
sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
pooled = sum_embeddings / sum_mask
return pooled
|