unified_cell_classifier.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import json
import os
from typing import Dict, Optional, Tuple, List
from huggingface_hub import hf_hub_download
from transformers.modeling_outputs import SequenceClassifierOutput
from safetensors.torch import load_file
from huggingface_hub import snapshot_download

class MLPBlock(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, dropout_rate: float = 0.2, use_residual: bool = False):
        super().__init__()
        self.use_residual = use_residual and (input_dim == output_dim)
        self.linear = nn.Linear(input_dim, output_dim)
        self.bn = nn.BatchNorm1d(output_dim)
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout_rate)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        x = self.linear(x)
        x = self.bn(x)
        x = self.activation(x)
        x = self.dropout(x)
        if self.use_residual:
            x = x + identity
        return x

class AdvancedMLPClassifier(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dims: List[int],
        output_dim: int,
        dropout_rate: float = 0.2,
        use_residual_in_hidden: bool = True,
        loss_fn: Optional[nn.Module] = None
    ):
        super().__init__()
        self.initial_bn = nn.BatchNorm1d(input_dim)
        
        all_dims = [input_dim] + hidden_dims
        mlp_layers = []
        for i in range(len(all_dims) - 1):
            mlp_layers.append(
                MLPBlock(
                    input_dim=all_dims[i],
                    output_dim=all_dims[i + 1],
                    dropout_rate=dropout_rate,
                    use_residual=use_residual_in_hidden and (all_dims[i] == all_dims[i + 1])
                )
            )
        
        self.hidden_network = nn.Sequential(*mlp_layers)
        self.output_projection = nn.Linear(all_dims[-1], output_dim)
        self.loss_fn = loss_fn if loss_fn is not None else nn.CrossEntropyLoss()
        self._initialize_weights()
    
    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = True,
        **kwargs
    ) -> SequenceClassifierOutput:
        if input_ids.ndim > 2:
            input_ids = input_ids.view(input_ids.size(0), -1)  # Flatten if necessary
        
        x = self.initial_bn(input_ids)
        x = self.hidden_network(x)
        logits = self.output_projection(x)
        
        loss = self.loss_fn(logits, labels) if labels is not None else None
        
        if not return_dict:
            return (logits, loss) if loss is not None else (logits,)
        
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=None,
            attentions=None
        )
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

