Fix handling of attention_mask in encoders (#14)

Author: fcogidi

mmlearn/modules/encoders/clip_encoders.py

@@ -327,8 +327,8 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> Tuple[torch.Tensor
             The text embeddings. Will be a tuple with a single element.
         """
         input_ids = inputs[Modalities.TEXT]
-        attention_mask = inputs.get("attention_mask") or inputs.get(
-            Modalities.TEXT.attention_mask
+        attention_mask: Optional[torch.Tensor] = inputs.get(
+            "attention_mask", inputs.get(Modalities.TEXT.attention_mask, None)
         )
         position_ids = inputs.get("position_ids")
 
@@ -568,8 +568,9 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """
         output = self.model(
             input_ids=inputs[Modalities.TEXT],
-            attention_mask=inputs.get("attention_mask")
-            or inputs.get(Modalities.TEXT.attention_mask),
+            attention_mask=inputs.get(
+                "attention_mask", inputs.get(Modalities.TEXT.attention_mask, None)
+            ),
             inputs_embeds=inputs.get("inputs_embeds"),
             output_attentions=inputs.get("output_attentions"),
             output_hidden_states=True,

mmlearn/modules/encoders/hf_text_encoders.py

@@ -156,8 +156,9 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """
         outputs = self.model(
             input_ids=inputs[Modalities.TEXT],
-            attention_mask=inputs.get("attention_mask")
-            or inputs.get(Modalities.TEXT.attention_mask),
+            attention_mask=inputs.get(
+                "attention_mask", inputs.get(Modalities.TEXT.attention_mask, None)
+            ),
             position_ids=inputs.get("position_ids"),
             output_attentions=inputs.get("output_attentions"),
             return_dict=True,

projects/bioscan_clip/encoders.py

@@ -140,8 +140,9 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """Run the forward pass."""
         outputs = self.model(
             input_ids=inputs[Modalities.DNA],
-            attention_mask=inputs.get("attention_mask")
-            or inputs.get(Modalities.DNA.attention_mask),
+            attention_mask=inputs.get(
+                "attention_mask", inputs.get(Modalities.DNA.attention_mask, None)
+            ),
             position_ids=inputs.get("position_ids"),
             output_attentions=inputs.get("output_attentions"),
             return_dict=True,