08 - Fine-tuning LLM Localmente: LoRA su Consumer GPU

from transformers import (
    AutoTokenizer,
    AutoModelForMaskedLM,
    DataCollatorForLanguageModeling,
    DataCollatorForWholeWordMask,
    TrainingArguments,
    Trainer
)
from datasets import load_dataset, Dataset
import torch

# Modello base da adattare
BASE_MODEL = "dbmdz/bert-base-italian-cased"
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForMaskedLM.from_pretrained(BASE_MODEL)

# Corpus del dominio (es. testi medici italiani, senza label)
# In pratica caricar da file/database, qui esemplificativo
domain_texts = [
    "Il paziente presenta sintomi di insufficienza cardiaca congestizia...",
    "La diagnosi differenziale include patologie neoplastiche e infiammatorie...",
    "La terapia farmacologica prevede la somministrazione di ACE-inibitori...",
    "Esame istologico evidenzia presenza di cellule atipiche a livello...",
    # ... migliaia di testi medici
]

# Tokenizza il corpus con chunking per testi lunghi
def tokenize_corpus(examples, chunk_size=512):
    """Tokenizza e divide in chunk da max 512 token."""
    tokenized = tokenizer(
        examples["text"],
        truncation=False,
        return_special_tokens_mask=True
    )
    # Crea chunk di lunghezza fissa
    all_input_ids = []
    all_attention_masks = []
    all_special_tokens_masks = []

    for ids, attn, stm in zip(
        tokenized["input_ids"],
        tokenized["attention_mask"],
        tokenized["special_tokens_mask"]
    ):
        for i in range(0, len(ids), chunk_size):
            chunk = ids[i:i+chunk_size]
            if len(chunk) >= 64:  # ignora chunk troppo corti
                # Padding al chunk_size
                padded = chunk + [tokenizer.pad_token_id] * (chunk_size - len(chunk))
                attn_chunk = [1] * len(chunk) + [0] * (chunk_size - len(chunk))
                stm_chunk = stm[i:i+chunk_size] + [1] * (chunk_size - len(chunk))

                all_input_ids.append(padded)
                all_attention_masks.append(attn_chunk)
                all_special_tokens_masks.append(stm_chunk)

    return {
        "input_ids": all_input_ids,
        "attention_mask": all_attention_masks,
        "special_tokens_mask": all_special_tokens_masks
    }

domain_dataset = Dataset.from_dict({"text": domain_texts})
tokenized_corpus = domain_dataset.map(
    tokenize_corpus,
    batched=True,
    remove_columns=["text"],
    batch_size=100
)

# Data collator per MLM standard (maschera il 15% dei token)
data_collator_mlm = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15
)

# Data collator per Whole Word Masking (più efficace per BERT)
data_collator_wwm = DataCollatorForWholeWordMask(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15
)

# Training DAPT con Whole Word Masking
dapt_args = TrainingArguments(
    output_dir="./models/bert-italian-medical-dapt",
    num_train_epochs=5,              # più epoche per DAPT
    per_device_train_batch_size=16,
    learning_rate=5e-5,             # più alto per DAPT che per fine-tuning
    warmup_ratio=0.05,
    weight_decay=0.01,
    save_steps=500,
    save_total_limit=2,
    fp16=True,
    report_to="none",
    logging_steps=100,
    # Strategia di valutazione opzionale
    eval_strategy="no",
)

dapt_trainer = Trainer(
    model=model,
    args=dapt_args,
    train_dataset=tokenized_corpus,
    data_collator=data_collator_wwm
)

print("Avvio DAPT training...")
dapt_trainer.train()

# Salva il modello DAPT (da usare come base per il fine-tuning task-specifico)
model.save_pretrained("./models/bert-italian-medical-dapt")
tokenizer.save_pretrained("./models/bert-italian-medical-dapt")

print("\nDAPT completato!")
print("Il modello ha ora acquisito il vocabolario medico italiano.")
print("Prossimo step: fine-tuning task-specifico (NER, classificazione, etc.)")

