Upload modeling_aetherstory.py with huggingface_hub

modeling_aetherstory.py

@@ -0,0 +1,133 @@
+"""Minimal load+generate helper for the AetherStory model on Hugging Face.
+
+Users can copy this file next to ``model.safetensors`` / ``config.json`` /
+``tokenizer.json`` and run::
+
+    from modeling_aetherstory import StoryTeller
+    t = StoryTeller.from_dir(".")
+    print(t("Once upon a time"))
+"""
+
+import json
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+
+
+# --- model definition (mirrors src/model.py) -------------------------------
+class _Attn(nn.Module):
+    def __init__(s, d, h, drop):
+        super().__init__()
+        s.h, s.hd, s.scale = h, d // h, (d // h) ** -0.5
+        s.qkv = nn.Linear(d, 3 * d, bias=False)
+        s.proj = nn.Linear(d, d, bias=False)
+        s.drop = nn.Dropout(drop)
+    def forward(s, x):
+        B, T, C = x.shape
+        qkv = s.qkv(x).reshape(B, T, 3, s.h, s.hd).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        att = (q @ k.transpose(-2, -1)) * s.scale
+        att = att.masked_fill(~torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool)), float("-inf"))
+        att = s.drop(F.softmax(att, dim=-1))
+        y = (att @ v).transpose(1, 2).contiguous().reshape(B, T, C)
+        return s.drop(s.proj(y))
+
+
+class _FFN(nn.Module):
+    def __init__(s, d, f, drop):
+        super().__init__()
+        s.fc1, s.fc2, s.drop = nn.Linear(d, f, bias=False), nn.Linear(f, d, bias=False), nn.Dropout(drop)
+    def forward(s, x):
+        return s.drop(s.fc2(F.gelu(s.fc1(x))))
+
+
+class _Block(nn.Module):
+    def __init__(s, d, h, f, drop):
+        super().__init__()
+        s.ln1, s.attn = nn.LayerNorm(d), _Attn(d, h, drop)
+        s.ln2, s.ffn = nn.LayerNorm(d), _FFN(d, f, drop)
+    def forward(s, x):
+        x = x + s.attn(s.ln1(x)); x = x + s.ffn(s.ln2(x)); return x
+
+
+class AetherStoryModel(nn.Module):
+    def __init__(s, cfg):
+        super().__init__()
+        s.cfg = cfg
+        s.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["d_model"])
+        s.pos_emb = nn.Parameter(torch.zeros(1, cfg["max_seq_len"], cfg["d_model"]))
+        s.blocks = nn.ModuleList([_Block(cfg["d_model"], cfg["n_heads"], cfg["ffn_dim"], cfg["dropout"]) for _ in range(cfg["n_layers"])])
+        s.ln_f = nn.LayerNorm(cfg["d_model"])
+        if cfg.get("tie_embeddings", True):
+            s.head_bias = nn.Parameter(torch.zeros(cfg["vocab_size"])); s.lm_head = None
+        else:
+            s.head_bias = None; s.lm_head = nn.Linear(cfg["d_model"], cfg["vocab_size"], bias=False)
+    def forward(s, idx, targets=None):
+        B, T = idx.shape
+        x = s.tok_emb(idx) + s.pos_emb[:, :T, :]
+        for b in s.blocks: x = b(x)
+        x = s.ln_f(x)
+        logits = s.lm_head(x) if s.lm_head is not None else (x @ s.tok_emb.weight.t() + s.head_bias)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.reshape(-1, s.cfg["vocab_size"]), targets.reshape(-1), ignore_index=s.cfg.get("pad_token_id", 0))
+        return logits, loss
+    @torch.no_grad()
+    def generate(s, idx, max_new, temperature=0.9, top_k=40, eos_token_id=None):
+        s.eval()
+        for _ in range(max_new):
+            ic = idx if idx.size(1) <= s.cfg["max_seq_len"] else idx[:, -s.cfg["max_seq_len"]:]
+            logits, _ = s(ic)
+            logits = logits[:, -1, :] / max(temperature, 1e-5)
+            if top_k and top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+            nxt = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+            idx = torch.cat([idx, nxt], dim=1)
+            if eos_token_id is not None and (nxt == eos_token_id).all(): break
+        return idx
+
+
+class StoryTeller:
+    @classmethod
+    def from_dir(cls, d="."):
+        d = Path(d)
+        cfg = json.loads((d / "config.json").read_text(encoding="utf-8"))
+        tok_payload = json.loads((d / "tokenizer.json").read_text(encoding="utf-8"))
+        model = AetherStoryModel(cfg)
+        state = load_file(str(d / "model.safetensors"))
+        model.load_state_dict({k: v for k, v in state.items()})
+        model.eval()
+        w2i = {}
+        for word, idx in tok_payload["special_tokens"].items(): w2i[word] = idx
+        for idx, word in enumerate(tok_payload["vocab"]):
+            if idx not in w2i: w2i[word] = idx
+        i2w = {v: k for k, v in w2i.items()}
+        inst = cls(); inst.model = model; inst.w2i = w2i; inst.i2w = i2w
+        inst.bos = tok_payload["special_tokens"]["<bos>"]; inst.eos = tok_payload["special_tokens"]["<eos>"]
+        return inst
+    def _encode(self, text):
+        import re
+        ids = [self.bos]
+        for t in re.findall(r"\w+|[^\w\s]|\s+", text.lower()):
+            if t.strip() and t in self.w2i: ids.append(self.w2i[t])
+        return ids
+    def _decode(self, ids):
+        out = []
+        for i in ids:
+            if i == self.bos or i == 0: continue
+            if i == self.eos: break
+            w = self.i2w.get(i, "")
+            if w.startswith("<") and w.endswith(">"): continue
+            out.append(w)
+        t = " ".join(out)
+        import re
+        return re.sub(r"\s+([,.;:!?\'\"()])", r"\1", t).strip()
+    def __call__(self, prompt, max_tokens=80, temperature=0.9, top_k=40, seed=None):
+        if seed is not None: torch.manual_seed(seed)
+        ids = self._encode(prompt)
+        idx = torch.tensor([ids], dtype=torch.long)
+        out = self.model.generate(idx, max_tokens, temperature, top_k, self.eos)
+        return self._decode(out[0].tolist())