Upload model

config.json

@@ -1,4 +1,5 @@
 {
+  "add_time_deltas": true,
   "architectures": [
     "CXRMateEDModel"
   ],
@@ -6,49 +7,84 @@
     "AutoConfig": "configuration_cxrmate_ed.CXRMateEDConfig",
     "AutoModel": "modelling_cxrmate_ed.CXRMateEDModel"
   },
-  "decoder": {
-    "add_time_deltas": true,
+  "hidden_size": 768,
+  "history": 0,
+  "ignore_index": -100,
+  "image_seq_length": 576,
+  "image_token_index": 32000,
+  "include_time_delta": true,
+  "index_value_encoder_intermediate_size": 2048,
+  "model_type": "cxrmate-ed",
+  "pad_token_id": 4,
+  "projector_hidden_act": "gelu",
+  "prompt_report_sections_filter": [
+    "indication",
+    "history"
+  ],
+  "tables_filter": [
+    "mimic_cxr_sectioned",
+    "triage",
+    "medrecon"
+  ],
+  "text_config": {
+    "head_dim": 64,
     "hidden_size": 768,
-    "history": 0,
-    "include_time_delta": true,
-    "index_value_encoder_intermediate_size": 2048,
     "intermediate_size": 3072,
-    "is_decoder": true,
     "model_type": "llama",
     "num_attention_heads": 12,
     "num_hidden_layers": 6,
     "num_key_value_heads": 12,
-    "pad_token_id": 4,
-    "prompt_report_sections_filter": [
-      "indication",
-      "history"
-    ],
-    "tables_filter": [
-      "mimic_cxr_sectioned",
-      "triage",
-      "medrecon"
-    ],
-    "time_delta_monotonic_inversion": true,
     "vocab_size": 30000
   },
-  "encoder": {
+  "time_delta_monotonic_inversion": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.0",
+  "vision_config": {
     "_name_or_path": "aehrc/uniformer_base_tl_384",
     "architectures": [
       "UniFormerModel"
     ],
+    "attn_drop_rate": 0.0,
     "auto_map": {
       "AutoConfig": "aehrc/uniformer_base_tl_384--configuration_uniformer.UniFormerWithProjectionHeadConfig",
       "AutoModel": "aehrc/uniformer_base_tl_384--modelling_uniformer.UniFormerModel"
     },
+    "conv_stem": false,
+    "depth": [
+      5,
+      8,
+      20,
+      7
+    ],
+    "drop_path_rate": 0.3,
+    "drop_rate": 0.0,
+    "embed_dim": [
+      64,
+      128,
+      320,
+      512
+    ],
+    "head_dim": 64,
+    "image_size": 384,
+    "in_chans": 3,
     "init_value": 1e-06,
+    "layer_norm_eps": 1e-06,
     "layer_scale": false,
+    "mlp_ratio": 4,
     "model_type": "uniformer",
-    "projection_size": 768,
+    "num_classes": 1000,
+    "patch_size": [
+      4,
+      2,
+      2,
+      2
+    ],
+    "projection_size": null,
+    "qk_scale": null,
+    "qkv_bias": true,
+    "representation_size": null,
     "torch_dtype": "float32"
   },
-  "is_encoder_decoder": false,
-  "model_type": "cxrmate-ed",
-  "tie_word_embeddings": false,
-  "torch_dtype": "float32",
-  "transformers_version": "4.39.3"
+  "vision_feature_layer": -2,
+  "vision_feature_select_strategy": "default"
 }

configuration_cxrmate_ed.py

@@ -1,61 +1,94 @@
-import transformers
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
+from typing import Any
 
-logger = logging.get_logger(__name__)
+from transformers import LlavaConfig
 
 
-class CXRMateEDConfig(PretrainedConfig):
+class CXRMateEDConfig(LlavaConfig):
 
-    model_type = "cxrmate-ed"
+    model_type = 'cxrmate-ed'
 
-    def __init__(self, **kwargs):
+    def __init__(
+        self,
+        index_value_encoder_intermediate_size: int = 2048,
+        include_time_delta: bool = True,
+        time_delta_monotonic_inversion: bool = True,
+        add_time_deltas: bool = True,
+        history: int = 0,
+        tables_filter: list = ['mimic_cxr_sectioned', 'triage', 'medrecon'],
+        prompt_report_sections_filter: list = ['indication', 'history'],
+        pad_token_id: int = 4,
+        **kwargs: Any,
+    ) -> None:
         super().__init__(**kwargs)
+        self.index_value_encoder_intermediate_size = index_value_encoder_intermediate_size
+        self.include_time_delta = include_time_delta
+        self.time_delta_monotonic_inversion = time_delta_monotonic_inversion
+        self.add_time_deltas = add_time_deltas
+        self.history = history
+        self.tables_filter = tables_filter
+        self.prompt_report_sections_filter = prompt_report_sections_filter
+        self.pad_token_id = pad_token_id
+        
+        self.hidden_size = self.text_config.hidden_size
+
+# import transformers
+# from transformers.configuration_utils import PretrainedConfig
+# from transformers.utils import logging
+
+# logger = logging.get_logger(__name__)
+
+
+# class CXRMateEDConfig(PretrainedConfig):
+
+#     model_type = "cxrmate-ed"
+
+#     def __init__(self, **kwargs):
+#         super().__init__(**kwargs)
+
+#         if 'decoder' not in kwargs:
+
+#             self.decoder = transformers.LlamaConfig(
+#                 vocab_size=30000,
+#                 hidden_size=768,
+#                 intermediate_size=3072,
+#                 num_attention_heads=12,
+#                 num_hidden_layers=6,
+#                 max_position_embeddings=2048,
+#             )
+#             self.decoder.is_decoder = True
+
+#             self.decoder.index_value_encoder_intermediate_size = 2048
+#             self.decoder.include_time_delta = True
+#             self.decoder.time_delta_monotonic_inversion = True
+#             self.decoder.add_time_deltas = True
+#             self.decoder.history = 0
+#             self.decoder.tables_filter = ["mimic_cxr_sectioned", "triage", "medrecon"]
+#             self.decoder.prompt_report_sections_filter = ["indication", "history"]
+#             self.decoder.pad_token_id = 4
+
+#         else:
+#             self.decoder = kwargs.pop("decoder")
+
 
-        if 'decoder' not in kwargs:
-
-            self.decoder = transformers.LlamaConfig(
-                vocab_size=30000,
-                hidden_size=768,
-                intermediate_size=3072,
-                num_attention_heads=12,
-                num_hidden_layers=6,
-                max_position_embeddings=2048,
-            )
-            self.decoder.is_decoder = True
-
-            self.decoder.index_value_encoder_intermediate_size = 2048
-            self.decoder.include_time_delta = True
-            self.decoder.time_delta_monotonic_inversion = True
-            self.decoder.add_time_deltas = True
-            self.decoder.history = 0
-            self.decoder.tables_filter = ["mimic_cxr_sectioned", "triage", "medrecon"]
-            self.decoder.prompt_report_sections_filter = ["indication", "history"]
-            self.decoder.pad_token_id = 4
-
-        else:
-            self.decoder = kwargs.pop("decoder")
-
-
-        if 'encoder' not in kwargs:
-            self.encoder = transformers.AutoConfig.from_pretrained(
-                'aehrc/uniformer_base_tl_384',
-                projection_size=768,
-                trust_remote_code=True,
-            )   
-        else:
-            self.encoder = kwargs.pop("encoder")
+#         if 'encoder' not in kwargs:
+#             self.encoder = transformers.AutoConfig.from_pretrained(
+#                 'aehrc/uniformer_base_tl_384',
+#                 projection_size=768,
+#                 trust_remote_code=True,
+#             )   
+#         else:
+#             self.encoder = kwargs.pop("encoder")
         
 
-        self.is_encoder_decoder = True
+#         self.is_encoder_decoder = True
 
-    @classmethod
-    def from_encoder_decoder_configs(
-        cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
-    ) -> PretrainedConfig:
+#     @classmethod
+#     def from_encoder_decoder_configs(
+#         cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
+#     ) -> PretrainedConfig:
 
-        logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
-        decoder_config.is_decoder = True
-        decoder_config.add_cross_attention = True
+#         logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
+#         decoder_config.is_decoder = True
+#         decoder_config.add_cross_attention = True
 
-        return cls(encoder=encoder_config, decoder=decoder_config, **kwargs)
+#         return cls(encoder=encoder_config, decoder=decoder_config, **kwargs)

generation_config.json

@@ -3,5 +3,5 @@
   "bos_token_id": 1,
   "eos_token_id": 2,
   "pad_token_id": 4,
-  "transformers_version": "4.39.3"
+  "transformers_version": "4.47.0"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:757ff7d2e55bf73d9a170d521fdfacc735b8226dcca11c32b5a20d2b2250ec48
-size 789964216
+oid sha256:00a9a6697b96ba73294054503626e877190b4c30b95d826d3ca3410d44739aed
+size 789967160

modelling_cxrmate_ed.py

@@ -14,7 +14,7 @@ from transformers import PreTrainedTokenizerFast, VisionEncoderDecoderModel
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_outputs import ModelOutput, Seq2SeqLMOutput
 from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import logging
+from transformers.utils import check_min_version, logging
 
 from .configuration_cxrmate_ed import CXRMateEDConfig
 from .dataset import PriorsDataset
@@ -108,74 +108,39 @@ class CXRStudyImagesEncoder(torch.nn.Module):
         return ModelOutput(last_hidden_state=last_hidden_state, attention_mask=attention_mask)
 
 
-
-class CXRMateEDModel(VisionEncoderDecoderModel):
+class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
 
     config_class = CXRMateEDConfig
 
-    def __init__(
-        self,
-        config: Optional[PretrainedConfig] = None,
-        encoder: Optional[PreTrainedModel] = None,
-        decoder: Optional[PreTrainedModel] = None,
-    ):
-
-        if decoder:
-            assert decoder.config.is_decoder, '"is_decoder" must be True for the given decoder'
-
-        if config is None and (encoder is None or decoder is None):
-            raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
-        if config is None:
-            config = CXRMateEDConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
-        else:
-            if not isinstance(config, self.config_class):
-                raise ValueError(f"Config: {config} has to be of type {self.config_class}")
-
-        config.tie_word_embeddings = False
-        config.is_encoder_decoder = False
-
-        # Initialize with config:
-        PreTrainedModel.__init__(self, config)
-
-        # Encoder:
-        if encoder is None:
-            encoder = transformers.AutoModel.from_pretrained(
-                'aehrc/uniformer_base_tl_384', 
-                config=config.encoder,
-                trust_remote_code=True,
-            )
-
-        # Decoder:
-        if decoder is None:
-            decoder = transformers.LlamaForCausalLM(config=config.decoder)
-
-        self.encoder = CXRStudyImagesEncoder(encoder, self.config.decoder)
-        self.decoder = decoder
-
-        if self.encoder.config.to_dict() != self.config.encoder.to_dict():
-            logger.warning(
-                f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
-                f" {self.config.encoder}"
-            )
-        if self.decoder.config.to_dict() != self.config.decoder.to_dict():
-            logger.warning(
-                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
-                f" {self.config.decoder}"
-            )
+    def __init__(self, config: CXRMateEDConfig):
+        
+        check_min_version("4.46.0.dev0")
 
-        self.encoder.config = self.config.encoder
-        self.decoder.config = self.config.decoder
+        super(transformers.LlavaPreTrainedModel, self).__init__(config)
+        
+        self.config = config
+        
+        self.vocab_size = config.text_config.vocab_size
+        
+        self.image_encoder = transformers.AutoModel.from_config(self.config.vision_config, trust_remote_code=True)
+        
+        self.language_model = transformers.AutoModelForCausalLM.from_config(
+            config.text_config,
+            attn_implementation=config._attn_implementation,
+        )
+        
+        self.image_encoder = CXRStudyImagesEncoder(self.image_encoder, config.text_config)
+    
+        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
 
-        assert config.decoder.is_decoder
-        assert not config.decoder.is_encoder_decoder
-        assert 'pad_token_id' in self.decoder.config.__dict__
-        assert 'time_delta_monotonic_inversion' in self.decoder.config.__dict__
-        assert 'add_time_deltas' in self.decoder.config.__dict__
-        assert 'history' in self.decoder.config.__dict__
-        assert 'tables_filter' in self.decoder.config.__dict__
-        assert 'prompt_report_sections_filter' in self.decoder.config.__dict__
+        # assert 'pad_token_id' in self.config.__dict__
+        # assert 'time_delta_monotonic_inversion' in self.config.__dict__
+        # assert 'add_time_deltas' in self.config.__dict__
+        # assert 'history' in self.config.__dict__
+        # assert 'tables_filter' in self.config.__dict__
+        # assert 'prompt_report_sections_filter' in self.config.__dict__
 
-        assert isinstance(self.decoder.config.time_delta_monotonic_inversion, bool)
+        # assert isinstance(self.config.time_delta_monotonic_inversion, bool)
 
         with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tables.json'), 'r') as f:
             self.tables = json.load(f)
@@ -186,8 +151,8 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'token_type_ids.json'), 'r') as f:
             self.token_type_to_token_type_id = json.load(f)
 