from peft import (
    LoraConfig,
    get_peft_model,
    TaskType,
    PeftModel,
    prepare_model_for_kbit_training
)
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer
from datasets import Dataset
import evaluate
import numpy as np

# Fine-tuning di BERT con LoRA per classificazione di contratti
# Usiamo il modello DAPT se disponibile, altrimenti il modello base
MODEL = "./models/bert-italian-medical-dapt"  # o "dbmdz/bert-base-italian-cased"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL,
    num_labels=5,   # 5 categorie di documenti medici
    id2label={
        0: "anamnesi",
        1: "diagnosi",
        2: "terapia",
        3: "referto_esame",
        4: "lettera_dimissione"
    },
    label2id={
        "anamnesi": 0,
        "diagnosi": 1,
        "terapia": 2,
        "referto_esame": 3,
        "lettera_dimissione": 4
    }
)

# Configurazione LoRA ottimizzata
lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=16,                        # rango ottimale per task di classificazione
    lora_alpha=32,               # scaling = lora_alpha / r = 2.0
    target_modules=[             # layer da modificare in BERT
        "query",                 # proiezione query in multi-head attention
        "key",                   # proiezione key
        "value",                 # proiezione value
        "dense"                  # layer denso nell'output di attention e FFN
    ],
    lora_dropout=0.05,
    bias="none",                 # "none" = nessun bias trainabile in LoRA
    modules_to_save=["classifier"]   # testa classificazione SEMPRE trainata completamente
)

peft_model = get_peft_model(model, lora_config)
peft_model.print_trainable_parameters()
# trainable params: 592,898 || all params: 124,647,170 || trainable%: 0.4756%

# Verifica la struttura del modello PEFT
print("\nLayer trainabili:")
for name, param in peft_model.named_parameters():
    if param.requires_grad:
        print(f"  {name}: {param.shape}")

# Dataset di addestramento (dominio medico)
train_texts = [
    "Il paziente riferisce dolore toracico da 3 giorni, ipertensione arteriosa nota...",
    "Diagnosi: fibrillazione atriale parossistica con risposta ventricolare elevata...",
    "Terapia: amoxicillina 1g x 2/die per 7 giorni, paracetamolo 1g al bisogno...",
    "RMN encefalo: presenza di lesione ischemica acuta in territorio della ACM destra...",
    "Si dimette in condizioni stabili con prescrizione di follow-up cardiologico...",
]
train_labels = [0, 1, 2, 3, 4]

def tokenize_fn(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=256
    )

train_ds = Dataset.from_dict({"text": train_texts, "label": train_labels})
train_ds = train_ds.map(tokenize_fn, batched=True, remove_columns=["text"])
train_ds.set_format("torch")

# Metriche
accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy.compute(predictions=preds, references=labels)["accuracy"],
        "f1_macro": f1.compute(predictions=preds, references=labels, average="macro")["f1"]
    }

# TrainingArguments per LoRA: LR più alto e più epoche rispetto a full fine-tuning
args = TrainingArguments(
    output_dir="./results/bert-medical-lora",
    num_train_epochs=20,          # più epoche per dataset piccoli con LoRA
    per_device_train_batch_size=8,
    learning_rate=3e-4,           # LoRA usa LR più alto (3e-4 invece di 2e-5)
    warmup_ratio=0.1,
    weight_decay=0.01,
    eval_strategy="no",           # no eval su dataset tiny
    save_strategy="epoch",
    save_total_limit=2,
    fp16=True,
    report_to="none",
    seed=42
)

trainer = Trainer(
    model=peft_model,
    args=args,
    train_dataset=train_ds,
    compute_metrics=compute_metrics
)

print("\nAvvio LoRA fine-tuning...")
trainer.train()

# Salva solo i pesi LoRA (~2MB invece di ~500MB!)
peft_model.save_pretrained("./models/bert-medical-lora-adapter")
tokenizer.save_pretrained("./models/bert-medical-lora-adapter")
print("\nSalvati adapter LoRA in ./models/bert-medical-lora-adapter")

