Update README.md

README.md

@@ -77,25 +77,142 @@ pip install transformers torch
 ## Esempio di utilizzo in Python
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-
-# Carica modello e tokenizer
-model_name = "Mattimax/PicoDAC"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Input di esempio
-input_text = "Ciao, come va?"
-inputs = tokenizer(input_text, return_tensors="pt")
-
-# Generazione
-with torch.no_grad():
-    outputs = model.generate(**inputs, max_length=64, do_sample=True, temperature=0.8)
-
-# Decodifica
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(response)
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from transformers import PreTrainedTokenizerFast
+import os
+
+# ----------------------------
+# Config modello
+# ----------------------------
+VOCAB_SIZE = 1920
+D_MODEL = 240
+N_LAYERS = 6
+N_HEADS = 6
+D_HEAD = 40
+D_FF = 960
+N_CTX = 64
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Usando device: {device}")
+
+# ----------------------------
+# Scarica tokenizer e modello da HF
+# ----------------------------
+repo_id = "Mattimax/PicoDAC-IT"
+
+tokenizer_path = hf_hub_download(repo_id, "tokenizer.json")
+model_path = hf_hub_download(repo_id, "model.safetensors")
+# opzionale: scales_path = hf_hub_download(repo_id, "best/scales.json")
+
+# ----------------------------
+# Carica tokenizer
+# ----------------------------
+tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
+tokenizer.add_special_tokens({
+    "pad_token": "<PAD>",
+    "bos_token": "<BOS>",
+    "eos_token": "<EOS>",
+    "sep_token": "<SEP>",
+    "unk_token": "<UNK>"
+})
+
+# ----------------------------
+# RMSNorm
+# ----------------------------
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    
+    def forward(self, x):
+        norm = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return x / norm * self.weight
+
+# ----------------------------
+# Transformer Block
+# ----------------------------
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.ln1 = RMSNorm(D_MODEL)
+        self.ln2 = RMSNorm(D_MODEL)
+        self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, batch_first=True)
+        self.ffn = nn.Sequential(
+            nn.Linear(D_MODEL, D_FF),
+            nn.SiLU(),
+            nn.Linear(D_FF, D_MODEL)
+        )
+    
+    def forward(self, x, mask=None):
+        attn_out, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=mask)
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+# ----------------------------
+# TinyGPT
+# ----------------------------
+class TinyGPT(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tok_emb = nn.Embedding(VOCAB_SIZE, D_MODEL)
+        self.pos_emb = nn.Embedding(N_CTX, D_MODEL)
+        self.layers = nn.ModuleList([TransformerBlock() for _ in range(N_LAYERS)])
+        self.ln_f = RMSNorm(D_MODEL)
+        self.head = nn.Linear(D_MODEL, VOCAB_SIZE, bias=False)
+        self.head.weight = self.tok_emb.weight
+    
+    def forward(self, idx):
+        B, T = idx.shape
+        pos = torch.arange(T, device=idx.device).unsqueeze(0)
+        x = self.tok_emb(idx) + self.pos_emb(pos)
+        mask = torch.triu(torch.ones(T, T, device=idx.device) * float('-inf'), diagonal=1)
+        for layer in self.layers:
+            x = layer(x, mask=mask)
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits
+
+# ----------------------------
+# Carica pesi
+# ----------------------------
+state_dict = load_file(model_path)
+model = TinyGPT()
+model.load_state_dict(state_dict, strict=False)
+model.to(device)
+model.eval()
+
+# ----------------------------
+# Funzione generazione
+# ----------------------------
+def generate(prompt, max_new_tokens=50, temperature=0.7):
+    model_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    for _ in range(max_new_tokens):
+        logits = model(model_input)
+        next_token_logits = logits[0, -1, :] / temperature
+        probs = F.softmax(next_token_logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+        model_input = torch.cat([model_input, next_token.unsqueeze(0)], dim=1)
+        if next_token.item() == tokenizer.eos_token_id:
+            break
+    return tokenizer.decode(model_input[0], skip_special_tokens=True)
+
+# ----------------------------
+# Loop chat
+# ----------------------------
+print("Chat PicoDAC-IT (scrivi 'exit' per uscire)")
+while True:
+    user_input = input("Tu: ").strip()
+    if user_input.lower() in ["exit", "quit"]:
+        print("Chiusura chat. Ciao!")
+        break
+    response = generate(user_input)
+    print(f"PicoDAC-IT: {response}")
 ```
 
 **Suggerimenti pratici:**