-        self.tables = {k: self.tables[k] for k in self.decoder.config.tables_filter}
-        self.tables['mimic_cxr_sectioned']['text_columns'] = self.decoder.config.prompt_report_sections_filter
+        self.tables = {k: self.tables[k] for k in self.config.tables_filter}
+        self.tables['mimic_cxr_sectioned']['text_columns'] = self.config.prompt_report_sections_filter
 
         for k in self.tables.keys():  
             if self.luts[k]['total'] > 0:        
@@ -196,179 +161,182 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
                     f'{k}_index_value_encoder', 
                     FNNEncoder(
                         num_features=self.luts[k]['total'], 
-                        intermediate_size=self.decoder.config.index_value_encoder_intermediate_size, 
-                        decoder_hidden_size=self.decoder.config.hidden_size,
+                        intermediate_size=self.config.index_value_encoder_intermediate_size, 
+                        decoder_hidden_size=self.config.hidden_size,
                     ),
                 )
                             
-        if self.decoder.config.add_time_deltas:
+        if self.config.add_time_deltas:
             self.time_delta_encoder = FNNEncoder(
                 num_features=1, 
-                intermediate_size=self.decoder.config.index_value_encoder_intermediate_size, 
-                decoder_hidden_size=self.decoder.config.hidden_size,
+                intermediate_size=self.config.index_value_encoder_intermediate_size, 
+                decoder_hidden_size=self.config.hidden_size,
             )
             
