Initial release: from-scratch char Transformer, 40.0% exact-match on held-out vocabulary
1982a09 verified | """Small char-level encoder-decoder Transformer, built from scratch | |
| (no pretrained weights) — the task and vocabulary are small enough that | |
| this is a cleaner research artifact than fine-tuning a large pretrained LM. | |
| """ | |
| from dataclasses import dataclass | |
| import torch | |
| from torch import nn | |
| from vocab import PAD, VOCAB_SIZE | |
| class ModelConfig: | |
| vocab_size: int = VOCAB_SIZE | |
| d_model: int = 128 | |
| nhead: int = 4 | |
| num_encoder_layers: int = 3 | |
| num_decoder_layers: int = 3 | |
| dim_feedforward: int = 256 | |
| dropout: float = 0.1 | |
| max_len: int = 32 | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, d_model: int, max_len: int): | |
| super().__init__() | |
| pos = torch.arange(max_len).unsqueeze(1) | |
| div = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model)) | |
| pe = torch.zeros(max_len, d_model) | |
| pe[:, 0::2] = torch.sin(pos * div) | |
| pe[:, 1::2] = torch.cos(pos * div) | |
| self.register_buffer("pe", pe) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return x + self.pe[: x.size(1)].unsqueeze(0) | |
| class SpellCorrectorTransformer(nn.Module): | |
| def __init__(self, cfg: ModelConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=PAD) | |
| self.pos_enc = PositionalEncoding(cfg.d_model, cfg.max_len) | |
| self.transformer = nn.Transformer( | |
| d_model=cfg.d_model, | |
| nhead=cfg.nhead, | |
| num_encoder_layers=cfg.num_encoder_layers, | |
| num_decoder_layers=cfg.num_decoder_layers, | |
| dim_feedforward=cfg.dim_feedforward, | |
| dropout=cfg.dropout, | |
| batch_first=True, | |
| ) | |
| self.out_proj = nn.Linear(cfg.d_model, cfg.vocab_size) | |
| def forward( | |
| self, | |
| src: torch.Tensor, | |
| tgt: torch.Tensor, | |
| src_key_padding_mask: torch.Tensor, | |
| tgt_key_padding_mask: torch.Tensor, | |
| ) -> torch.Tensor: | |
| src_emb = self.pos_enc(self.embed(src)) | |
| tgt_emb = self.pos_enc(self.embed(tgt)) | |
| tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1), device=tgt.device) | |
| out = self.transformer( | |
| src_emb, | |
| tgt_emb, | |
| tgt_mask=tgt_mask, | |
| src_key_padding_mask=src_key_padding_mask, | |
| tgt_key_padding_mask=tgt_key_padding_mask, | |
| memory_key_padding_mask=src_key_padding_mask, | |
| ) | |
| return self.out_proj(out) | |
| def num_parameters(self) -> int: | |
| return sum(p.numel() for p in self.parameters()) | |