Refactor from_pretrained method for robust model loading and improved error handling

unified_cell_classifier.py

@@ -221,78 +221,78 @@ class UnifiedCellClassifier(nn.Module):
         """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 path locale
-        
-        Args:
-            repo_id_or_path: ID del repository HF o path locale
+        Carica il modello da HuggingFace Hub o da un path locale in modo robusto.
         """
-        # 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")
+        # 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)
         
-        # 2. Istanzia il modello
+        # 3. Istanzia il "contenitore" del modello
         model = cls(**config)
         
-        # 3. Carica il classificatore principale
-        main_config = config['main_classifier_config']
-        model.main_classifier = model._create_classifier_from_config(main_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)
         
-        # Carica i pesi del main classifier
-        main_weights_path = os.path.join(repo_id_or_path, "main_classifier/main_classifier.safetensors")
+        # 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)
         
-        # Carica le label del main classifier
-        main_labels_path = get_file_path("main_classifier/id2label_main.json")
+        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)
-        
-        # 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['sub_classifiers_config'][sub_name]
-                model.sub_classifiers[sub_name] = model._create_classifier_from_config(sub_config)
-                
-                sub_weights_path = os.path.join(repo_id_or_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)
-                
-                # Carica le label del sub-classificatore
-                sub_labels_path = get_file_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)
+            
+        # 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)
                     
-            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")
+                    # 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)
-        except:
-            print("File macro_to_sub.json non trovato, uso mapping di default")
+        else:
+            print("File macro_to_sub.json non trovato, uso mapping di default.")
         
         model.eval()
         return model
+
     
     def save_pretrained(self, save_directory: str):
         """