| import os |
| import json |
| from huggingface_hub import snapshot_download |
|
|
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class DenseTransformerBlock(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(n_embd) |
| self.attn = nn.MultiheadAttention(n_embd, n_head, dropout=dropout, batch_first=True) |
| self.ln2 = nn.LayerNorm(n_embd) |
| self.mlp = nn.Sequential( |
| nn.Linear(n_embd, 4 * n_embd), |
| nn.GELU(), |
| nn.Linear(4 * n_embd, n_embd), |
| nn.Dropout(dropout), |
| ) |
| self.register_buffer("mask", torch.triu(torch.full((block_size, block_size), float("-inf")), diagonal=1)) |
|
|
| def forward(self, x): |
| t = x.shape[1] |
| a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=self.mask[:t, :t]) |
| x = x + a |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
| class SLM(nn.Module): |
| def __init__(self, vocab_size, n_embd=384, n_head=8, n_layer=10, block_size=64, dropout=0.1): |
| super().__init__() |
| self.token_emb = nn.Embedding(vocab_size, n_embd) |
| self.pos_emb = nn.Embedding(block_size, n_embd) |
| self.blocks = nn.Sequential(*[DenseTransformerBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]) |
| self.ln_f = nn.LayerNorm(n_embd) |
| self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) |
| self.token_emb.weight = self.lm_head.weight |
| self.block_size = block_size |
|
|
| def forward(self, idx, targets=None): |
| b, t = idx.shape |
| x = self.token_emb(idx) + self.pos_emb(torch.arange(t, device=idx.device)) |
| x = self.blocks(x) |
| logits = self.lm_head(self.ln_f(x)) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, new_tokens, temperature=1.0): |
| for _ in range(new_tokens): |
| idx_cond = idx[:, -self.block_size :] |
| logits, _ = self.forward(idx_cond) |
| logits = logits[:, -1, :] / temperature |
| probs = F.softmax(logits, dim=-1) |
| idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1) |
| return idx |
|
|
| def load_model(repo_or_dir): |
| ckpt_dir = repo_or_dir if os.path.isdir(repo_or_dir) else snapshot_download(repo_or_dir) |
| cfg = json.load(open(os.path.join(ckpt_dir, "config.json"))) |
| vocab = cfg["vocab"] |
| arch = cfg["arch"] |
| model = SLM(len(vocab), **arch) |
| model.load_state_dict(torch.load(os.path.join(ckpt_dir, "pytorch_model.bin"), map_location="cpu")) |
| return model, vocab |
|
|
| def sample(model, vocab, prompt, new_tokens=200, temperature=1.0): |
| stoi = {c: i for i, c in enumerate(vocab)} |
| itos = {i: c for c, i in stoi.items()} |
| idx = torch.tensor([[stoi[c] for c in prompt]], dtype=torch.long) |
| out = model.generate(idx, new_tokens, temperature=temperature)[0].tolist() |
| return "".join(itos[i] for i in out) |
|
|
