Fix handling of attention_mask in encoders (#14)

VectorInstitute · Sep 9, 2024 · 1ad7199 · 1ad7199
1 parent 039b5d0
commit 1ad7199
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/mmlearn/modules/encoders/clip_encoders.py b/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,

diff --git a/mmlearn/modules/encoders/hf_text_encoders.py b/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,

diff --git a/projects/bioscan_clip/encoders.py b/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,