Update unified_cell_classifier.py

unified_cell_classifier.py

@@ -1,163 +1,249 @@
-import torch
-import torch.nn as nn
-import json
-import os
-from typing import Dict, Optional, Tuple
-
-class UnifiedCellClassifier(nn.Module):
-    def __init__(self, models_base_path: str, sub_classifier_names: list):
-        """
-        Args:
-            models_base_path: Path base dove sono salvati i modelli
-            sub_classifier_names: Lista nomi cartelle sub-classificatori
-                Struttura attesa:
-                - main_classifier/model.pth + id2label.json
-                - B_cells_classifier/model.pth + id2label.json
-                - T_cells_classifier/model.pth + id2label.json
-                - ...
-        """
-        super().__init__()
-        
-        # Carica classificatore principale
-        main_path = os.path.join(models_base_path, "main_classifier")
-        self.main_classifier = torch.load(os.path.join(main_path, "complete_model.pth"), map_location='cpu', weights_only=False)
-        with open(os.path.join(main_path, "id2label.json")) as f:
-            self.main_labels = json.load(f)
-        
-        
-        # Carica sub-classificatori
-        self.sub_classifiers = nn.ModuleDict()
-        self.sub_labels = {}
-        
-        for sub_name in sub_classifier_names:
-            sub_path = os.path.join(models_base_path, sub_name)
-            
-            if os.path.exists(sub_path):
-                self.sub_classifiers[sub_name] = torch.load(
-                    os.path.join(sub_path, "complete_model.pth"), map_location='cpu', weights_only=False
-                )
-                with open(os.path.join(sub_path, "id2label.json")) as f:
-                    self.sub_labels[sub_name] = json.load(f)
-        
-        # Mapping macrocategoria -> sub-classificatore (da configurare)
-        self.macro_to_sub = self._build_macro_to_sub_mapping()
-        
-        # Imposta modalità eval
-        self.eval()
-    
-    def _build_macro_to_sub_mapping(self):
-
-        return {
-            "0": "B_cells_classifier",      
-            "1": "CD4plus_T_cells_classifier",     
-            "4": "Myeloid_cells_classifier",     
-            "5": "NK_cells_classifier",
-            "7": "TRAV1_2_CD8plus_T_cells",
-            "8": "gd_T_cells_classfier"
-        }
-    
-    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
-        """
-        # Classificazione principale
-        with torch.no_grad():
-            main_logits = self.main_classifier(x)
-            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_logits = self.sub_classifiers[sub_classifier_name](x[i:i+1])
-                        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
-    def from_pretrained(cls, repo_path: str, **kwargs):
-        """
-        carica la struttura dal repo di HF: aspetta
-          - config.json
-          - id2label_main.json
-          - macro_to_sub.json
-          - sub_classifiers/<name>.bin + id2label
-        """
-        # 1. leggi config
-        with open(os.path.join(repo_path, "config.json")) as f:
-            config = json.load(f)
-
-        # 2. istanzia l'oggetto
-        model = cls(**config)  
-
-        # 3. carica main
-        main_sd = torch.load(os.path.join(repo_path, "main_classifier.bin"), map_location="cpu")
-        model.main_classifier.load_state_dict(main_sd)
-        model.main_labels = json.load(open(os.path.join(repo_path, "id2label_main.json")))
-
-        # 4. carica sub
-        model.sub_classifiers = nn.ModuleDict()
-        model.sub_labels = {}
-        for name in model.sub_classifier_names:
-            bin_path = os.path.join(repo_path, "sub_classifiers", f"{name}.bin")
-            model.sub_classifiers[name] = model._build_submodule(name)  # metodo helper che crea l’istanza
-            model.sub_classifiers[name].load_state_dict(torch.load(bin_path, map_location="cpu"))
-            model.sub_labels[name] = json.load(open(
-                os.path.join(repo_path, "sub_classifiers", f"{name}_id2label.json")))
-        
-        model.macro_to_sub = json.load(open(os.path.join(repo_path, "macro_to_sub.json")))
-        model.eval()
-        return model
+import torch
+import torch.nn as nn
+import json
+import os
+from typing import Dict, Optional, Tuple
+from huggingface_hub import hf_hub_download
+
+class UnifiedCellClassifier(nn.Module):
+    def __init__(self, 
+                 main_classifier_config: Dict = None,
+                 sub_classifiers_config: Dict = None,
+                 sub_classifier_names: list = None,
+                 **kwargs):
+        """
+        Args:
+            main_classifier_config: Configurazione per il classificatore principale
+            sub_classifiers_config: Configurazioni per i sub-classificatori
+            sub_classifier_names: Lista nomi 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",
+            "8": "gd_T_cells_classfier"
+        }
+    
+    def _create_classifier_from_config(self, config: Dict):
+        """Crea un classificatore dalla configurazione"""
+        # Esempio di configurazione base - adatta secondo i tuoi modelli
+        input_dim = config.get('input_dim', 512)
+        hidden_dim = config.get('hidden_dim', 256)
+        num_classes = config.get('num_classes', 10)
+        dropout = config.get('dropout', 0.1)
+        
+        return nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_classes)
+        )
+    
+    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
+        with torch.no_grad():
+            main_logits = self.main_classifier(x)
+            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_logits = self.sub_classifiers[sub_classifier_name](x[i:i+1])
+                        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
+    def from_pretrained(cls, repo_id_or_path: str, **kwargs):
+        """
+        Carica il modello da HuggingFace Hub o da path locale
+        
+        Args:
+            repo_id_or_path: ID del repository HF o path locale
+        """
+        # Determina se è un path locale o repo HF
+        is_local = os.path.exists(repo_id_or_path)
+        
+        def get_file_path(filename):
+            if is_local:
+                return os.path.join(repo_id_or_path, filename)
+            else:
+                return hf_hub_download(repo_id=repo_id_or_path, filename=filename)
+        
+        # 1. Carica configurazione
+        config_path = get_file_path("config.json")
+        with open(config_path) as f:
+            config = json.load(f)
+        
+        # 2. Istanzia il modello
+        model = cls(**config)
+        
+        # 3. Carica il classificatore principale
+        # Crea l'architettura del main classifier
+        main_config = config.get('main_classifier_config', {})
+        model.main_classifier = model._create_classifier_from_config(main_config)
+        
+        # Carica i pesi del main classifier
+        main_weights_path = get_file_path("main_classifier.bin")
+        main_state_dict = torch.load(main_weights_path, map_location="cpu")
+        model.main_classifier.load_state_dict(main_state_dict)
+        
+        # Carica le label del main classifier
+        main_labels_path = get_file_path("id2label_main.json")
+        with open(main_labels_path) as f:
+            model.main_labels = json.load(f)
+        
+        # 4. Carica i sub-classificatori
+        model.sub_classifiers = nn.ModuleDict()
+        model.sub_labels = {}
+        
+        for sub_name in model.sub_classifier_names:
+            try:
+                # Crea l'architettura del sub-classificatore
+                sub_config = config.get('sub_classifiers_config', {}).get(sub_name, {})
+                model.sub_classifiers[sub_name] = model._create_classifier_from_config(sub_config)
+                
+                # Carica i pesi del sub-classificatore
+                sub_weights_path = get_file_path(f"sub_classifiers/{sub_name}.bin")
+                sub_state_dict = torch.load(sub_weights_path, map_location="cpu") 
+                model.sub_classifiers[sub_name].load_state_dict(sub_state_dict)
+                
+                # Carica le label del sub-classificatore
+                sub_labels_path = get_file_path(f"sub_classifiers/{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"Errore nel caricamento del sub-classificatore {sub_name}: {e}")
+                continue
+        
+        # 5. Carica il mapping macro_to_sub se esiste
+        try:
+            macro_to_sub_path = get_file_path("macro_to_sub.json")
+            with open(macro_to_sub_path) as f:
+                model.macro_to_sub = json.load(f)
+        except:
+            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)
+