Spaces:
Runtime error
Runtime error
| import math | |
| from dataclasses import dataclass | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| class ModelConfig: | |
| vocab_size: int | |
| block_size: int = 128 | |
| n_embd: int = 128 | |
| n_head: int = 4 | |
| n_layer: int = 4 | |
| dropout: float = 0.1 | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, config: ModelConfig): | |
| super().__init__() | |
| assert config.n_embd % config.n_head == 0 | |
| self.n_head = config.n_head | |
| self.head_size = config.n_embd // config.n_head | |
| self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd) | |
| self.proj = nn.Linear(config.n_embd, config.n_embd) | |
| self.dropout = nn.Dropout(config.dropout) | |
| mask = torch.tril(torch.ones(config.block_size, config.block_size)) | |
| self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size)) | |
| def forward(self, x): | |
| bsz, seq_len, channels = x.shape | |
| qkv = self.qkv(x) | |
| q, k, v = qkv.split(channels, dim=2) | |
| q = q.view(bsz, seq_len, self.n_head, self.head_size).transpose(1, 2) | |
| k = k.view(bsz, seq_len, self.n_head, self.head_size).transpose(1, 2) | |
| v = v.view(bsz, seq_len, self.n_head, self.head_size).transpose(1, 2) | |
| scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_size) | |
| scores = scores.masked_fill(self.mask[:, :, :seq_len, :seq_len] == 0, float("-inf")) | |
| weights = F.softmax(scores, dim=-1) | |
| weights = self.dropout(weights) | |
| out = weights @ v | |
| out = out.transpose(1, 2).contiguous().view(bsz, seq_len, channels) | |
| return self.dropout(self.proj(out)) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, config: ModelConfig): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(config.n_embd) | |
| self.attn = CausalSelfAttention(config) | |
| self.ln2 = nn.LayerNorm(config.n_embd) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(config.n_embd, 4 * config.n_embd), | |
| nn.GELU(), | |
| nn.Linear(4 * config.n_embd, config.n_embd), | |
| nn.Dropout(config.dropout), | |
| ) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln1(x)) | |
| x = x + self.mlp(self.ln2(x)) | |
| return x | |
| class SuperLilLM(nn.Module): | |
| def __init__(self, config: ModelConfig): | |
| super().__init__() | |
| self.config = config | |
| self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd) | |
| self.position_embedding = nn.Embedding(config.block_size, config.n_embd) | |
| self.blocks = nn.Sequential(*[TransformerBlock(config) for _ in range(config.n_layer)]) | |
| self.ln_f = nn.LayerNorm(config.n_embd) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size) | |
| self.dropout = nn.Dropout(config.dropout) | |
| self.apply(self._init_weights) | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| if module.bias is not None: | |
| nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, idx, targets=None): | |
| _, seq_len = idx.shape | |
| if seq_len > self.config.block_size: | |
| raise ValueError(f"Sequence length {seq_len} exceeds block size {self.config.block_size}") | |
| positions = torch.arange(seq_len, device=idx.device) | |
| x = self.token_embedding(idx) + self.position_embedding(positions) | |
| x = self.dropout(x) | |
| x = self.blocks(x) | |
| x = self.ln_f(x) | |
| logits = self.lm_head(x) | |
| loss = None | |
| if targets is not None: | |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100) | |
| return logits, loss | |