"""
Standalone inference script for AxionLab-official/RizzAura-100k (RizzLLaMA arch).
No trust_remote_code. Requires: torch, huggingface_hub, tokenizers
"""
import json
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
REPO_ID = "AxionLab-official/RizzAura-100k"
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x):
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return x * self.weight
def rope_freqs(head_dim, max_len, base=10000.0, device="cpu"):
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
t = torch.arange(max_len, device=device).float()
freqs = torch.outer(t, inv_freq)
return torch.cos(freqs), torch.sin(freqs)
def apply_rope(x, cos, sin):
x1, x2 = x[..., 0::2], x[..., 1::2]
cos = cos[None, None, :x.shape[2], :]
sin = sin[None, None, :x.shape[2], :]
rx1 = x1 * cos - x2 * sin
rx2 = x1 * sin + x2 * cos
out = torch.stack([rx1, rx2], dim=-1).flatten(-2)
return out
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.proj = nn.Linear(dim, dim, bias=False)
def forward(self, x, cos, sin, mask):
B, T, C = x.shape
qkv = self.qkv(x)
q, k, v = qkv.split(C, dim=-1)
q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
att = att.masked_fill(mask[:, :, :T, :T] == 0, float("-inf"))
att = F.softmax(att, dim=-1)
out = att @ v
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.proj(out)
class FFN(nn.Module):
def __init__(self, dim, hidden_mult):
super().__init__()
hidden = int(dim * hidden_mult)
self.w1 = nn.Linear(dim, hidden, bias=False)
self.w3 = nn.Linear(dim, hidden, bias=False)
self.w2 = nn.Linear(hidden, dim, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class Block(nn.Module):
def __init__(self, dim, n_heads, hidden_mult):
super().__init__()
self.attn_norm = RMSNorm(dim)
self.attn = Attention(dim, n_heads)
self.ffn_norm = RMSNorm(dim)
self.ffn = FFN(dim, hidden_mult)
def forward(self, x, cos, sin, mask):
x = x + self.attn(self.attn_norm(x), cos, sin, mask)
x = x + self.ffn(self.ffn_norm(x))
return x
class RizzLLaMA(nn.Module):
def __init__(self, vocab_size, dim, n_layers, n_heads, hidden_mult, max_len):
super().__init__()
self.max_len = max_len
self.tok_emb = nn.Embedding(vocab_size, dim)
self.blocks = nn.ModuleList(
[Block(dim, n_heads, hidden_mult) for _ in range(n_layers)]
)
self.norm = RMSNorm(dim)
self.head = nn.Linear(dim, vocab_size, bias=False)
head_dim = dim // n_heads
cos, sin = rope_freqs(head_dim, max_len)
self.register_buffer("cos", cos, persistent=False)
self.register_buffer("sin", sin, persistent=False)
self.register_buffer(
"mask", torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len),
persistent=False,
)
def forward(self, idx):
x = self.tok_emb(idx)
for blk in self.blocks:
x = blk(x, self.cos, self.sin, self.mask)
x = self.norm(x)
return self.head(x)
def load_model(device="cpu"):
model_path = hf_hub_download(REPO_ID, "model.pt")
ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
cfg = ckpt["config"]
model = RizzLLaMA(
vocab_size=ckpt["vocab_size"],
dim=cfg["dim"],
n_layers=cfg["n_layers"],
n_heads=cfg["n_heads"],
hidden_mult=cfg["hidden_mult"],
max_len=ckpt["max_len"],
)
model.load_state_dict(ckpt["model_state"])
model.to(device).eval()
return model, ckpt
def load_tokenizer():
tok_path = hf_hub_download(REPO_ID, "tokenizer.json")
return Tokenizer.from_file(tok_path)
@torch.no_grad()
def generate(model, tokenizer, prompt, max_new_tokens=64, temperature=0.8, top_p=0.9, device="cpu"):
bos = tokenizer.token_to_id("<|bos|>")
eos = tokenizer.token_to_id("<|eos|>")
user_tok = tokenizer.token_to_id("<|user|>")
assistant_tok = tokenizer.token_to_id("<|assistant|>")
ids = [bos, user_tok] + tokenizer.encode(prompt).ids + [assistant_tok]
idx = torch.tensor([ids], dtype=torch.long, device=device)
for _ in range(max_new_tokens):
idx_cond = idx[:, -model.max_len:]
logits = model(idx_cond)[:, -1, :]
logits = logits / max(temperature, 1e-5)
probs = F.softmax(logits, dim=-1)
sorted_probs, sorted_idx = torch.sort(probs, descending=True)
cum_probs = torch.cumsum(sorted_probs, dim=-1)
cutoff = cum_probs > top_p
cutoff[..., 1:] = cutoff[..., :-1].clone()
cutoff[..., 0] = False
sorted_probs[cutoff] = 0.0
sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
next_sorted = torch.multinomial(sorted_probs, num_samples=1)
next_id = sorted_idx.gather(-1, next_sorted)
idx = torch.cat([idx, next_id], dim=1)
if next_id.item() == eos or idx.shape[1] >= model.max_len:
break
out_ids = idx[0].tolist()
gen_ids = out_ids[out_ids.index(assistant_tok) + 1:]
if eos in gen_ids:
gen_ids = gen_ids[: gen_ids.index(eos)]
return tokenizer.decode(gen_ids)
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
model, ckpt = load_model(device)
tokenizer = load_tokenizer()
print(f"RizzAura-100k carregado ({sum(p.numel() for p in model.parameters())} params) em {device}")
while True:
prompt = input("\nuser> ")
if not prompt.strip():
break
reply = generate(model, tokenizer, prompt, device=device)
print(f"assistant> {reply}")