-        self.token_type_embeddings = torch.nn.Embedding(max(self.token_type_to_token_type_id.values()) + 1, self.decoder.config.hidden_size)
+        self.token_type_embeddings = torch.nn.Embedding(max(self.token_type_to_token_type_id.values()) + 1, self.config.hidden_size)
 
         self.time_delta_map = lambda x: 1 / math.sqrt(x + 1)
         self.zero_time_delta_value = self.time_delta_map(0)
         
         self.inf_time_delta_value = self.time_delta_map(float('inf'))
-
-    @classmethod
-    def from_encoder_decoder_pretrained(
-        cls,
-        encoder_pretrained_model_name_or_path: str = None,
-        decoder_pretrained_model_name_or_path: str = None,
-        *model_args,
-        **kwargs,
-    ) -> PreTrainedModel:
-        r"""
-        Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
-        checkpoints.
-
-
-        The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
-        the model, you need to first set it back in training mode with `model.train()`.
-
-        Params:
-            encoder_pretrained_model_name_or_path (`str`, *optional*):
-                Information necessary to initiate the image encoder. Can be either:
-
-                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. An
-                      example is `google/vit-base-patch16-224-in21k`.
-                    - A path to a *directory* containing model weights saved using
-                      [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
-                    - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
-                      this case, `from_tf` should be set to `True` and a configuration object should be provided as
-                      `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
-                      PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
-
-            decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
-                Information necessary to initiate the text decoder. Can be either:
-
-                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                    - A path to a *directory* containing model weights saved using
-                      [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
-                    - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
-                      this case, `from_tf` should be set to `True` and a configuration object should be provided as
-                      `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
-                      PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
-
-            model_args (remaining positional arguments, *optional*):
-                All remaning positional arguments will be passed to the underlying model's `__init__` method.
-
-            kwargs (remaining dictionary of keyword arguments, *optional*):
-                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
-                `output_attentions=True`).
-
-                - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
-                - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
-                - To update the parent model configuration, do not use a prefix for each configuration parameter.
-
-                Behaves differently depending on whether a `config` is provided or automatically loaded.
-
-        Example:
-
-        ```python
-        >>> from transformers import VisionEncoderDecoderModel
-
-        >>> # initialize a vit-bert from a pretrained ViT and a pretrained BERT model. Note that the cross-attention layers will be randomly initialized
-        >>> model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(
-        ...     "google/vit-base-patch16-224-in21k", "google-bert/bert-base-uncased"
-        ... )
-        >>> # saving model after fine-tuning
-        >>> model.save_pretrained("./vit-bert")
-        >>> # load fine-tuned model
-        >>> model = VisionEncoderDecoderModel.from_pretrained("./vit-bert")
-        ```"""
-
-        kwargs_encoder = {
-            argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
-        }
-
-        kwargs_decoder = {
-            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
-        }
-
-        # remove encoder, decoder kwargs from kwargs
-        for key in kwargs_encoder.keys():
-            del kwargs["encoder_" + key]
-        for key in kwargs_decoder.keys():
-            del kwargs["decoder_" + key]
-
-        # Load and initialize the encoder and decoder
-        # The distinction between encoder and decoder at the model level is made
-        # by the value of the flag `is_decoder` that we need to set correctly.
-        encoder = kwargs_encoder.pop("model", None)
-        if encoder is None:
-            if encoder_pretrained_model_name_or_path is None:
-                raise ValueError(
-                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
-                    "to be defined."
-                )
-
-            if "config" not in kwargs_encoder:
-                encoder_config, kwargs_encoder = transformers.AutoConfig.from_pretrained(
-                    encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
-                )
-
-                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
-                    logger.info(
-                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
-                        "from a decoder model. Cross-attention and casual mask are disabled."
-                    )
-                    encoder_config.is_decoder = False
-                    encoder_config.add_cross_attention = False
-
-                kwargs_encoder["config"] = encoder_config
-
-            encoder = transformers.AutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)
-
-        decoder = kwargs_decoder.pop("model", None)
-        if decoder is None:
-            if decoder_pretrained_model_name_or_path is None:
-                raise ValueError(
-                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
-                    "to be defined."
-                )
-
-            if "config" not in kwargs_decoder:
-                decoder_config, kwargs_decoder = transformers.AutoConfig.from_pretrained(
-                    decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
-                )
-
-                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
-                    logger.info(
-                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
-                        f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
-                        f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
-                    )
-                    decoder_config.is_decoder = True
-                    decoder_config.add_cross_attention = False
-
-                kwargs_decoder["config"] = decoder_config
-
-            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
-                logger.warning(
-                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
-                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
-                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
-                    "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
-                    "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
-                )
-
-            decoder = transformers.AutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
-
-        # instantiate config with corresponding kwargs
-        config = CXRMateEDConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
-
-        # make sure input & output embeddings is not tied
-        config.tie_word_embeddings = False
         
