Upload 2 files

app.py

@@ -97,6 +97,7 @@ class DrugTargetInteractionApp:
             ).to(self.device)
 
             # Make prediction
+            self.model.INTERPR_DISABLE_MODE()
             with torch.no_grad():
                 prediction = self.model(target_inputs, drug_inputs)
                 
@@ -143,6 +144,7 @@ class DrugTargetInteractionApp:
                 max_length=512,
                 return_tensors="pt"
             ).to(self.device)
+            self.model.INTERPR_ENABLE_MODE()
 
             # Make prediction and extract visualization data
             with torch.no_grad():
@@ -163,6 +165,7 @@ class DrugTargetInteractionApp:
                 b = self.model.model.b
                 scaler = self.model.model.scaler
                 
+            logger.info(target_inputs, drug_inputs)
             # Generate visualizations
             try:
                 # 1. Cross-attention heatmap

modeling_dlmberta.py

@@ -52,6 +52,8 @@ class InteractionModelATTNForRegression(PreTrainedModel):
     def INTERPR_ENABLE_MODE(self):
         self.model.INTERPR_ENABLE_MODE()
 
+    def INTERPR_DISABLE_MODE(self):
+        self.model.INTERPR_DISABLE_MODE()
 
     def INTERPR_OVERRIDE_ATTN(self, new_weights):
         self.model.INTERPR_OVERRIDE_ATTN(new_weights)
@@ -295,6 +297,13 @@ class InteractionModelATTN(nn.Module):
             raise RuntimeError("Cannot enable interpretability mode while the model is training.")
         self.INTERPR_MODE = True
 
+    def INTERPR_DISABLE_MODE(self):
+        """
+        Disables the interpretability mode for the model.
+        """
+        if self.training:
+            raise RuntimeError("Cannot disable interpretability mode while the model is training.")
+        self.INTERPR_MODE = False
 
     def INTERPR_OVERRIDE_ATTN(self, new_weights):
         self.replace_weights = new_weights