File size: 3,998 Bytes
d8c733f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """Load GUIDO-small 200M (bugfix) e genera testo.
Run: python inference_example.py
"""
import sys
from pathlib import Path
import torch
from transformers import AutoTokenizer
# Aggiungi questa cartella al sys.path così Vathos viene trovato
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "vathos")) # bundled Vathos package
sys.path.insert(0, str(HERE)) # modeling_guido.py
# Il modeling è in modeling_guido.py (fork di train_guido_small.py — usa la classe PiCOFormerLM).
# `import modeling_guido as m` registra Vathos in production mode + carica le classi.
import modeling_guido as m
def load_model(ckpt_path: str, device="cuda"):
"""Carica il ckpt fp32, infera l'architettura dai pesi, restituisce model+tokenizer."""
state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
sd = state["model"] if "model" in state else state
# strip eventuali prefissi
sd = {k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k: v for k, v in sd.items()}
# Infer arch dal state_dict (così funziona anche se il config.json non c'è)
hp = m.HP()
hp.head_dim = 64
hp.d_model = sd["backbone.embedder.embedding.weight"].shape[1]
hp.n_heads = hp.d_model // hp.head_dim
hp.d_ff = sd["backbone.blocks.0.channel_mixer.expand.weight"].shape[0]
hp.n_layers = max(int(k.split("blocks.")[1].split(".")[0])
for k in sd if "blocks." in k) + 1
vocab = sd["backbone.embedder.embedding.weight"].shape[0]
print(f"[load] arch inferred: {hp.n_layers}L × {hp.d_model} × {hp.n_heads}h × {hp.d_ff}ff vocab={vocab}")
model = m.PiCOFormerLM(hp, vocab_size=vocab).to(device).bfloat16()
missing, unexpected = model.load_state_dict(sd, strict=False)
if missing: print(f"[load] missing keys ({len(missing)}): {missing[:3]}...")
if unexpected: print(f"[load] unexpected ({len(unexpected)}): {unexpected[:3]}...")
model.eval()
tok = AutoTokenizer.from_pretrained("mistralai/Mathstral-7B-v0.1")
return model, tok
@torch.no_grad()
def generate(model, tok, prompt: str, max_new=200, temperature=0.6, top_p=0.95, top_k=40, device="cuda"):
"""Sampling semplice (T, top-p, top-k)."""
import torch.nn.functional as F
ids = tok(prompt, return_tensors="pt").input_ids.to(device)
eos = tok.eos_token_id or 2
softcap = model.logit_softcap
for _ in range(max_new):
h = model._hidden(ids)
logits = F.linear(h[:, -1, :], model.classifier_weight)
if softcap > 0:
logits = softcap * torch.tanh(logits / softcap)
logits = logits.float() / temperature
if top_k and top_k < logits.size(-1):
kth = logits.topk(top_k, dim=-1).values[:, -1:]
logits = logits.masked_fill(logits < kth, float("-inf"))
probs = F.softmax(logits, dim=-1)
if top_p < 1.0:
sp, si = probs.sort(dim=-1, descending=True)
csum = sp.cumsum(dim=-1)
mask = csum - sp > top_p
sp = sp.masked_fill(mask, 0.0); sp /= sp.sum(dim=-1, keepdim=True)
choice = torch.multinomial(sp, 1)
nxt = si.gather(-1, choice)
else:
nxt = torch.multinomial(probs, 1)
ids = torch.cat([ids, nxt], dim=1)
if nxt.item() == eos:
break
return tok.decode(ids[0], skip_special_tokens=False)
if __name__ == "__main__":
# Adatta il path al checkpoint
CKPT = HERE / "pytorch_model.bin"
device = "cuda" if torch.cuda.is_available() else "cpu"
model, tok = load_model(str(CKPT), device=device)
prompts = [
"Question: Sarah has 5 apples and gives 2 to Tom. How many does she have left?\nAnswer: ",
"Solve: $x^2 + 3x + 2 = 0$. Step 1:",
"The largest planet in the solar system is",
]
for p in prompts:
print("=" * 60)
print("PROMPT:", p)
out = generate(model, tok, p, max_new=120, device=device)
print("OUT: ", out)
|