Upload model

config.json

@@ -32,11 +32,19 @@
     "vocab_size": 30000
   },
   "encoder": {
-    "_name_or_path": "",
-    "architectures": null,
-    "model_type": "vit",
+    "_name_or_path": "aehrc/uniformer_base_tl_384",
+    "architectures": [
+      "UniFormerModel"
+    ],
+    "auto_map": {
+      "AutoConfig": "aehrc/uniformer_base_tl_384--configuration_uniformer.UniFormerWithProjectionHeadConfig",
+      "AutoModel": "aehrc/uniformer_base_tl_384--modelling_uniformer.UniFormerModel"
+    },
+    "init_value": 1e-06,
+    "layer_scale": false,
+    "model_type": "uniformer",
     "projection_size": 768,
-    "torch_dtype": null
+    "torch_dtype": "float32"
   },
   "is_encoder_decoder": false,
   "model_type": "cxrmate-ed",

configuration_cxrmate_ed.py

@@ -2,8 +2,6 @@ import transformers
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
-from .configuration_uniformer import UniFormerWithProjectionHeadConfig
-
 logger = logging.get_logger(__name__)
 
 
@@ -40,9 +38,10 @@ class CXRMateEDConfig(PretrainedConfig):
 
 
         if 'encoder' not in kwargs:
-            self.encoder = UniFormerWithProjectionHeadConfig.from_pretrained(
+            self.encoder = transformers.AutoConfig.from_pretrained(
                 'aehrc/uniformer_base_tl_384',
                 projection_size=768,
+                trust_remote_code=True,
             )   
         else:
             self.encoder = kwargs.pop("encoder")

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3479bf42a0f90e144362b0785b0fe9a11078562f61217230e4340a5519e56f48
-size 789958760
+oid sha256:38661d70c87174cf130d3b60c278cd5a07491742aed00eac34b2bca5d795d564
+size 789964216

modelling_cxrmate_ed.py

@@ -12,13 +12,12 @@ from torch.utils.data import Subset
 from torchvision.io import decode_image
 from transformers import PreTrainedTokenizerFast, VisionEncoderDecoderModel
 from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_outputs import Seq2SeqLMOutput
+from transformers.modeling_outputs import ModelOutput, Seq2SeqLMOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 
 from .configuration_cxrmate_ed import CXRMateEDConfig
 from .dataset import PriorsDataset
-from .modelling_uniformer import MultiUniFormerWithProjectionHead
 from .prepare_dataset import prepare_dataset
 from .utils import compute_time_delta
 
@@ -46,6 +45,70 @@ class FNNEncoder(torch.nn.Module):
         return self.down_proj(self.act_fn(self.up_proj(x)))
 
 
+class ProjectionHead(torch.nn.Module):
+
+    def __init__(self, input_size, hidden_size) -> None:
+        super().__init__()
+
+        # Layer normalisation before projection:
+        self.layer_norm = torch.nn.LayerNorm(input_size, eps=1e-6)
+
+        # No bias as following layer normalisation with bias:
+        self.projection = torch.nn.Linear(input_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.layer_norm(x)
+        x = self.projection(x)
+        return x
+
+
+class CXRStudyImagesEncoder(torch.nn.Module):
+    def __init__(self, encoder, decoder_config):
+        super().__init__()
+
+        self.encoder = encoder
+        self.config = encoder.config
+        self.adapter = ProjectionHead(self.config.embed_dim[-1], decoder_config.hidden_size)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, ModelOutput]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Flatten the batch and study_id dimensions:
+        assert len(pixel_values.shape) == 5, 'pixel_values must be B, S, C, H, W, where S is the max number of images for a study in the batch.'
+        last_hidden_state = self.encoder(pixel_values.view(-1, *pixel_values.shape[2:])).last_hidden_state
+
+        # Flatten h x w:
+        last_hidden_state = torch.flatten(last_hidden_state, 2) if last_hidden_state.dim() > 3 else last_hidden_state
+        
+        # Project the features for each spatial position to the decoder's hidden size using the adapter network:
+        last_hidden_state = self.adapter(last_hidden_state)
+        
+        # Concatenate the features for each chest X-ray:
+        last_hidden_state = last_hidden_state.view(pixel_values.shape[0], -1, last_hidden_state.shape[-1])
+
+        # Derive the attention mask from the pixel values:
+        mask = (pixel_values[:, :, 0, 0, 0] != 0.0)[:, :, None]
+        attention_mask = torch.ones(
+            [last_hidden_state.shape[0], pixel_values.shape[1], last_hidden_state.shape[1] // pixel_values.shape[1]], 
+            dtype=torch.long,
+            device=mask.device,
+        )
+        attention_mask = attention_mask * mask
+        attention_mask = attention_mask.view(attention_mask.shape[0], -1)
+
+        if not return_dict:
+            return last_hidden_state
+
+        return ModelOutput(last_hidden_state=last_hidden_state, attention_mask=attention_mask)
+
+
+
 class CXRMateEDModel(VisionEncoderDecoderModel):
 
     config_class = CXRMateEDConfig
@@ -77,14 +140,18 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
 
         # Encoder:
         if encoder is None:
-            encoder = MultiUniFormerWithProjectionHead(config=config.encoder)
+            encoder = transformers.AutoModel.from_pretrained(
+                'aehrc/uniformer_base_tl_384', 
+                config=config.encoder,
+                trust_remote_code=True,
+            )
 
         # Decoder:
         if decoder is None:
             assert not config.decoder.add_cross_attention
             decoder = transformers.LlamaForCausalLM(config=config.decoder)
 
-        self.encoder = encoder
+        self.encoder = CXRStudyImagesEncoder(encoder, self.config.decoder)
         self.decoder = decoder
 
         if self.encoder.config.to_dict() != self.config.encoder.to_dict():