-        config.is_encoder_decoder = False
+        self.post_init()
+
+
+    # @classmethod
+    # def from_encoder_decoder_pretrained(
+    #     cls,
+    #     encoder_pretrained_model_name_or_path: str = None,
+    #     decoder_pretrained_model_name_or_path: str = None,
+    #     *model_args,
+    #     **kwargs,
+    # ) -> PreTrainedModel:
+    #     r"""
+    #     Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
+    #     checkpoints.
+
+
+    #     The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
+    #     the model, you need to first set it back in training mode with `model.train()`.
+
+    #     Params:
+    #         encoder_pretrained_model_name_or_path (`str`, *optional*):
+    #             Information necessary to initiate the image encoder. Can be either:
+
+    #                 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. An
+    #                   example is `google/vit-base-patch16-224-in21k`.
+    #                 - A path to a *directory* containing model weights saved using
+    #                   [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+    #                 - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
+    #                   this case, `from_tf` should be set to `True` and a configuration object should be provided as
+    #                   `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
+    #                   PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
+
+    #         decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+    #             Information necessary to initiate the text decoder. Can be either:
+
+    #                 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+    #                 - A path to a *directory* containing model weights saved using
+    #                   [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+    #                 - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
+    #                   this case, `from_tf` should be set to `True` and a configuration object should be provided as
+    #                   `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
+    #                   PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
+
+    #         model_args (remaining positional arguments, *optional*):
+    #             All remaning positional arguments will be passed to the underlying model's `__init__` method.
+
+    #         kwargs (remaining dictionary of keyword arguments, *optional*):
+    #             Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+    #             `output_attentions=True`).
+
+    #             - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
+    #             - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
+    #             - To update the parent model configuration, do not use a prefix for each configuration parameter.
+
+    #             Behaves differently depending on whether a `config` is provided or automatically loaded.
+
+    #     Example:
+
+    #     ```python
+    #     >>> from transformers import VisionEncoderDecoderModel
+
+    #     >>> # initialize a vit-bert from a pretrained ViT and a pretrained BERT model. Note that the cross-attention layers will be randomly initialized
+    #     >>> model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(
+    #     ...     "google/vit-base-patch16-224-in21k", "google-bert/bert-base-uncased"
+    #     ... )
+    #     >>> # saving model after fine-tuning
+    #     >>> model.save_pretrained("./vit-bert")
+    #     >>> # load fine-tuned model
+    #     >>> model = VisionEncoderDecoderModel.from_pretrained("./vit-bert")
+    #     ```"""
+
+    #     kwargs_encoder = {
+    #         argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
+    #     }
+
+    #     kwargs_decoder = {
+    #         argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+    #     }
+
+    #     # remove encoder, decoder kwargs from kwargs
+    #     for key in kwargs_encoder.keys():
+    #         del kwargs["encoder_" + key]
+    #     for key in kwargs_decoder.keys():
+    #         del kwargs["decoder_" + key]
+
+    #     # Load and initialize the encoder and decoder
+    #     # The distinction between encoder and decoder at the model level is made
+    #     # by the value of the flag `is_decoder` that we need to set correctly.
+    #     encoder = kwargs_encoder.pop("model", None)
+    #     if encoder is None:
+    #         if encoder_pretrained_model_name_or_path is None:
+    #             raise ValueError(
+    #                 "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
+    #                 "to be defined."
+    #             )
+
+    #         if "config" not in kwargs_encoder:
+    #             encoder_config, kwargs_encoder = transformers.AutoConfig.from_pretrained(
+    #                 encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
+    #             )
+
+    #             if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+    #                 logger.info(
+    #                     f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
+    #                     "from a decoder model. Cross-attention and casual mask are disabled."
+    #                 )
+    #                 encoder_config.is_decoder = False
+    #                 encoder_config.add_cross_attention = False
+
+    #             kwargs_encoder["config"] = encoder_config
+
+    #         encoder = transformers.AutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)
+
+    #     decoder = kwargs_decoder.pop("model", None)
+    #     if decoder is None:
+    #         if decoder_pretrained_model_name_or_path is None:
+    #             raise ValueError(
+    #                 "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
+    #                 "to be defined."
+    #             )
+
+    #         if "config" not in kwargs_decoder:
+    #             decoder_config, kwargs_decoder = transformers.AutoConfig.from_pretrained(
+    #                 decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
+    #             )
+
+    #             if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
+    #                 logger.info(
+    #                     f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
+    #                     f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
+    #                     f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
+    #                 )
+    #                 decoder_config.is_decoder = True
+    #                 decoder_config.add_cross_attention = False
+
+    #             kwargs_decoder["config"] = decoder_config
+
+    #         if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
+    #             logger.warning(
+    #                 f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
+    #                 f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
+    #                 "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
+    #                 "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
+    #                 "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
+    #             )
+
+    #         decoder = transformers.AutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
+
+    #     # instantiate config with corresponding kwargs
+    #     config = CXRMateEDConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
+
+    #     # make sure input & output embeddings is not tied
+    #     config.tie_word_embeddings = False
+        
+    #     config.is_encoder_decoder = False
 
