model.py

import torch
import torch.nn as nn

# pad=0, unk=1, a-z=2..27, hyphen=28, apostrophe=29
VOCAB = ["<PAD>", "<UNK>"] + list("abcdefghijklmnopqrstuvwxyz") + ["-", "'"]
CHAR2IDX: dict[str, int] = {c: i for i, c in enumerate(VOCAB)}

MAX_LEN = 30
NUM_CLASSES = 10  # syllable counts 1-10


def tokenize(word: str) -> list[int]:
    chars = word.lower().strip()[:MAX_LEN]
    tokens = [CHAR2IDX.get(c, CHAR2IDX["<UNK>"]) for c in chars]
    tokens += [0] * (MAX_LEN - len(tokens))
    return tokens


class SyllableTransformer(nn.Module):
    def __init__(
        self,
        vocab_size: int = len(VOCAB),
        d_model: int = 128,
        nhead: int = 4,
        num_layers: int = 3,
        dim_feedforward: int = 256,
        dropout: float = 0.1,
        max_len: int = MAX_LEN,
        num_classes: int = NUM_CLASSES,
    ):
        super().__init__()
        self.char_emb = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_emb = nn.Embedding(max_len, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers, enable_nested_tensor=False)
        self.head = nn.Linear(d_model, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L = x.shape
        pad_mask = x == 0
        positions = torch.arange(L, device=x.device).unsqueeze(0).expand(B, -1)
        emb = self.char_emb(x) + self.pos_emb(positions)
        out = self.transformer(emb, src_key_padding_mask=pad_mask)
        # mean-pool over non-pad positions
        not_pad = (~pad_mask).unsqueeze(-1).float()
        pooled = (out * not_pad).sum(1) / not_pad.sum(1).clamp(min=1e-9)
        return self.head(pooled)


MODEL_CONFIG = dict(
    vocab_size=len(VOCAB),
    d_model=256,
    nhead=8,
    num_layers=4,
    dim_feedforward=2048,
    dropout=0.1,
    max_len=MAX_LEN,
    num_classes=NUM_CLASSES,
)