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