generate.py

"""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()