158 lines
6.2 KiB
Python
158 lines
6.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Slice 19 model zoo: three genuinely sequence-sensitive tiny models, trained
|
|
from scratch on Maven's narrow binary pragmatics task.
|
|
|
|
A. CharCNN — codepoint ids → char embedding → parallel small 1D convs
|
|
(several kernel widths) → global max-pool → linear head
|
|
B. BiGRU — subword ids → token embedding → 1-layer BiGRU →
|
|
maxpool[final] → linear head
|
|
C. TinyTransformer — subword ids → token embedding + sine position →
|
|
N self-attention encoder blocks (heads, FFN 4x, PreNorm) →
|
|
CLS → linear head
|
|
|
|
All expose :forward(ids) returning the binary logit, plus .n_params().
|
|
Deterministic: everything is plain torch ops.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class CharCNN(nn.Module):
|
|
def __init__(self, vocab_size, embed_dim, filters, widths, pad_idx=0, dropout=0.3):
|
|
super().__init__()
|
|
self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
|
|
self.convs = nn.ModuleList([
|
|
nn.Conv1d(embed_dim, filters, k, padding=(k - 1) // 2)
|
|
for k in widths
|
|
])
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.head = nn.Linear(filters * len(widths), 1)
|
|
|
|
def forward(self, ids):
|
|
# ids: (B, T)
|
|
x = self.embed(ids).transpose(1, 2) # (B, D, T)
|
|
hiddens = [F.relu(conv(x)) for conv in self.convs] # each (B, F, T)
|
|
pooled = torch.cat([h.max(dim=2).values for h in hiddens], dim=1) # (B, F*W)
|
|
return self.head(self.dropout(pooled)).squeeze(-1)
|
|
|
|
def n_params(self):
|
|
return sum(p.numel() for p in self.parameters())
|
|
|
|
|
|
class BiGRU(nn.Module):
|
|
def __init__(self, vocab_size, embed_dim, hidden, pad_idx=0, dropout=0.3):
|
|
super().__init__()
|
|
self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
|
|
self.encoder = nn.GRU(embed_dim, hidden, num_layers=1, bidirectional=True,
|
|
batch_first=True)
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.head = nn.Linear(hidden * 2, 1)
|
|
|
|
def forward(self, ids):
|
|
mask = (ids != 0).float() # (B, T)
|
|
x = self.embed(ids)
|
|
lens = mask.sum(dim=1).clamp(min=1).long()
|
|
x_p = nn.utils.rnn.pack_padded_sequence(x, lens.cpu(), batch_first=True,
|
|
enforce_sorted=False)
|
|
out, _ = self.encoder(x_p)
|
|
out, _ = nn.utils.rnn.pad_packed_sequence(out, batch_first=True,
|
|
total_length=mask.size(1))
|
|
out = out * mask.unsqueeze(-1)
|
|
maxed = out.max(dim=1).values # (B, 2H)
|
|
return self.head(self.dropout(maxed)).squeeze(-1)
|
|
|
|
def n_params(self):
|
|
return sum(p.numel() for p in self.parameters())
|
|
|
|
|
|
class TinyTransformer(nn.Module):
|
|
def __init__(self, vocab_size, d_model, n_layers, n_heads, ff_mult=4,
|
|
max_len=64, pad_idx=0, dropout=0.1):
|
|
super().__init__()
|
|
self.d_model = d_model
|
|
self.embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_idx)
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.pos = nn.Parameter(torch.empty(1, max_len, d_model))
|
|
nn.init.normal_(self.pos, std=0.02)
|
|
blocks = []
|
|
for _ in range(n_layers):
|
|
blocks.append(TransformerBlock(d_model, n_heads, ff_mult, dropout))
|
|
self.blocks = nn.ModuleList(blocks)
|
|
self.ln_out = nn.LayerNorm(d_model)
|
|
self.head = nn.Linear(d_model, 1)
|
|
|
|
def forward(self, ids):
|
|
B, T = ids.shape
|
|
mask = (ids != 0)
|
|
x = self.embed(ids) * math.sqrt(self.d_model) + self.pos[:, :T, :]
|
|
x = self.dropout(x)
|
|
for blk in self.blocks:
|
|
x = blk(x, mask)
|
|
x = self.ln_out(x)
|
|
pooled = x.masked_fill(~mask.unsqueeze(-1), float("-inf")).max(dim=1).values
|
|
return self.head(pooled).squeeze(-1)
|
|
|
|
def n_params(self):
|
|
return sum(p.numel() for p in self.parameters())
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(self, d_model, n_heads, ff_mult, dropout):
|
|
super().__init__()
|
|
self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout,
|
|
batch_first=True)
|
|
self.ln1 = nn.LayerNorm(d_model)
|
|
self.ff = nn.Sequential(
|
|
nn.Linear(d_model, d_model * ff_mult),
|
|
nn.GELU(),
|
|
nn.Linear(d_model * ff_mult, d_model),
|
|
)
|
|
self.ln2 = nn.LayerNorm(d_model)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x, mask):
|
|
# mask: (B, T) bool of non-pad; MultiheadAttention uses key_padding_mask
|
|
h = self.ln1(x)
|
|
h = self.attn(h, h, h, key_padding_mask=~mask,
|
|
need_weights=False, is_causal=False)[0]
|
|
x = x + self.dropout(h)
|
|
h = self.ln2(x)
|
|
x = x + self.dropout(self.ff(h))
|
|
return x
|
|
|
|
|
|
# ─── Sizes ladder ───────────────────────────────────────────────────────────
|
|
|
|
def make_model(arch, size, char_vocab, bpe_vocab):
|
|
if arch == "char_cnn":
|
|
configs = {
|
|
"tiny": dict(embed_dim=32, filters=64, widths=[3, 4, 5]),
|
|
"medium": dict(embed_dim=64, filters=160, widths=[2, 3, 4, 5]),
|
|
}
|
|
c = configs[size]
|
|
return CharCNN(char_vocab, c["embed_dim"], c["filters"], c["widths"])
|
|
if arch == "bigru":
|
|
configs = {
|
|
"tiny": dict(embed_dim=64, hidden=64),
|
|
"medium": dict(embed_dim=128, hidden=128),
|
|
"large": dict(embed_dim=256, hidden=256),
|
|
}
|
|
c = configs[size]
|
|
return BiGRU(bpe_vocab, c["embed_dim"], c["hidden"])
|
|
if arch == "tiny_transformer":
|
|
configs = {
|
|
"small": dict(d_model=128, n_layers=2, n_heads=4),
|
|
"medium": dict(d_model=192, n_layers=4, n_heads=4),
|
|
}
|
|
c = configs[size]
|
|
return TinyTransformer(bpe_vocab, c["d_model"], c["n_layers"], c["n_heads"])
|
|
raise ValueError(arch)
|
|
|
|
|
|
def n_params_of(arch, size, char_vocab, bpe_vocab):
|
|
return make_model(arch, size, char_vocab, bpe_vocab).n_params() |