# pip install bitsandbytes accelerate peft trl transformers
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig
from datasets import Dataset
import torch

# =========================================
# Configurazione quantizzazione 4-bit
# =========================================
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # Normal Float 4 (ottimale per LLM)
    bnb_4bit_compute_dtype=torch.bfloat16,  # compute in bfloat16 per stabilità
    bnb_4bit_use_double_quant=True,     # double quantization (risparmia ~0.4 bit/param)
)

# Carica modello in 4-bit
# Riduzione VRAM: Mistral-7B da ~14GB a ~5GB!
MODEL_NAME = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # necessario per batch padding
tokenizer.padding_side = "right"           # padding a destra per generazione

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",              # auto-distribuzione su GPU disponibili
    trust_remote_code=True,
    attn_implementation="flash_attention_2"  # Flash Attention 2 se disponibile
)

print(f"Memoria GPU: {torch.cuda.memory_allocated()/1e9:.2f}GB")

# Prepara per il training con kbit quantization
model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True  # risparmia memoria aggiuntiva
)

# Configurazione LoRA per LLM (tutti i layer attention + MLP)
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj",    # attention layers
        "o_proj",                          # output projection
        "gate_proj", "up_proj", "down_proj"  # MLP (SwiGLU) layers di Mistral
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

peft_model = get_peft_model(model, lora_config)
peft_model.print_trainable_parameters()
# Per Mistral-7B: trainable params ~83M || all params ~3.75B || trainable%: 2.24%

# =========================================
# Dataset in formato instruction-following
# =========================================
def format_instruction(instruction: str, input_text: str, output: str) -> str:
    """Formatta un esempio nel formato Alpaca per instruction tuning."""
    if input_text:
        return (
            f"### Istruzione:\n{instruction}\n\n"
            f"### Input:\n{input_text}\n\n"
            f"### Risposta:\n{output}"
        )
    return (
        f"### Istruzione:\n{instruction}\n\n"
        f"### Risposta:\n{output}"
    )

train_examples = [
    {
        "text": format_instruction(
            instruction="Classifica questo testo medico nella categoria appropriata.",
            input_text="Il paziente presenta febbre a 38.5C, tosse secca persistente da 5 giorni...",
            output="anamnesi"
        )
    },
    {
        "text": format_instruction(
            instruction="Estrai i farmaci prescritti e le relative posologie.",
            input_text="Terapia: paracetamolo 500mg x 3/die, amoxicillina 1g x 2/die per 7gg.",
            output="Farmaci: paracetamolo 500mg (3 volte al giorno), amoxicillina 1g (2 volte al giorno per 7 giorni)"
        )
    },
    {
        "text": format_instruction(
            instruction="Riassumi la diagnosi principale in una frase.",
            input_text="RMN cerebrale: lesione iperintensa in T2/FLAIR in sede occipito-parietale destra...",
            output="Ictus ischemico acuto in territorio della arteria cerebrale posteriore destra."
        )
    },
]

train_dataset = Dataset.from_list(train_examples)

# =========================================
# SFTTrainer per supervised fine-tuning
# =========================================
sft_config = SFTConfig(
    output_dir="./models/mistral-medical-qlora",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,    # effective batch = 4*4 = 16
    warmup_ratio=0.1,
    learning_rate=2e-4,               # QLoRA usa LR alto
    fp16=False,
    bf16=True,                        # bfloat16 più stabile di fp16 per LLM
    logging_steps=10,
    optim="paged_adamw_32bit",        # ottimizzatore paginato (risparmia ~8GB!)
    lr_scheduler_type="cosine",
    max_seq_length=512,               # lunghezza massima sequenza
    dataset_text_field="text",
    packing=True,                     # packing: concatena esempi corti per efficienza
    report_to="none",
    save_steps=100,
    save_total_limit=2
)

trainer = SFTTrainer(
    model=peft_model,
    train_dataset=train_dataset,
    args=sft_config,
)

