File size: 4,261 Bytes
7a7d768 | 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """
TaxiLM — Vanilla transformer for Hassaniya.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from config import TaxiConfig
class Attention(nn.Module):
def __init__(self, config):
super().__init__()
self.n_heads = config.n_heads
self.head_dim = config.d_model // config.n_heads
self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
self.out = nn.Linear(config.d_model, config.d_model)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x, mask=None):
B, T, C = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
attn = attn.masked_fill(mask == 0, float("-inf"))
attn = self.dropout(F.softmax(attn, dim=-1))
return self.out((attn @ v).transpose(1, 2).contiguous().view(B, T, C))
class FFN(nn.Module):
def __init__(self, config):
super().__init__()
self.up = nn.Linear(config.d_model, config.ffn_hidden)
self.down = nn.Linear(config.ffn_hidden, config.d_model)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x):
return self.dropout(self.down(F.relu(self.up(x))))
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.norm1 = nn.LayerNorm(config.d_model)
self.attn = Attention(config)
self.norm2 = nn.LayerNorm(config.d_model)
self.ffn = FFN(config)
def forward(self, x, mask=None):
x = x + self.attn(self.norm1(x), mask)
x = x + self.ffn(self.norm2(x))
return x
class TaxiLM(nn.Module): # ← Changé
def __init__(self, config: TaxiConfig): # ← Changé
super().__init__()
self.config = config
self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
self.pos_emb = nn.Embedding(config.max_seq_len, config.d_model)
self.drop = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
self.norm = nn.LayerNorm(config.d_model)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
self.lm_head.weight = self.tok_emb.weight
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Embedding):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
def forward(self, idx, targets=None):
B, T = idx.shape
pos = torch.arange(T, device=idx.device)
x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
mask = torch.tril(torch.ones(T, T, device=idx.device)).unsqueeze(0).unsqueeze(0)
for block in self.blocks:
x = block(x, mask)
logits = self.lm_head(self.norm(x))
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1), ignore_index=0)
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new_tokens=64, temperature=0.7, top_k=50, **kwargs):
self.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.config.max_seq_len:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / temperature
if top_k > 0:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = float("-inf")
probs = F.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
idx = torch.cat([idx, next_id], dim=1)
if next_id.item() == self.config.eos_id:
break
return idx, []
def param_count(self):
total = sum(p.numel() for p in self.parameters())
return total, 0
def param_summary(self):
total, _ = self.param_count()
return f"TaxiLM: {total:,} params ({total/1e6:.1f}M)" # ← Changé |