Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import math | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from tiny_transformer.config import ModelConfig | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, config: ModelConfig) -> None: | |
| super().__init__() | |
| self.n_head = config.n_head | |
| self.head_dim = 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.attn_dropout = nn.Dropout(config.dropout) | |
| self.resid_dropout = nn.Dropout(config.dropout) | |
| mask = torch.tril(torch.ones(config.block_size, config.block_size)) | |
| self.register_buffer("causal_mask", mask.view(1, 1, config.block_size, config.block_size)) | |
| def forward( | |
| self, x: torch.Tensor, return_attention: bool = False | |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | |
| batch, seq_len, channels = x.shape | |
| qkv = self.qkv(x) | |
| query, key, value = qkv.split(channels, dim=2) | |
| query = query.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| key = key.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| value = value.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim) | |
| scores = scores.masked_fill(self.causal_mask[:, :, :seq_len, :seq_len] == 0, float("-inf")) | |
| weights = F.softmax(scores, dim=-1) | |
| weights = self.attn_dropout(weights) | |
| out = weights @ value | |
| out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels) | |
| out = self.resid_dropout(self.proj(out)) | |
| if return_attention: | |
| return out, weights | |
| return out | |
| def forward_cached( | |
| self, | |
| x: torch.Tensor, | |
| past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None, | |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: | |
| batch, seq_len, channels = x.shape | |
| query, key, value = self.qkv(x).split(channels, dim=2) | |
| query = query.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| key = key.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| value = value.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) | |
| past_len = 0 | |
| if past_key_value is not None: | |
| past_key, past_value = past_key_value | |
| past_len = past_key.size(2) | |
| key = torch.cat((past_key, key), dim=2) | |
| value = torch.cat((past_value, value), dim=2) | |
| total_len = key.size(2) | |
| scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim) | |
| mask = self.causal_mask[:, :, past_len : past_len + seq_len, :total_len] | |
| scores = scores.masked_fill(mask == 0, float("-inf")) | |
| weights = self.attn_dropout(F.softmax(scores, dim=-1)) | |
| out = weights @ value | |
| out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels) | |
| return self.resid_dropout(self.proj(out)), (key, value) | |
| class FeedForward(nn.Module): | |
| def __init__(self, config: ModelConfig) -> None: | |
| super().__init__() | |
| self.net = 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: torch.Tensor) -> torch.Tensor: | |
| return self.net(x) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, config: ModelConfig) -> None: | |
| super().__init__() | |
| self.ln_1 = nn.LayerNorm(config.n_embd) | |
| self.attn = CausalSelfAttention(config) | |
| self.ln_2 = nn.LayerNorm(config.n_embd) | |
| self.ffwd = FeedForward(config) | |
| def forward( | |
| self, x: torch.Tensor, return_attention: bool = False | |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | |
| if return_attention: | |
| attn_out, weights = self.attn(self.ln_1(x), return_attention=True) | |
| x = x + attn_out | |
| x = x + self.ffwd(self.ln_2(x)) | |
| return x, weights | |
| x = x + self.attn(self.ln_1(x)) | |
| x = x + self.ffwd(self.ln_2(x)) | |
| return x | |
| def forward_cached( | |
| self, | |
| x: torch.Tensor, | |
| past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None, | |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: | |
| attn_out, present = self.attn.forward_cached(self.ln_1(x), past_key_value) | |
| x = x + attn_out | |
| x = x + self.ffwd(self.ln_2(x)) | |
| return x, present | |
| class TinyTransformer(nn.Module): | |
| def __init__(self, config: ModelConfig) -> None: | |
| 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.dropout = nn.Dropout(config.dropout) | |
| self.blocks = nn.ModuleList([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.apply(self._init_weights) | |
| self.lm_head.weight = self.token_embedding.weight | |
| def _init_weights(self, module: nn.Module) -> None: | |
| 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: torch.Tensor, targets: torch.Tensor | None = None | |
| ) -> tuple[torch.Tensor, torch.Tensor | None]: | |
| logits, loss, _ = self._forward(idx, targets, capture_attention=False) | |
| return logits, loss | |
| def _forward( | |
| self, | |
| idx: torch.Tensor, | |
| targets: torch.Tensor | None = None, | |
| capture_attention: bool = False, | |
| ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: | |
| batch, seq_len = idx.shape | |
| if seq_len > self.config.block_size: | |
| raise ValueError("Sequence length exceeds block_size") | |
| attentions: list[torch.Tensor] = [] | |
| positions = torch.arange(seq_len, device=idx.device) | |
| x = self.token_embedding(idx) + self.position_embedding(positions) | |
| x = self.dropout(x) | |
| for block in self.blocks: | |
| if capture_attention: | |
| x, weights = block(x, return_attention=True) | |
| attentions.append(weights) | |
| else: | |
| x = block(x) | |
| x = self.ln_f(x) | |
| logits = self.lm_head(x) | |
| loss = None | |
| if targets is not None: | |
| loss = F.cross_entropy(logits.view(batch * seq_len, -1), targets.view(batch * seq_len)) | |
| return logits, loss, attentions | |
| def _forward_cached( | |
| self, | |
| idx: torch.Tensor, | |
| past_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = None, | |
| ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]: | |
| _, seq_len = idx.shape | |
| past_len = 0 if past_key_values is None else past_key_values[0][0].size(2) | |
| if past_len + seq_len > self.config.block_size: | |
| raise ValueError("Cached sequence length exceeds block_size") | |
| positions = torch.arange(past_len, past_len + seq_len, device=idx.device) | |
| x = self.dropout(self.token_embedding(idx) + self.position_embedding(positions)) | |
| presents = [] | |
| for layer, block in enumerate(self.blocks): | |
| past = None if past_key_values is None else past_key_values[layer] | |
| x, present = block.forward_cached(x, past) | |
| presents.append(present) | |
| return self.lm_head(self.ln_f(x)), presents | |
| def parameter_count(self, trainable_only: bool = True) -> int: | |
| parameters = self.parameters() | |
| if trainable_only: | |
| parameters = (parameter for parameter in parameters if parameter.requires_grad) | |
| return sum(parameter.numel() for parameter in parameters) | |
| def attention_maps(self, idx: torch.Tensor) -> list[torch.Tensor]: | |
| self.eval() | |
| _, _, attentions = self._forward(idx, capture_attention=True) | |
| return attentions | |
| def generate( | |
| self, | |
| idx: torch.Tensor, | |
| max_new_tokens: int, | |
| temperature: float = 1.0, | |
| top_k: int | None = None, | |
| use_cache: bool = True, | |
| ) -> torch.Tensor: | |
| if max_new_tokens < 0: | |
| raise ValueError("max_new_tokens must be non-negative") | |
| if temperature <= 0: | |
| raise ValueError("temperature must be positive") | |
| if top_k is not None and top_k <= 0: | |
| raise ValueError("top_k must be positive") | |
| past_key_values = None | |
| for _ in range(max_new_tokens): | |
| if not use_cache: | |
| idx_cond = idx[:, -self.config.block_size :] | |
| logits, _ = self(idx_cond) | |
| elif past_key_values is None: | |
| idx_cond = idx[:, -self.config.block_size :] | |
| logits, past_key_values = self._forward_cached(idx_cond) | |
| else: | |
| cached_length = past_key_values[0][0].size(2) | |
| if cached_length >= self.config.block_size: | |
| idx_cond = idx[:, -self.config.block_size :] | |
| logits, past_key_values = self._forward_cached(idx_cond) | |
| else: | |
| logits, past_key_values = self._forward_cached( | |
| idx[:, -1:], past_key_values | |
| ) | |
| logits = logits[:, -1, :] / temperature | |
| if top_k is not None: | |
| values, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| logits[logits < values[:, [-1]]] = -float("inf") | |
| probs = F.softmax(logits, dim=-1) | |
| next_idx = torch.multinomial(probs, num_samples=1) | |
| idx = torch.cat((idx, next_idx), dim=1) | |
| return idx | |