Update explainer.py

explainer.py

@@ -22,6 +22,7 @@ class SigLIPBBoxRegressor(nn.Module):
         text_dim = self.siglip.text_model.config.hidden_size
         if giant: text_dim = 1536
 
+         # Feature fusion layers
         self.vision_projector = nn.Sequential(
             nn.Linear(vision_dim, hidden_dim),
             nn.ReLU(),
@@ -32,6 +33,8 @@ class SigLIPBBoxRegressor(nn.Module):
             nn.ReLU(),
             nn.Dropout(0.1)
         )
+        
+        # Cross-modal fusion
         self.fusion_layer = nn.Sequential(
             nn.Linear(hidden_dim*2, hidden_dim),
             nn.ReLU(),
@@ -46,7 +49,7 @@ class SigLIPBBoxRegressor(nn.Module):
             nn.Dropout(0.1),
             nn.Linear(256, 128),
             nn.ReLU(),
-            nn.Linear(128, 2),
+            nn.Linear(128, 2), # (x1, y1)
         )
         self.bottomright_regressor = nn.Sequential(
             nn.Linear(hidden_dim//2, 256),
@@ -54,22 +57,30 @@ class SigLIPBBoxRegressor(nn.Module):
             nn.Dropout(0.1),
             nn.Linear(256, 128),
             nn.ReLU(),
-            nn.Linear(128, 2),
+            nn.Linear(128, 2), # (x2, y2)
         )
 
     def forward(self, pixel_values, input_ids):
         with torch.no_grad():
             outputs = self.siglip(pixel_values=pixel_values, input_ids=input_ids, return_dict=True)
+            
+        # Extract pooled features
         vision_features = outputs.image_embeds.float()
         text_features = outputs.text_embeds.float()
+        
+        # Project features
 
         vision_proj = self.vision_projector(vision_features)
         text_proj = self.text_projector(text_features)
+        
+        # Fuse modalities
         fused = torch.cat([vision_proj, text_proj], dim=1)
         fused_features = self.fusion_layer(fused)
-
+        
+        # Predict bbox
         topleft_pred = self.topleft_regressor(fused_features)
         bottomright_pred = self.bottomright_regressor(fused_features)
+        
         return torch.cat([topleft_pred, bottomright_pred], dim=1)
 
 class Explainer(PreTrainedModel):
@@ -104,7 +115,6 @@ class Explainer(PreTrainedModel):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-    # Load config automatically (HF passes `config` here sometimes)
         config = kwargs.pop("config", None)
         if config is None:
             config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path)