print("\nAvvio QLoRA fine-tuning di Mistral-7B...")
trainer.train()
trainer.save_model("./models/mistral-medical-qlora")
print("QLoRA fine-tuning completato!")

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import TrainingArguments, Trainer, EarlyStoppingCallback
from datasets import Dataset
import numpy as np
import torch

# =========================================
# Strategia 1: SetFit per few-shot learning (2-64 esempi!)
# =========================================
from setfit import SetFitModel, SetFitTrainer

# SetFit addestra sentence transformer + classificatore con POCHISSIMI esempi
setfit_model = SetFitModel.from_pretrained(
    "nickprock/sentence-bert-base-italian-uncased"
)

# Solo 8 esempi per classe (64 totali per 8 classi)!
train_data = {
    "text": ["Testo di esempio 1", "Testo di esempio 2"] * 4,
    "label": [0, 1, 0, 1, 0, 1, 0, 1]
}

setfit_trainer = SetFitTrainer(
    model=setfit_model,
    train_dataset=Dataset.from_dict(train_data),
    num_iterations=20,          # numero di coppie di contrasting
    num_epochs=1,               # epoche per la testa di classificazione
    batch_size=16,
)
setfit_trainer.train()

# =========================================
# Strategia 2: Layer freezing progressivo
# =========================================
def progressive_unfreeze(model, epoch, total_epochs, num_layers=12):
    """
    Gradual unfreezing: sblocca i layer dall'ultimo al primo man mano
    che il training avanza. Questo previene catastrophic forgetting
    e migliora le performance con pochi dati.
    """
    # Quanti layer sbloccare in questa epoch
    layers_to_unfreeze = max(1, int(num_layers * epoch / total_epochs))
    first_layer_to_unfreeze = num_layers - layers_to_unfreeze

    # Congela/scongela in modo progressivo
    for i, layer in enumerate(model.bert.encoder.layer):
        if i >= first_layer_to_unfreeze:
            for param in layer.parameters():
                param.requires_grad = True
        else:
            for param in layer.parameters():
                param.requires_grad = False

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  Epoch {epoch}: sblocati layer {first_layer_to_unfreeze}-{num_layers-1}, "
          f"trainable params: {trainable:,}")

# =========================================
# Strategia 3: Learning rate differenziali per layer
# =========================================
from torch.optim import AdamW

def get_layerwise_lr(model, base_lr=2e-5, lr_decay=0.75):
    """
    Learning rate decrescente per i layer più bassi.
    I layer più bassi (syntax, basic semantics) cambiano poco,
    i layer alti (task-specific features) cambiano molto.
    """
    # Embedding layer
    params = [{"params": model.bert.embeddings.parameters(), "lr": base_lr * (lr_decay ** 13)}]

    # Encoder layers (da 0 a 11 per BERT-base)
    for i, layer in enumerate(model.bert.encoder.layer):
        lr = base_lr * (lr_decay ** (12 - i))   # LR crescente per layer più alti
        params.append({"params": layer.parameters(), "lr": lr})

    # Pooler e classifier: LR massimo
    params.append({"params": model.bert.pooler.parameters(), "lr": base_lr})
    params.append({"params": model.classifier.parameters(), "lr": base_lr * 10})

    return params

model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=3)
layer_params = get_layerwise_lr(model, base_lr=2e-5, lr_decay=0.75)
optimizer = AdamW(layer_params)
print("Layer-wise LR configurato correttamente")

# =========================================
# Strategia 4: Data Augmentation per NLP
# =========================================
import random

def easy_data_augmentation(text, num_aug=4, alpha_rs=0.1, alpha_ri=0.1, alpha_sr=0.1):
    """
    Easy Data Augmentation (EDA):
    - RS: Random Swap di parole
    - RI: Random Insertion di sinonimi
    - SR: Synonym Replacement
    """
    words = text.split()
    augmented = []

    for _ in range(num_aug):
        new_words = words.copy()

        # Random Swap
        if len(new_words) >= 2 and random.random() < alpha_rs:
            i, j = random.sample(range(len(new_words)), 2)
            new_words[i], new_words[j] = new_words[j], new_words[i]

        augmented.append(" ".join(new_words))

    return augmented

