Update unified_cell_classifier.py

unified_cell_classifier.py

@@ -1,21 +1,110 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import json
 import os
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, List
 from huggingface_hub import hf_hub_download
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+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,
-                 sub_classifier_names: list = 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
-            sub_classifier_names: Lista nomi sub-classificatori
         """
         super().__init__()
         
@@ -40,23 +129,19 @@ class UnifiedCellClassifier(nn.Module):
             "1": "CD4plus_T_cells_classifier",     
             "4": "Myeloid_cells_classifier",     
             "5": "NK_cells_classifier",
-            "7": "TRAV1_2_CD8plus_T_cells",
+            "7": "TRAV1_2_CD8plus_T_cells_classifier",
             "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)
+        """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):
@@ -73,9 +158,10 @@ class UnifiedCellClassifier(nn.Module):
         if self.main_classifier is None:
             raise RuntimeError("Modello non caricato. Usa from_pretrained() per caricare il modello.")
         
-        # Classificazione principale
+        # Classificazione principale - usa il metodo del classificatore per avere solo logits
         with torch.no_grad():
-            main_logits = self.main_classifier(x)
+            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)
         
@@ -94,7 +180,8 @@ class UnifiedCellClassifier(nn.Module):
                 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_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)
                     
@@ -159,8 +246,7 @@ class UnifiedCellClassifier(nn.Module):
         model = cls(**config)
         
         # 3. Carica il classificatore principale
-        # Crea l'architettura del main classifier
-        main_config = config.get('main_classifier_config', {})
+        main_config = config['main_classifier_config']
         model.main_classifier = model._create_classifier_from_config(main_config)
         
         # Carica i pesi del main classifier
@@ -180,7 +266,7 @@ class UnifiedCellClassifier(nn.Module):
         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, {})
+                sub_config = config['sub_classifiers_config'][sub_name]
                 model.sub_classifiers[sub_name] = model._create_classifier_from_config(sub_config)
                 
                 # Carica i pesi del sub-classificatore
@@ -246,4 +332,3 @@ class UnifiedCellClassifier(nn.Module):
         # 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)
-