| """Test del modello distillato con prompt personalizzati da riga di comando.""" |
| from __future__ import annotations |
|
|
| import os |
| import sys |
|
|
| |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| sys.stdout.reconfigure(encoding="utf-8", errors="replace") |
|
|
| import torch |
| from tokenizers import Tokenizer |
|
|
| from src.model.config import BoccaccioConfig |
| from src.model.transformer import BoccaccioForCausalLM |
|
|
|
|
| def load_model(path: str, config: BoccaccioConfig) -> BoccaccioForCausalLM: |
| model = BoccaccioForCausalLM(config).to("cuda") |
| sd = torch.load(path, map_location="cuda", weights_only=True) |
| if isinstance(sd, dict) and "state_dict" in sd: |
| sd = sd["state_dict"] |
| cleaned = {} |
| for k, v in sd.items(): |
| nk = k |
| |
| if nk.startswith("_orig_mod."): |
| nk = nk[len("_orig_mod."):] |
| |
| |
| elif nk.startswith("model.model."): |
| nk = nk[len("model."):] |
| cleaned[nk] = v |
| missing, unexpected = model.load_state_dict(cleaned, strict=False) |
| if missing: |
| print(f"WARNING: {len(missing)} missing keys") |
| if unexpected: |
| print(f"WARNING: {len(unexpected)} unexpected keys") |
| model.eval() |
| return model |
|
|
|
|
| def generate(model, tokenizer, prompt, max_new=300, temp=0.3, rep=1.2): |
| ids = tokenizer.encode(prompt).ids |
| x = torch.tensor([ids], device="cuda") |
| with torch.no_grad(): |
| for _ in range(max_new): |
| out = model(x) |
| logits = out["logits"][0, -1, :] |
| for tid in set(x[0].tolist()): |
| logits[tid] /= rep |
| if temp > 0: |
| logits = logits / temp |
| v, _ = torch.topk(logits, 40) |
| logits[logits < v[-1]] = float("-inf") |
| probs = torch.softmax(logits, dim=-1) |
| tid = torch.multinomial(probs, num_samples=1).item() |
| else: |
| tid = torch.argmax(logits).item() |
| if tid == 2: |
| break |
| x = torch.cat([x, torch.tensor([[tid]], device="cuda")], dim=1) |
| new = x[0].cpu().tolist()[len(ids):] |
| return tokenizer.decode(new).strip() |
|
|
|
|
| def main(): |
| config = BoccaccioConfig.from_yaml("configs/model.yaml", variant="model") |
| tok = Tokenizer.from_file("tokenizer/boccaccio-32k.json") |
|
|
| import os |
| ckpt_dir = "checkpoints/finetune_distill" |
| |
| model_pt = os.path.join(ckpt_dir, "model.pt") |
| if os.path.exists(model_pt): |
| ckpt_path = model_pt |
| else: |
| ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".ckpt")]) |
| if not ckpts: |
| print("Nessun checkpoint trovato in checkpoints/finetune_distill/") |
| return |
| ckpt_path = os.path.join(ckpt_dir, ckpts[-1]) |
| print(f"Using: {ckpt_path}\n", flush=True) |
| model = load_model(ckpt_path, config) |
|
|
| print("Modalita' chat. Digita 'quit' per uscire, 'ctx' per aggiungere contesto.\n") |
|
|
| context = "" |
| while True: |
| try: |
| user_input = input("Tu> ").strip() |
| except (EOFError, KeyboardInterrupt): |
| break |
| if not user_input: |
| continue |
| if user_input.lower() == "quit": |
| break |
| if user_input.lower().startswith("ctx "): |
| context = user_input[4:].strip() |
| print(f"[Contesto impostato: {context[:80]}...]\n") |
| continue |
| if user_input.lower() == "ctx": |
| context = "" |
| print("[Contesto cancellato]\n") |
| continue |
|
|
| if context: |
| prompt = f"### Contesto: {context}\n\n### Utente: {user_input}\n### Assistente:" |
| else: |
| prompt = f"### Utente: {user_input}\n### Assistente:" |
| resp = generate(model, tok, prompt) |
| print(f"\nAssistente: {resp}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