# Back-translation: traduzione IT->EN->IT per generare varianti semanticamente simili
# Richiede modelli di traduzione (es. Helsinki-NLP/opus-mt-it-en e opus-mt-en-it)
def back_translate(text: str, it_to_en, en_to_it) -> str:
    """Traduzione inversa per data augmentation."""
    en_text = it_to_en(text, max_length=512)[0]['translation_text']
    it_back = en_to_it(en_text, max_length=512)[0]['translation_text']
    return it_back

# Esempio di utilizzo (richiede pipeline di traduzione configurate)
from transformers import pipeline
# it_en = pipeline("translation_it_to_en", model="Helsinki-NLP/opus-mt-it-en")
# en_it = pipeline("translation_en_to_it", model="Helsinki-NLP/opus-mt-en-it")

# augmented_text = back_translate("Il paziente riferisce dolore toracico.", it_en, en_it)
print("Data augmentation configurata!")

import torch
from torch import nn
from typing import Dict, Iterator
import copy

# =========================================
# Elastic Weight Consolidation (EWC)
# =========================================
class EWC:
    """
    Elastic Weight Consolidation per prevenire catastrophic forgetting.
    Penalizza grandi cambiamenti ai parametri importanti per i task precedenti.

    Riferimento: Kirkpatrick et al. (2017) "Overcoming catastrophic forgetting in NNs"
    """
    def __init__(self, model: nn.Module, dataset: Iterator, lambda_ewc: float = 0.4):
        self.model = model
        self.lambda_ewc = lambda_ewc

        # Salva i pesi originali
        self._means: Dict[str, torch.Tensor] = {
            n: p.data.clone()
            for n, p in model.named_parameters()
            if p.requires_grad
        }

        # Calcola Fisher Information Matrix (diagonale)
        self._fisher = self._compute_fisher(dataset)

    def _compute_fisher(self, dataset: Iterator) -> Dict[str, torch.Tensor]:
        """
        Stima la Fisher Information Matrix diagonale come media dei gradienti al quadrato.
        Più alto il valore, più importante e quel parametro.
        """
        fisher = {n: torch.zeros_like(p) for n, p in self.model.named_parameters() if p.requires_grad}

        self.model.eval()
        n_samples = 0

        for batch in dataset:
            self.model.zero_grad()
            outputs = self.model(**batch)
            loss = outputs.loss
            loss.backward()

            for n, p in self.model.named_parameters():
                if p.grad is not None and n in fisher:
                    fisher[n] += p.grad.detach() ** 2
            n_samples += 1

        # Normalizza per il numero di batch
        for n in fisher:
            fisher[n] /= n_samples

        return fisher

    def penalty(self, model: nn.Module) -> torch.Tensor:
        """Calcola la penalita EWC da aggiungere alla task loss."""
        penalty = torch.tensor(0.0, device=next(model.parameters()).device)
        for n, p in model.named_parameters():
            if n in self._fisher and n in self._means:
                penalty += (self._fisher[n] * (p - self._means[n]) ** 2).sum()
        return self.lambda_ewc * penalty

# Uso nel training loop:
# ewc = EWC(model, old_task_dataloader, lambda_ewc=0.4)
# loss = task_loss + ewc.penalty(model)

# =========================================
# L2 Regularization verso i pesi originali (alternativa più semplice)
# =========================================
def l2_penalty_to_pretrained(model: nn.Module, original_params: dict, lambda_l2: float = 0.01) -> torch.Tensor:
    """
    Penalizza la distanza L2 dai pesi originali.
    Più semplice di EWC ma meno preciso (non tiene conto dell'importanza dei parametri).
    """
    penalty = torch.tensor(0.0)
    for n, p in model.named_parameters():
        if n in original_params:
            penalty += ((p - original_params[n]) ** 2).sum()
    return lambda_l2 * penalty