class UnifiedCellClassifier(nn.Module):
    def __init__(self, 
                 sub_classifier_names: list = None,
                 main_classifier_config: Dict = None,
                 sub_classifiers_config: Dict = None,
                 **kwargs):
        """
        Args:
            sub_classifier_names: Lista nomi sub-classificatori
            main_classifier_config: Configurazione per il classificatore principale
            sub_classifiers_config: Configurazioni per i sub-classificatori
        """
        super().__init__()
        
        # Salva configurazione
        self.sub_classifier_names = sub_classifier_names or []
        self.main_classifier_config = main_classifier_config or {}
        self.sub_classifiers_config = sub_classifiers_config or {}
        
        # Inizializza placeholder (verranno caricati in from_pretrained)
        self.main_classifier = None
        self.sub_classifiers = nn.ModuleDict()
        self.main_labels = {}
        self.sub_labels = {}
        
        # Mapping macrocategoria -> sub-classificatore
        self.macro_to_sub = self._build_default_macro_to_sub_mapping()
    
    def _build_default_macro_to_sub_mapping(self):
        """Mapping di default - può essere sovrascritto dal file macro_to_sub.json"""
        return {
            "0": "B_cells_classifier",      
            "1": "CD4plus_T_cells_classifier",     
            "4": "Myeloid_cells_classifier",     
            "5": "NK_cells_classifier",
            "7": "TRAV1_2_CD8plus_T_cells_classifier",
            "8": "gd_T_cells_classifier"
        }
    
    def _create_classifier_from_config(self, config: Dict):
        """Crea un AdvancedMLPClassifier dalla configurazione"""
        return AdvancedMLPClassifier(
            input_dim=config['input_dim'],
            hidden_dims=config['hidden_dims'],
            output_dim=config['output_dim'],
            dropout_rate=config.get('dropout_rate', 0.2),
            use_residual_in_hidden=config.get('use_residual_in_hidden', True),
            loss_fn=nn.CrossEntropyLoss()
        )
    
    def forward(self, x: torch.Tensor, return_probabilities: bool = False):
        """
        Forward pass gerarchico
        
        Args:
            x: Input embeddings [batch_size, embedding_dim]
            return_probabilities: Se True, restituisce anche le probabilità
            
        Returns:
            Dict con macro_prediction, sub_prediction, final_prediction
        """
        if self.main_classifier is None:
            raise RuntimeError("Modello non caricato. Usa from_pretrained() per caricare il modello.")
        
        # Classificazione principale - usa il metodo del classificatore per avere solo logits
        with torch.no_grad():
            main_output = self.main_classifier(x, return_dict=True)
            main_logits = main_output.logits
            main_probs = torch.softmax(main_logits, dim=-1)
            main_pred = torch.argmax(main_logits, dim=-1)
        
        # Classificazione secondaria
        batch_size = x.shape[0]
        sub_predictions = []
        sub_probabilities = [] if return_probabilities else None
        
        for i in range(batch_size):
            macro_idx = str(main_pred[i].item())
            macro_label = self.main_labels.get(macro_idx, f"unknown_{macro_idx}")
            
            # Controlla se esiste sub-classificatore per questa macro
            if macro_idx in self.macro_to_sub:
                sub_classifier_name = self.macro_to_sub[macro_idx]
                if sub_classifier_name in self.sub_classifiers:
                    # Usa sub-classificatore
                    with torch.no_grad():
                        sub_output = self.sub_classifiers[sub_classifier_name](x[i:i+1], return_dict=True)
                        sub_logits = sub_output.logits
                        sub_probs = torch.softmax(sub_logits, dim=-1)
                        sub_pred = torch.argmax(sub_logits, dim=-1)
                    
                    sub_idx = str(sub_pred.item())
                    sub_label = self.sub_labels[sub_classifier_name].get(sub_idx, f"unknown_{sub_idx}")
                    final_prediction = f"{macro_label}_{sub_label}"
                    
                    if return_probabilities:
                        sub_probabilities.append(sub_probs[0])
                else:
                    # Sub-classificatore non trovato, usa solo macro
                    final_prediction = macro_label
                    if return_probabilities:
                        sub_probabilities.append(None)
            else:
                # Nessun sub-classificatore per questa macro, usa solo macro
                final_prediction = macro_label
                if return_probabilities:
                    sub_probabilities.append(None)
            
            sub_predictions.append(final_prediction)
        
        result = {
            'macro_predictions': [self.main_labels.get(str(idx.item()), f"unknown_{idx.item()}") 
                                for idx in main_pred],
            'final_predictions': sub_predictions
        }
        
        if return_probabilities:
            result['macro_probabilities'] = main_probs
            result['sub_probabilities'] = sub_probabilities
        
        return result
    
    def predict(self, x: torch.Tensor):
        """Metodo semplificato per predizione"""
        return self.forward(x, return_probabilities=False)['final_predictions']
    
    @classmethod
    @classmethod
    def from_pretrained(cls, repo_id_or_path: str, **kwargs):
        """
        Carica il modello da HuggingFace Hub o da un path locale in modo robusto.
        """
        
        # 1. Ottieni un path locale unificato
        if os.path.isdir(repo_id_or_path):
            local_model_path = repo_id_or_path
        else:
            local_model_path = snapshot_download(repo_id=repo_id_or_path)

        # 2. Carica la configurazione generale
        config_path = os.path.join(local_model_path, "config.json")
        with open(config_path) as f:
            config = json.load(f)
        
        # 3. Istanzia il "contenitore" del modello
        model = cls(**config)
        
        # 4. Carica il classificatore principale
        # --- PASSAGGIO MANCANTE INSERITO QUI ---
        # a) Crea l'architettura del classificatore principale
        if 'main_classifier_config' in config:
            main_config = config['main_classifier_config']
            # Assumo che tu abbia un metodo per creare un classificatore dalla sua config
            model.main_classifier = model._create_classifier_from_config(main_config)
        
        # b) Ora carica i pesi, perché model.main_classifier esiste
        main_weights_path = os.path.join(local_model_path, "main_classifier/main_classifier.safetensors")
        main_state_dict = load_file(main_weights_path)
        model.main_classifier.load_state_dict(main_state_dict, strict=False)
        
        main_labels_path = os.path.join(local_model_path, "main_classifier/id2label_main.json")
        with open(main_labels_path) as f:
            model.main_labels = json.load(f)
            
        # 5. Carica i sub-classificatori
        if 'sub_classifiers_config' in config:
            for sub_name in model.sub_classifier_names:
                try:
                    # --- PASSAGGIO MANCANTE INSERITO QUI ---
                    # a) Crea l'architettura del sub-classificatore
                    sub_config = config['sub_classifiers_config'][sub_name]
                    model.sub_classifiers[sub_name] = model._create_classifier_from_config(sub_config)

                    # b) Ora carica i suoi pesi
                    sub_weights_path = os.path.join(local_model_path, f"sub_classifiers/{sub_name}/{sub_name}.safetensors")
                    sub_state_dict = load_file(sub_weights_path)
                    model.sub_classifiers[sub_name].load_state_dict(sub_state_dict, strict=False)
                    
                    # c) Carica le sue label
                    sub_labels_path = os.path.join(local_model_path, f"sub_classifiers/{sub_name}/{sub_name}_id2label.json")
                    with open(sub_labels_path) as f:
                        model.sub_labels[sub_name] = json.load(f)
                        
                except Exception as e:
                    print(f"⚠️ Avviso: impossibile caricare il sub-classificatore {sub_name}. Errore: {e}")
                    continue
        
        # 6. Carica il mapping macro_to_sub se esiste
        macro_to_sub_path = os.path.join(local_model_path, "macro_to_sub.json")
        if os.path.exists(macro_to_sub_path):
            with open(macro_to_sub_path) as f:
                model.macro_to_sub = json.load(f)
        else:
            print("File macro_to_sub.json non trovato, uso mapping di default.")
        
        model.eval()
        return model

    
    def save_pretrained(self, save_directory: str):
        """
        Salva il modello in formato HuggingFace
        """
        os.makedirs(save_directory, exist_ok=True)
        
        # Salva configurazione
        config = {
            'sub_classifier_names': self.sub_classifier_names,
            'main_classifier_config': self.main_classifier_config,
            'sub_classifiers_config': self.sub_classifiers_config
        }
        
        with open(os.path.join(save_directory, "config.json"), 'w') as f:
            json.dump(config, f, indent=2)
        
        # Salva main classifier
        if self.main_classifier is not None:
            torch.save(self.main_classifier.state_dict(), 
                      os.path.join(save_directory, "main_classifier.bin"))
            
            with open(os.path.join(save_directory, "id2label_main.json"), 'w') as f:
                json.dump(self.main_labels, f, indent=2)
        
        # Salva sub-classifiers
        sub_classifiers_dir = os.path.join(save_directory, "sub_classifiers")
        os.makedirs(sub_classifiers_dir, exist_ok=True)
        
        for name, classifier in self.sub_classifiers.items():
            torch.save(classifier.state_dict(), 
                      os.path.join(sub_classifiers_dir, f"{name}.bin"))
            
            with open(os.path.join(sub_classifiers_dir, f"{name}_id2label.json"), 'w') as f:
                json.dump(self.sub_labels[name], f, indent=2)
        
        # Salva mapping
        with open(os.path.join(save_directory, "macro_to_sub.json"), 'w') as f:
            json.dump(self.macro_to_sub, f, indent=2)