Spaces:
Sleeping
Sleeping
| # models/architecture.py | |
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| class LocalConfig: | |
| def __init__(self, **kwargs): | |
| for key, value in kwargs.items(): | |
| setattr(self, key, value) | |
| # Fallback defaults if not in JSON | |
| if not hasattr(self, 'dropout'): | |
| self.dropout = 0.0 | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.n_head = config.n_head | |
| self.n_embd = config.n_embd | |
| self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False) | |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False) | |
| self.attn_dropout = nn.Dropout(config.dropout) | |
| self.resid_dropout = nn.Dropout(config.dropout) | |
| def forward(self, x): | |
| B, T, C = x.size() | |
| qkv = self.c_attn(x) | |
| q, k, v = qkv.split(self.n_embd, dim=2) | |
| head_dim = C // self.n_head | |
| q = q.view(B, T, self.n_head, head_dim).transpose(1, 2) | |
| k = k.view(B, T, self.n_head, head_dim).transpose(1, 2) | |
| v = v.view(B, T, self.n_head, head_dim).transpose(1, 2) | |
| # PyTorch Flash Attention via scaled_dot_product_attention | |
| y = F.scaled_dot_product_attention( | |
| q, k, v, | |
| is_causal=True, | |
| dropout_p=self.config.dropout if self.training else 0.0 | |
| ) | |
| y = y.transpose(1, 2).contiguous().view(B, T, C) | |
| return self.resid_dropout(self.c_proj(y)) | |
| class MLP(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) | |
| self.gelu = nn.GELU() | |
| self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) | |
| self.dropout = nn.Dropout(config.dropout) | |
| def forward(self, x): | |
| return self.dropout(self.c_proj(self.gelu(self.c_fc(x)))) | |
| class Block(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.ln_1 = nn.LayerNorm(config.n_embd) | |
| self.attn = CausalSelfAttention(config) | |
| self.ln_2 = nn.LayerNorm(config.n_embd) | |
| self.mlp = MLP(config) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln_1(x)) | |
| x = x + self.mlp(self.ln_2(x)) | |
| return x | |
| class CRAB(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.transformer = nn.ModuleDict(dict( | |
| wte = nn.Embedding(config.vocab_size, config.n_embd), | |
| wpe = nn.Embedding(config.block_size, config.n_embd), | |
| drop = nn.Dropout(config.dropout), | |
| h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), | |
| ln_f = nn.LayerNorm(config.n_embd), | |
| )) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) | |
| self.transformer.wte.weight = self.lm_head.weight | |
| def forward(self, idx, targets=None): | |
| b, t = idx.size() | |
| pos = torch.arange(0, t, dtype=torch.long, device=idx.device) | |
| x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos)) | |
| for block in self.transformer.h: | |
| x = block(x) | |
| x = self.transformer.ln_f(x) | |
| logits = self.lm_head(x) | |
| loss = None | |
| if targets is not None: | |
| # -100 ignore_index for Target Masking in instruction tuning | |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100) | |
| return logits, loss |