# =========================================
# Mixout: dropout dai pesi originali (alternativa moderna)
# =========================================
class MixoutLinear(nn.Module):
    """
    Mixout: durante il training, con probabilità p usa i pesi originali
    invece dei pesi aggiornati. Questo evita overfitting ai dati di fine-tuning
    e mantiene la conoscenza del pre-training.

    Riferimento: Lee et al. (2020) "Mixout: Effective Regularization to Finetune LLMs"
    """
    def __init__(self, linear_layer: nn.Linear, p: float = 0.9):
        super().__init__()
        self.original_weight = linear_layer.weight.data.clone()
        self.original_bias = linear_layer.bias.data.clone() if linear_layer.bias is not None else None
        self.linear = linear_layer
        self.p = p

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.training:
            # Mask casuale: usa pesi originali con probabilità p
            mask = torch.bernoulli(torch.full_like(self.linear.weight, self.p))
            weight = mask * self.original_weight + (1 - mask) * self.linear.weight
            bias = self.original_bias if self.original_bias is not None else self.linear.bias
            return nn.functional.linear(x, weight, bias)
        return self.linear(x)

print("EWC, L2 regularization e Mixout configurati!")

from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    roc_auc_score,
    precision_recall_curve,
    average_precision_score
)
import numpy as np
import pandas as pd
import torch

def comprehensive_evaluation(
    model,
    tokenizer,
    test_texts: list,
    test_labels: list,
    label_names: list,
    batch_size: int = 32,
    device: str = "cuda"
):
    """
    Valutazione completa: metriche aggregate, per classe, analisi degli errori,
    calibrazione e esempi incerti.
    """
    model.eval()
    all_logits, all_labels_list = [], []

    for i in range(0, len(test_texts), batch_size):
        batch_texts = test_texts[i:i+batch_size]
        batch_labels = test_labels[i:i+batch_size]

        inputs = tokenizer(
            batch_texts, return_tensors='pt',
            truncation=True, padding=True, max_length=256
        ).to(device)

        with torch.no_grad():
            outputs = model(**inputs)

        all_logits.append(outputs.logits.cpu().numpy())
        all_labels_list.extend(batch_labels)

    all_logits = np.vstack(all_logits)
    # Softmax per probabilità calibrate
    all_probs = np.exp(all_logits) / np.exp(all_logits).sum(axis=1, keepdims=True)
    all_preds = np.argmax(all_logits, axis=1)
    all_labels_arr = np.array(all_labels_list)

    # =========================================
    # 1. Report di classificazione per classe
    # =========================================
    print("=" * 60)
    print("CLASSIFICATION REPORT")
    print("=" * 60)
    print(classification_report(all_labels_arr, all_preds, target_names=label_names, digits=4))

    # =========================================
    # 2. Metriche di confidenza
    # =========================================
    max_probs = all_probs.max(axis=1)
    correct_mask = (all_preds == all_labels_arr)

    print("\n=== CONFIDENZA MODELLO ===")
    print(f"Confidenza media (corretti): {max_probs[correct_mask].mean():.4f}")
    print(f"Confidenza media (sbagliati): {max_probs[~correct_mask].mean():.4f}")

    # =========================================
    # 3. Esempi incerti (alta entropia)
    # =========================================
    entropies = -np.sum(all_probs * np.log(all_probs + 1e-10), axis=1)
    uncertain_threshold = np.percentile(entropies, 80)  # top 20% più incerti
    uncertain_mask = entropies > uncertain_threshold

    print(f"\n=== ESEMPI INCERTI ({uncertain_mask.sum()}/{len(all_labels_arr)}) ===")
    print(f"Accuracy sugli incerti: {correct_mask[uncertain_mask].mean():.4f}")
    print(f"Accuracy sui certi: {correct_mask[~uncertain_mask].mean():.4f}")

    # =========================================
    # 4. Analisi degli errori per classe
    # =========================================
    error_mask = ~correct_mask
    error_df = pd.DataFrame({
        "text": [test_texts[i] for i in range(len(test_texts))],
        "true_label": [label_names[l] for l in all_labels_arr],
        "pred_label": [label_names[p] for p in all_preds],
        "confidence": max_probs,
        "entropy": entropies,
        "correct": correct_mask
    })

    print("\n=== ESEMPI ERRATI CON ALTA CONFIDENZA ===")
    high_conf_errors = error_df[
        (~error_df["correct"]) &
        (error_df["confidence"] > 0.9)
    ].head(5)
    print(high_conf_errors[["text", "true_label", "pred_label", "confidence"]].to_string())

    return {
        "predictions": all_preds,
        "probabilities": all_probs,
        "errors": error_df[~error_df["correct"]]
    }