-        return cls(encoder=encoder, decoder=decoder, config=config)
+    #     return cls(encoder=encoder, decoder=decoder, config=config)
 
     def forward(
         self,
@@ -393,14 +361,17 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
             argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
         }
   
-        assert decoder_attention_mask.dtype == torch.long, f'The dtype for {decoder_attention_mask} was {decoder_attention_mask.dtype}. It should be torch.long'
-
         if decoder_inputs_embeds is None:
-            decoder_inputs_embeds = self.decoder.get_input_embeddings()(decoder_input_ids)
+            decoder_inputs_embeds = self.language_model.get_input_embeddings()(decoder_input_ids)
         decoder_inputs_embeds += self.token_type_embeddings(decoder_token_type_ids)
 
+        if decoder_attention_mask.dim() == 4:
+            assert decoder_attention_mask.dtype == decoder_inputs_embeds.dtype, f'The dtype for {decoder_attention_mask} was {decoder_attention_mask.dtype}. It should be {decoder_inputs_embeds.dtype}'
+        else:
+            assert decoder_attention_mask.dtype == torch.long, f'The dtype for {decoder_attention_mask} was {decoder_attention_mask.dtype}. It should be torch.long'
+
         # Generation:
-        decoder_outputs = self.decoder(
+        decoder_outputs = self.language_model(
             inputs_embeds=decoder_inputs_embeds,
             attention_mask=decoder_attention_mask,
             position_ids=decoder_position_ids,
@@ -417,7 +388,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         if labels is not None:
             logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
             loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.reshape(-1))
+            loss = loss_fct(logits.reshape(-1, self.vocab_size), labels.reshape(-1))
 
         if not return_dict:
             if loss is not None:
@@ -448,20 +419,22 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
             https://github.com/huggingface/transformers/blob/main/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py#L660
         """
 
-        report_attention_mask = (input_ids != self.decoder.config.pad_token_id).long()
+        report_attention_mask = (input_ids != self.config.pad_token_id).long()
 
-        if past_key_values is None:
+        if len(past_key_values) == 0:
             
             # 4D attention mask:
-            decoder_attention_mask = self.create_4d_attention_mask_mixed_causality(prompt_attention_mask, report_attention_mask)
-
+            decoder_attention_mask = self.create_4d_attention_mask_mixed_causality(
+                prompt_attention_mask, report_attention_mask, dtype=kwargs['decoder_inputs_embeds'].dtype,
+            )
+            
             # Position identifiers accounting for padding:
             report_position_ids = report_attention_mask.cumsum(-1) + prompt_position_ids.max(dim=1).values[:, None]
             report_position_ids.masked_fill_(report_attention_mask == 0, 1)
             decoder_position_ids = torch.cat([prompt_position_ids, report_position_ids], dim=1)
 
             # `inputs_embeds` are only to be used in the 1st generation step:
-            inputs_embeds = torch.cat([kwargs['decoder_inputs_embeds'], self.decoder.get_input_embeddings()(input_ids)], dim=1)
+            inputs_embeds = torch.cat([kwargs['decoder_inputs_embeds'], self.language_model.get_input_embeddings()(input_ids)], dim=1)
 
             decoder_token_type_ids = self.token_ids_to_token_type_ids(
                 input_ids, special_token_ids, 
@@ -483,7 +456,9 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         else:
             
             # 4D attention mask:
-            decoder_attention_mask = self.create_4d_attention_mask_mixed_causality_past_key_values(prompt_attention_mask, report_attention_mask)
+            decoder_attention_mask = self.create_4d_attention_mask_mixed_causality_past_key_values(
+                prompt_attention_mask, report_attention_mask, dtype=kwargs['decoder_inputs_embeds'].dtype,
+            )
 
             # Position identifiers accounting for padding:
             decoder_position_ids = report_attention_mask.cumsum(-1) + prompt_position_ids.max(dim=1).values[:, None]
@@ -863,7 +838,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         time_delta.append(tokenized['time_delta'])
 
         # Image encoder:
-        encoder_outputs = self.encoder(images)  
+        encoder_outputs = self.image_encoder(images)  
         inputs_embeds.append(encoder_outputs[0])
         
         inputs_per_image = encoder_outputs[0].shape[-2] // images.shape[1]
@@ -883,14 +858,14 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
 
         # Compute embeddings from token identifiers:
         input_ids = torch.cat(input_ids, dim=1)
-        inputs_embeds.append(self.decoder.get_input_embeddings()(input_ids))
+        inputs_embeds.append(self.language_model.get_input_embeddings()(input_ids))
         
         # Concatentate time deltas and input embeddings before adding time delta embedding to prompt:
         time_delta = torch.cat(time_delta, dim=1)
         inputs_embeds = torch.cat(inputs_embeds, dim=1)
 
         # Add time delta embeddings to prompt:
-        if time_delta.shape[1] > 0 and self.decoder.config.add_time_deltas:
+        if time_delta.shape[1] > 0 and self.config.add_time_deltas:
             time_delta = time_delta.to(dtype=inputs_embeds.dtype)
             inputs_embeds += self.time_delta_encoder(time_delta)
             
@@ -902,7 +877,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
     
         # Tokenize report:
         if tokenized_report is not None:
-            inputs_embeds = torch.cat([inputs_embeds, self.decoder.get_input_embeddings()(tokenized_report['decoder_input_ids'])], dim=1)
+            inputs_embeds = torch.cat([inputs_embeds, self.language_model.get_input_embeddings()(tokenized_report['decoder_input_ids'])], dim=1)
             
             report_token_type_ids = self.token_ids_to_token_type_ids(
                 token_ids=tokenized_report['decoder_input_ids'], 
@@ -917,7 +892,8 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
             position_ids = torch.cat([position_ids, report_position_ids], dim=1)
             
             # 4D attention mask:
-            attention_mask = self.create_4d_attention_mask_mixed_causality(attention_mask, tokenized_report['decoder_attention_mask'])
+            attention_mask = self.create_4d_attention_mask_mixed_causality(attention_mask, tokenized_report['decoder_attention_mask'], dtype=inputs_embeds.dtype)
+            # attention_mask = self.create_4d_attention_mask_mixed_causality(attention_mask, tokenized_report['decoder_attention_mask'])
             # attention_mask_diagonal = torch.diagonal(attention_mask[:, 0], dim1=1, dim2=2)
 
         else:
@@ -934,7 +910,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         return inputs_embeds, attention_mask, token_type_ids, position_ids, bos_token_ids
     
     @staticmethod
-    def create_4d_attention_mask_mixed_causality(non_causal_2d_attention_mask, causal_2d_attention_mask):
+    def create_4d_attention_mask_mixed_causality(non_causal_2d_attention_mask, causal_2d_attention_mask, dtype):
     
         prompt_seq_len = non_causal_2d_attention_mask.shape[-1] 
         report_seq_len = causal_2d_attention_mask.shape[-1]
@@ -982,22 +958,91 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
 
         mixed_causality_4d_attention_mask = torch.cat((left, right), dim=-1)
         
+        mixed_causality_4d_attention_mask = mixed_causality_4d_attention_mask.to(dtype=dtype)
+        mixed_causality_4d_attention_mask[mixed_causality_4d_attention_mask == 0] = torch.finfo(mixed_causality_4d_attention_mask.dtype).min
+        mixed_causality_4d_attention_mask[mixed_causality_4d_attention_mask == 1] = 0.0
+        
         return mixed_causality_4d_attention_mask
     
     @staticmethod
-    def create_4d_attention_mask_mixed_causality_past_key_values(non_causal_2d_attention_mask, causal_2d_attention_mask):
+    def create_4d_attention_mask_mixed_causality_past_key_values(non_causal_2d_attention_mask, causal_2d_attention_mask, dtype):
     
         non_causal_2d_attention_mask = non_causal_2d_attention_mask[:, None, None, :]
         causal_2d_attention_mask = causal_2d_attention_mask[:, None, None, :]
 
         mixed_causality_4d_attention_mask = torch.cat((non_causal_2d_attention_mask, causal_2d_attention_mask), dim=-1)
+        
+        mixed_causality_4d_attention_mask = mixed_causality_4d_attention_mask.to(dtype=dtype)
+        mixed_causality_4d_attention_mask[mixed_causality_4d_attention_mask == 0] = torch.finfo(mixed_causality_4d_attention_mask.dtype).min
+        mixed_causality_4d_attention_mask[mixed_causality_4d_attention_mask == 1] = 0.0
+        
         return mixed_causality_4d_attention_mask
+        
+    # @staticmethod
+    # def create_4d_attention_mask_mixed_causality(non_causal_2d_attention_mask, causal_2d_attention_mask):
+    
+    #     prompt_seq_len = non_causal_2d_attention_mask.shape[-1] 
+    #     report_seq_len = causal_2d_attention_mask.shape[-1]
+        
+    #     non_causal_2d_attention_mask = non_causal_2d_attention_mask[:, None, None, :]
+    #     causal_2d_attention_mask = causal_2d_attention_mask[:, None, None, :]
+    
+    #     # Upper left of attention matrix:
+    #     upper_left = non_causal_2d_attention_mask.expand(-1, -1, prompt_seq_len, -1)
+    #     upper_left = upper_left * non_causal_2d_attention_mask
+    #     upper_left = upper_left * non_causal_2d_attention_mask.permute(0, 1, 3, 2)
+        
+    #     causal_mask = torch.tril(
+    #         torch.ones(
+    #             (
+    #                 report_seq_len, 
+    #                 report_seq_len,
+    #             ), 
+    #             dtype=torch.long, 
+    #             device=causal_2d_attention_mask.device,
+    #         ),
+    #     )   
+        
+    #     # Lower right of attention matrix:
+    #     lower_right = causal_2d_attention_mask.expand(-1, -1, report_seq_len, -1)
+    #     lower_right = lower_right * causal_2d_attention_mask.permute(0, 1, 3, 2)
+    #     lower_right = lower_right * causal_mask
+        
+    #     # Upper right of attention matrix:
+    #     upper_right = torch.zeros(
+    #         causal_2d_attention_mask.shape[0], 
+    #         1, 
+    #         prompt_seq_len, 
+    #         report_seq_len, 
+    #         dtype=torch.long, 
+    #         device=causal_2d_attention_mask.device,
+    #     )
+        
+    #     # Lower left of attention matrix:
+    #     lower_left = non_causal_2d_attention_mask.expand(-1, -1, report_seq_len, -1)
+    #     lower_left = lower_left * causal_2d_attention_mask.permute(0, 1, 3, 2)
+            
+    #     left = torch.cat((upper_left, lower_left), dim=2)
+    #     right = torch.cat((upper_right, lower_right), dim=2)
+
+    #     mixed_causality_4d_attention_mask = torch.cat((left, right), dim=-1)
+        
+    #     return mixed_causality_4d_attention_mask
+    
+    # @staticmethod
+    # def create_4d_attention_mask_mixed_causality_past_key_values(non_causal_2d_attention_mask, causal_2d_attention_mask):
+    
+    #     non_causal_2d_attention_mask = non_causal_2d_attention_mask[:, None, None, :]
+    #     causal_2d_attention_mask = causal_2d_attention_mask[:, None, None, :]
+
+    #     mixed_causality_4d_attention_mask = torch.cat((non_causal_2d_attention_mask, causal_2d_attention_mask), dim=-1)
+    #     return mixed_causality_4d_attention_mask
     
     def position_ids_from_time_deltas_and_attention_mask(self, time_deltas, attention_mask):
-        mask_value = torch.finfo(time_deltas.dtype).max if self.decoder.config.time_delta_monotonic_inversion else torch.finfo(time_deltas.dtype).min
+        mask_value = torch.finfo(time_deltas.dtype).max if self.config.time_delta_monotonic_inversion else torch.finfo(time_deltas.dtype).min
         
         masked_time_deltas = torch.where(attention_mask == 1, time_deltas[:, :, 0], mask_value)
-        _, col_indices = torch.sort(masked_time_deltas, descending=not self.decoder.config.time_delta_monotonic_inversion)
+        _, col_indices = torch.sort(masked_time_deltas, descending=not self.config.time_delta_monotonic_inversion)
         
         num_rows, num_cols, _ = time_deltas.shape
 
@@ -1081,7 +1126,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
             index_map = {study_id: idx for idx, study_id in enumerate(train_set_study_ids)}
             indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
             indices.sort()        
-            train_set = PriorsDataset(train_set, self.decoder.config.history, self.time_delta_map)
+            train_set = PriorsDataset(train_set, self.config.history, self.time_delta_map)
             train_set.set_transform(train_set_transform)
             train_set = Subset(train_set, indices)
         else:
@@ -1096,7 +1141,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
             index_map = {study_id: idx for idx, study_id in enumerate(val_set_study_ids)}
             indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
             indices.sort()    
-            val_set = PriorsDataset(val_set, self.decoder.config.history, self.time_delta_map)    
+            val_set = PriorsDataset(val_set, self.config.history, self.time_delta_map)    
             val_set.set_transform(test_set_transform)
             val_set = Subset(val_set, indices)
         else:
@@ -1110,7 +1155,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         index_map = {study_id: idx for idx, study_id in enumerate(test_set_study_ids)}
         indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
         indices.sort()        
-        test_set = PriorsDataset(test_set, self.decoder.config.history, self.time_delta_map)    
+        test_set = PriorsDataset(test_set, self.config.history, self.time_delta_map)    
         test_set.set_transform(test_set_transform)
         test_set = Subset(test_set, indices)
         
@@ -1163,7 +1208,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         index_map = {study_id: idx for idx, study_id in enumerate(train_set_study_ids)}
         indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
         indices.sort()        
-        train_set = PriorsDataset(train_set, self.decoder.config.history, self.time_delta_map)
+        train_set = PriorsDataset(train_set, self.config.history, self.time_delta_map)
         train_set.set_transform(train_set_transform)
         train_set = Subset(train_set, indices)
 
@@ -1175,7 +1220,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         index_map = {study_id: idx for idx, study_id in enumerate(val_set_study_ids)}
         indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
         indices.sort()        
-        val_set = PriorsDataset(val_set, self.decoder.config.history, self.time_delta_map)    
+        val_set = PriorsDataset(val_set, self.config.history, self.time_delta_map)    
         val_set.set_transform(test_set_transform)
         val_set = Subset(val_set, indices)
 
@@ -1187,7 +1232,7 @@ class CXRMateEDModel(VisionEncoderDecoderModel):
         index_map = {study_id: idx for idx, study_id in enumerate(test_set_study_ids)}
         indices = [index_map[study_id] for study_id in study_ids if study_id in index_map]
         indices.sort()        
-        test_set = PriorsDataset(test_set, self.decoder.config.history, self.time_delta_map)    
+        test_set = PriorsDataset(test_set, self.config.history, self.time_delta_map)    
         test_set.set_transform(test_set_transform)
         test_set = Subset(test_set, indices)