Add BIOSCAN-CLIP project (#8)

Author: fcogidi

mmlearn/datasets/core/modalities.py

@@ -38,6 +38,7 @@ class Modality(str):
 
     _default_properties = {
         "target": "{}_target",
+        "attention_mask": "{}_attention_mask",
         "mask": "{}_mask",
         "embedding": "{}_embedding",
         "masked_embedding": "{}_masked_embedding",

mmlearn/modules/encoders/clip_encoders.py

@@ -24,6 +24,7 @@
     group="modules/encoders",
     provider="mmlearn",
     model_name_or_path="openai/clip-vit-base-patch16",
+    hydra_convert="object",  # required for `peft_config` to be converted to a `PeftConfig` object
 )
 class HFCLIPTextEncoder(nn.Module):
     """Wrapper around the `CLIPTextModel` from HuggingFace.
@@ -103,7 +104,8 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """
         outputs = self.model(
             input_ids=inputs[Modalities.TEXT],
-            attention_mask=inputs.get("attention_mask"),
+            attention_mask=inputs.get("attention_mask")
+            or inputs.get(Modalities.TEXT.attention_mask),
             position_ids=inputs.get("position_ids"),
             output_attentions=inputs.get("output_attentions"),
             return_dict=True,
@@ -123,6 +125,7 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
     group="modules/encoders",
     provider="mmlearn",
     model_name_or_path="openai/clip-vit-base-patch16",
+    hydra_convert="object",
 )
 class HFCLIPVisionEncoder(nn.Module):
     """Wrapper around the `CLIPVisionModel` from HuggingFace.
@@ -247,6 +250,7 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
     group="modules/encoders",
     provider="mmlearn",
     model_name_or_path="openai/clip-vit-base-patch16",
+    hydra_convert="object",
 )
 class HFCLIPTextEncoderWithProjection(nn.Module):
     """Wrapper around the `CLIPTextModelWithProjection` from HuggingFace.
@@ -323,7 +327,9 @@ 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")
+        attention_mask = inputs.get("attention_mask") or inputs.get(
+            Modalities.TEXT.attention_mask
+        )
         position_ids = inputs.get("position_ids")
 
         if self.use_all_token_embeddings:
@@ -350,6 +356,7 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> Tuple[torch.Tensor
     group="modules/encoders",
     provider="mmlearn",
     model_name_or_path="openai/clip-vit-base-patch16",
+    hydra_convert="object",
 )
 class HFCLIPVisionEncoderWithProjection(nn.Module):
     """Wrapper around the `CLIPVisionModelWithProjection` class from HuggingFace.
@@ -463,7 +470,7 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> Tuple[torch.Tensor
         return (self.model.visual_projection(pooled_output),)
 
 
-@store(group="modules/encoders", provider="mmlearn")
+@store(group="modules/encoders", provider="mmlearn", hydra_convert="object")
 class PubMedBERTForCLIPTextEncoding(nn.Module):
     """BiomedNLP's PubMedBERT model for CLIP text encoding.
 
@@ -561,7 +568,8 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """
         output = self.model(
             input_ids=inputs[Modalities.TEXT],
-            attention_mask=inputs.get("attention_mask"),
+            attention_mask=inputs.get("attention_mask")
+            or inputs.get(Modalities.TEXT.attention_mask),
             inputs_embeds=inputs.get("inputs_embeds"),
             output_attentions=inputs.get("output_attentions"),
             output_hidden_states=True,

mmlearn/modules/encoders/hf_text_encoders.py

@@ -17,7 +17,7 @@
     from peft import PeftConfig
 
 
-@store(group="modules/encoders", provider="mmlearn")
+@store(group="modules/encoders", provider="mmlearn", hydra_convert="object")
 class HFTextEncoder(nn.Module):
     """Wrapper around huggingface models in the `AutoModelForTextEncoding` class.
 
@@ -66,7 +66,6 @@ def __init__(  # noqa: PLR0912
         super().__init__()
         if model_config_kwargs is None:
             model_config_kwargs = {}
-        model_config_kwargs["use_return_dict"] = True
         model_config_kwargs["output_hidden_states"] = True
         model_config_kwargs["add_pooling_layer"] = False
         model = hf_utils.load_huggingface_model(
@@ -157,7 +156,8 @@ def forward(self, inputs: Dict[Union[str, Modality], Any]) -> BaseModelOutput:
         """
         outputs = self.model(
             input_ids=inputs[Modalities.TEXT],
-            attention_mask=inputs.get("attention_mask"),
+            attention_mask=inputs.get("attention_mask")
+            or inputs.get(Modalities.TEXT.attention_mask),
             position_ids=inputs.get("position_ids"),
             output_attentions=inputs.get("output_attentions"),
             return_dict=True,