Add standalone inference (torch + tokenizers + safetensors)

generate.py

@@ -0,0 +1,192 @@
+"""catgirl-1m standalone inference. Needs only: torch, tokenizers, safetensors.
+
+A full fine-tune of Glint-Research/Glint-2 (1,064,737 params) into a
+tsundere-catgirl persona.
+
+Usage:
+    python generate.py "My name is"
+    python generate.py --max-new-tokens 120 --temperature 0.7 "User: hey there"
+
+The base is one shared transformer block looped 8 times, with a tiny per-loop
+LoRA and a loop-embedding that tells it which iteration it is on. The
+checkpoint's loop tables have capacity 16; run it at 8 loops (the value it was
+trained at). More loops produce gibberish.
+"""
+import argparse
+import json
+import math
+import os
+
+import torch
+from tokenizers import Tokenizer
+from torch import nn
+from torch.nn import functional
+from safetensors.torch import load_file
+
+LOOPS = 8
+ATTENTION_WINDOW = 256
+
+
+class SwiGlu(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
+        self.down = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x):
+        gate, up = self.gate_up(x).chunk(2, dim=-1)
+        return self.down(functional.silu(gate) * up)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, cos, sin, qkv_delta):
+        batch, seq_len, dim = x.shape
+        qkv = self.qkv(x) + qkv_delta
+        q, k, v = qkv.split(dim, dim=-1)
+        q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        pos = torch.arange(seq_len, device=x.device)
+        mask = (pos[None, :] <= pos[:, None]) & (pos[:, None] - pos[None, :] < ATTENTION_WINDOW)
+        attended = functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, scale=1.0 / math.sqrt(self.head_dim)
+        )
+        return self.out(attended.transpose(1, 2).reshape(batch, seq_len, dim))
+
+
+def apply_rope(x, cos, sin):
+    x_even, x_odd = x[..., 0::2], x[..., 1::2]
+    return torch.stack(
+        (x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1
+    ).flatten(-2)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, n_heads, ffn_hidden):
+        super().__init__()
+        self.attn_norm = nn.RMSNorm(dim)
+        self.attn = Attention(dim, n_heads)
+        self.ffn_norm = nn.RMSNorm(dim)
+        self.ffn = SwiGlu(dim, ffn_hidden)
+
+    def forward(self, x, cos, sin, qkv_delta):
+        x = x + self.attn(self.attn_norm(x), cos, sin, qkv_delta)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class LoopLora(nn.Module):
+    def __init__(self, dim, rank, max_loops):
+        super().__init__()
+        self.down = nn.ModuleList(nn.Linear(dim, rank, bias=False) for _ in range(max_loops))
+        self.up = nn.ModuleList(nn.Linear(rank, 3 * dim, bias=False) for _ in range(max_loops))
+
+    def forward(self, x, loop_index):
+        i = min(loop_index, len(self.down) - 1)
+        return self.up[i](self.down[i](x))
+
+
+class Indexer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gate = nn.Parameter(torch.tensor([0.1]))
+
+
+class Glint2(nn.Module):
+    def __init__(self, cfg, max_loops):
+        super().__init__()
+        self.cfg = cfg
+        self.max_loops = max_loops
+        self.embed = nn.Embedding(cfg["vocab_size"], cfg["dim"])
+        self.indexer = Indexer()
+        self.shared = Block(cfg["dim"], cfg["n_heads"], cfg["ffn_hidden"])
+        self.loop_lora = LoopLora(cfg["dim"], cfg["lora_rank"], max_loops)
+        self.loop_embed = nn.Embedding(max_loops, cfg["dim"])
+        self.final_norm = nn.RMSNorm(cfg["dim"])
+
+    def rope(self, seq_len):
+        head_dim = self.cfg["dim"] // self.cfg["n_heads"]
+        positions = torch.arange(seq_len, dtype=torch.float32)
+        inv_freq = 1.0 / (
+            self.cfg["rope_base"]
+            ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+        )
+        angles = torch.outer(positions, inv_freq)
+        return torch.cos(angles), torch.sin(angles)
+
+    def forward(self, tokens, loops=LOOPS, return_hidden=False):
+        cos, sin = self.rope(tokens.shape[1])
+        x = self.embed(tokens)
+        for loop_index in range(loops):
+            clamped = min(loop_index, self.max_loops - 1)
+            gated = x + self.loop_embed.weight[clamped]
+            delta = self.loop_lora(gated, loop_index)
+            x = self.shared(gated, cos, sin, delta)
+        x = self.final_norm(x)
+        if return_hidden:
+            return x
+        return functional.linear(x, self.embed.weight)
+
+
+def load_model(model_dir):
+    with open(os.path.join(model_dir, "config.json")) as f:
+        cfg = json.load(f)
+    state = load_file(os.path.join(model_dir, "catgirl-1m.safetensors"))
+    model = Glint2(cfg, max_loops=cfg["max_loops"])
+    model.load_state_dict(state)
+    model.eval()
+    return model
+
+
+def sample(logits, ids, temperature, top_k, repetition_penalty, repetition_window):
+    if repetition_penalty != 1.0:
+        for token in set(ids[-repetition_window:]):
+            if logits[token] > 0:
+                logits[token] /= repetition_penalty
+            else:
+                logits[token] *= repetition_penalty
+    if temperature <= 0:
+        return int(logits.argmax())
+    values, indices = logits.topk(top_k)
+    probs = functional.softmax(values / temperature, dim=-1)
+    return int(indices[torch.multinomial(probs, 1)])
+
+
+def main():
+    parser = argparse.ArgumentParser(description="catgirl-1m text generation")
+    parser.add_argument("prompt")
+    parser.add_argument("--model-dir", default=os.path.dirname(os.path.abspath(__file__)))
+    parser.add_argument("--max-new-tokens", type=int, default=100)
+    parser.add_argument("--temperature", type=float, default=0.15, help="0 = greedy")
+    parser.add_argument("--top-k", type=int, default=5)
+    parser.add_argument("--repetition-penalty", type=float, default=1.05)
+    parser.add_argument("--repetition-window", type=int, default=128)
+    parser.add_argument("--loops", type=int, default=LOOPS)
+    args = parser.parse_args()
+
+    tokenizer = Tokenizer.from_file(os.path.join(args.model_dir, "tokenizer.json"))
+    model = load_model(args.model_dir)
+    ids = tokenizer.encode(args.prompt).ids
+    print(args.prompt, end="", flush=True)
+    with torch.no_grad():
+        for _ in range(args.max_new_tokens):
+            logits = model(torch.tensor([ids]), loops=args.loops)[0, -1]
+            token = sample(
+                logits, ids, args.temperature, args.top_k,
+                args.repetition_penalty, args.repetition_window,
+            )
+            new_text = tokenizer.decode(ids + [token])[len(tokenizer.decode(ids)):]
+            ids.append(token)
+            print(new_text, end="", flush=True)
+    print()
+
+
+if __name__ == "__main__":
+    main()