from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
from peft import PeftModel
import json
import os
from pathlib import Path
from datetime import datetime

class ModelDeploymentManager:
    """
    Gestisce il deployment di modelli fine-tuned con LoRA.
    Supporta: salvataggio versioni, merge, export ONNX, serving.
    """

    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def save_version(
        self,
        base_model_name: str,
        adapter_path: str,
        metadata: dict,
        merge: bool = True
    ) -> str:
        """Salva una versione del modello con metadata completi."""
        version = datetime.now().strftime("%Y%m%d_%H%M%S")
        version_dir = self.output_dir / f"v_{version}"
        version_dir.mkdir()

        # Carica modelli
        base_model = AutoModelForSequenceClassification.from_pretrained(base_model_name)
        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        peft_model = PeftModel.from_pretrained(base_model, adapter_path)

        # Salva adapter (leggero, ~1-5MB)
        adapter_dir = version_dir / "adapter"
        peft_model.save_pretrained(str(adapter_dir))
        tokenizer.save_pretrained(str(adapter_dir))

        if merge:
            # Merge e salva modello completo (per inference veloce)
            merged_dir = version_dir / "merged"
            merged_model = peft_model.merge_and_unload()
            merged_model.save_pretrained(str(merged_dir))
            tokenizer.save_pretrained(str(merged_dir))

        # Metadata del deployment
        deploy_metadata = {
            "version": version,
            "base_model": base_model_name,
            "created_at": datetime.now().isoformat(),
            "adapter_path": str(adapter_dir),
            "merged_path": str(merged_dir) if merge else None,
            "adapter_size_mb": sum(
                os.path.getsize(f) for f in adapter_dir.rglob("*") if f.is_file()
            ) / 1e6,
            **metadata
        }

        with open(version_dir / "metadata.json", "w") as f:
            json.dump(deploy_metadata, f, indent=2)

        print(f"Versione {version} salvata:")
        print(f"  Adapter: {deploy_metadata['adapter_size_mb']:.1f}MB")
        if merge:
            merged_size = sum(os.path.getsize(f) for f in merged_dir.rglob("*") if f.is_file()) / 1e6
            print(f"  Merged: {merged_size:.1f}MB")

        return str(version_dir)

    def load_for_inference(self, version_dir: str, use_merged: bool = True):
        """Carica il modello per inference in produzione."""
        version_path = Path(version_dir)
        with open(version_path / "metadata.json") as f:
            meta = json.load(f)

        if use_merged and meta.get("merged_path"):
            model = AutoModelForSequenceClassification.from_pretrained(meta["merged_path"])
            tokenizer = AutoTokenizer.from_pretrained(meta["merged_path"])
        else:
            base = AutoModelForSequenceClassification.from_pretrained(meta["base_model"])
            model = PeftModel.from_pretrained(base, meta["adapter_path"])
            tokenizer = AutoTokenizer.from_pretrained(meta["adapter_path"])

        # Crea pipeline production-ready
        clf_pipeline = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            device=0 if __import__("torch").cuda.is_available() else -1
        )
        return clf_pipeline, meta

# Esempio di utilizzo
manager = ModelDeploymentManager("./deployed_models")
# version_dir = manager.save_version(
#     base_model_name="dbmdz/bert-base-italian-cased",
#     adapter_path="./models/bert-medical-lora-adapter",
#     metadata={"eval_f1": 0.912, "eval_accuracy": 0.924, "domain": "medical-it"}
# )
print("Deployment manager configurato!")