File size: 4,070 Bytes
9cd7b60 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import torch
import torch.nn as nn
class CompactGenderClassifier(nn.Module):
"""
Ultra-compact transformer binary gender classifier trained from scratch.
Features:
- Factorized Embeddings (vocab_size -> emb_dim -> hidden_dim) for ~50% embedding weight reduction
- Pre-LN TransformerEncoder for fast, stable training dynamics
- Masked Dual-Pooling: Concat(MeanPool, MaxPool) for subword affix sensitivity
- Single-task binary classification head: MALE (0) vs FEMALE (1)
- Zero BertPooler / token_type_ids bloat for clean ONNX execution
"""
def __init__(
self,
vocab_size=4000,
emb_dim=64,
hidden_dim=128,
num_layers=2,
num_heads=4,
dim_feedforward=256,
max_len=20,
dropout=0.1
):
super().__init__()
self.vocab_size = vocab_size
self.emb_dim = emb_dim
self.hidden_dim = hidden_dim
self.max_len = max_len
# Factorized Token Embeddings: cuts parameters from 4000x128=512k to 4000x64+64x128=264k
self.token_emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
self.emb_proj = nn.Linear(emb_dim, hidden_dim, bias=False)
self.pos_emb = nn.Embedding(max_len, hidden_dim)
# Transformer Encoder with Pre-LN (norm_first=True)
encoder_layer = nn.TransformerEncoderLayer(
d_model=hidden_dim,
nhead=num_heads,
dim_feedforward=dim_feedforward,
dropout=dropout,
activation="gelu",
batch_first=True,
norm_first=True
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.norm = nn.LayerNorm(hidden_dim)
self.dropout = nn.Dropout(dropout)
# Dual Pooling: 128 (MeanPool) + 128 (MaxPool) = 256
self.classifier = nn.Linear(hidden_dim * 2, 2)
self._init_weights()
def _init_weights(self):
"""Standard truncated normal initialization for training from scratch."""
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Embedding):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.padding_idx is not None:
with torch.no_grad():
m.weight[m.padding_idx].fill_(0)
elif isinstance(m, nn.LayerNorm):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, input_ids, attention_mask):
# input_ids: [batch_size, seq_len] (int32 or int64)
# attention_mask: [batch_size, seq_len] (int32 or int64, 1 for real tokens, 0 for pad)
seq_len = input_ids.size(1)
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
# Embeddings: token + positional
x = self.emb_proj(self.token_emb(input_ids.long())) + self.pos_emb(positions)
# PyTorch TransformerEncoder expects src_key_padding_mask=True for PAD tokens
padding_mask = (attention_mask == 0)
hidden = self.encoder(x, src_key_padding_mask=padding_mask)
hidden = self.norm(hidden) # [batch_size, seq_len, hidden_dim]
# 1. Masked Mean Pooling
mask_expanded = attention_mask.unsqueeze(-1).float()
sum_embeddings = torch.sum(hidden * mask_expanded, dim=1)
sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
mean_pooled = sum_embeddings / sum_mask
# 2. Masked Max Pooling (padded positions set to -1e9 so they are never max)
hidden_masked = hidden.masked_fill(padding_mask.unsqueeze(-1), -1e9)
max_pooled = torch.max(hidden_masked, dim=1).values
# Concatenate features: 128 + 128 = 256
features = torch.cat([mean_pooled, max_pooled], dim=-1)
features = self.dropout(features)
logits = self.classifier(features)
return logits
|