Delete modelling_variable.py

modelling_variable.py

@@ -1,425 +0,0 @@
-import os
-from typing import Any, Optional, Tuple, Union
-
-import torch
-import transformers
-from torch.nn import CrossEntropyLoss
-from transformers import PreTrainedTokenizerFast, VisionEncoderDecoderModel
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
-from transformers.modeling_utils import PreTrainedModel
-from transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder import \
-    VisionEncoderDecoderConfig
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class CvtWithProjectionHeadConfig(transformers.CvtConfig):
-    def __init__(self, projection_size: int = None, **kwargs: Any) -> None:
-        super().__init__(**kwargs)
-        self.projection_size = projection_size
-
-
-class ModelOutputWithProjectionEmbedding(transformers.modeling_outputs.ModelOutput):
-    last_hidden_state: torch.FloatTensor
-    attention_mask: torch.FloatTensor
-
-
-class CvtProjectionHead(torch.nn.Module):
-
-    def __init__(self, config) -> None:
-        super().__init__()
-
-        # https://github.com/huggingface/transformers/blob/68287689f2f0d8b7063c400230b3766987abf18d/src/transformers/models/cvt/modeling_cvt.py#L657
-        self.layer_norm = torch.nn.LayerNorm(config.embed_dim[-1], eps=config.layer_norm_eps)
-
-        # No bias as following layer normalisation with bias:
-        self.projection = torch.nn.Linear(config.embed_dim[-1], config.projection_size, bias=False)
-
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.layer_norm(x)
-        x = self.projection(x)
-        return x
-
-
-class VariableCvtWithProjectionHead(transformers.CvtPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.cvt = transformers.CvtModel(config, add_pooling_layer=False)
-        self.projection_head = CvtProjectionHead(config)
-
-        # Initialize weights and apply final processing:
-        self.post_init()
-
-    def forward(
-        self,
-        pixel_values: Optional[torch.Tensor] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, ModelOutputWithProjectionEmbedding]:
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # Flatten the batch and study_id dimensions:
-        outputs = self.cvt(
-            pixel_values.view(-1, *pixel_values.shape[2:]),
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        # Flatten h x w:
-        last_hidden_state = torch.flatten(outputs.last_hidden_state, 2)
-
-        # Project the features for each spatial position to the decoder's hidden size:
-        projection = self.projection_head(torch.permute(last_hidden_state, [0, 2, 1]))
-
-        # Concatenate the features for each chest X-ray:
-        projection = projection.view(pixel_values.shape[0], -1, projection.shape[-1])
-
-        # Derive the attention mask from the pixel values:
-        attention_mask = (pixel_values[:, :, 0, 0, 0] != 0.0).repeat_interleave(last_hidden_state.shape[-1], dim=1)
-
-        if not return_dict:
-            return projection
-
-        return ModelOutputWithProjectionEmbedding(
-            last_hidden_state=projection, attention_mask=attention_mask,
-        )
-    
-
-class VariableCXREncoderDecoderModel(VisionEncoderDecoderModel):
-
-    config_class = VisionEncoderDecoderConfig
-    base_model_prefix = "vision_encoder_decoder"
-    main_input_name = "pixel_values"
-    supports_gradient_checkpointing = True
-
-    def __init__(        
-        self,
-        config: Optional[PretrainedConfig] = None,
-        encoder: Optional[PreTrainedModel] = None,
-        decoder: Optional[PreTrainedModel] = None,
-    ):
-
-        if decoder:
-            assert decoder.config.add_cross_attention, '"add_cross_attention" must be True for the given 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 = VisionEncoderDecoderConfig.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
-
-        # initialize with config
-        PreTrainedModel.__init__(self, config)
-
-        # Encoder:
-        if encoder is None:
-            encoder = VariableCvtWithProjectionHead(config=config.encoder)
-
-        # Decoder:
-        if decoder is None:
-            decoder = transformers.BertLMHeadModel(config=config.decoder)
-
-        self.encoder = encoder
-        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}"
-            )
-            
-        self.encoder.config = self.config.encoder
-        self.decoder.config = self.config.decoder
-
-        # config.add_cross_attention = True
-        # config.is_decoder = True
-
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        decoder_input_ids: Optional[torch.LongTensor] = None,
-        decoder_attention_mask: Optional[torch.BoolTensor] = None,
-        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
-        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
-
-        kwargs_decoder = {
-            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
-        }
-
-        if encoder_outputs is None:
-            if pixel_values is None:
-                raise ValueError("You have to specify pixel_values")
-
-            encoder_outputs = self.encoder(
-                pixel_values,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-                **kwargs_encoder,
-            )  # CvT does not support output_attentions.
-
-        elif isinstance(encoder_outputs, tuple):
-            encoder_outputs = BaseModelOutput(*encoder_outputs)
-
-        encoder_hidden_states = encoder_outputs[0]
-        
-        decoder_outputs = self.decoder(
-            input_ids=decoder_input_ids,
-            attention_mask=decoder_attention_mask,
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_attention_mask=encoder_outputs.attention_mask,
-            inputs_embeds=decoder_inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            use_cache=use_cache,
-            past_key_values=past_key_values,
-            return_dict=return_dict,
-            **kwargs_decoder,
-        )
-
-        # Loss:
-        loss = None
-        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))
-
-        if not return_dict:
-            if loss is not None:
-                return (loss,) + decoder_outputs + encoder_outputs
-            else:
-                return decoder_outputs + encoder_outputs
-
-        return Seq2SeqLMOutput(
-            loss=loss,
-            logits=decoder_outputs.logits,
-            past_key_values=decoder_outputs.past_key_values,
-            decoder_hidden_states=decoder_outputs.hidden_states,
-            decoder_attentions=decoder_outputs.attentions,
-            cross_attentions=decoder_outputs.cross_attentions,
-            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
-            # encoder_hidden_states=encoder_outputs.hidden_states,
-            # encoder_attentions=encoder_outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        special_token_ids,
-        past_key_values=None,
-        attention_mask=None,
-        use_cache=None,
-        encoder_outputs=None,
-        **kwargs,
-    ):
-        """
-        Modification of: 
-            https://github.com/huggingface/transformers/blob/main/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py#L660
-        """
-
-        decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)
-        decoder_attention_mask = decoder_inputs['attention_mask'] if 'attention_mask' in decoder_inputs else None
-
-        if not past_key_values:
-            token_type_ids = self.token_ids_to_token_type_ids(input_ids, special_token_ids)
-        else:
-            token_type_ids = self.token_ids_to_token_type_ids_past(input_ids, special_token_ids)
-
-        input_dict = {
-            'attention_mask': attention_mask,
-            'decoder_attention_mask': decoder_attention_mask,
-            'decoder_input_ids': decoder_inputs['input_ids'],
-            'decoder_token_type_ids': token_type_ids,
-            'encoder_outputs': encoder_outputs,
-            'past_key_values': decoder_inputs['past_key_values'],
-            'use_cache': use_cache,
-        }
-        return input_dict
-    
-    def token_ids_to_token_type_ids(self, token_ids, special_token_ids, token_type_id_sections=None):
-        """
-        Extract token type identifiers from the token identifiers.
-
-        Argument/s:
-            token_ids - token identifiers.
-            special_token_ids - special token identifiers that indicate the separation between sections.
-            token_type_id_section - token type identifier for each section.
-
-        Returns:
-            token_type_ids - token type identifiers.
-        """
-
-        token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))
-
-        mbatch_size, seq_len = token_ids.shape
-        token_type_ids = torch.full_like(token_ids, token_type_id_sections[0], dtype=torch.long, device=token_ids.device)
-
-        for i, j in enumerate(special_token_ids):
-            # Find first occurrence of special tokens that indicate the boundary between sections:
-            cols = (token_ids == j).int().argmax(dim=1)
-            rows = torch.arange(mbatch_size, device=token_ids.device)
-
-            # https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer.create_token_type_ids_from_sequences.example
-            cols += 1
-
-            # Ensure that the column index is not out of bounds. If 0, then token_id not present.
-            # This is safe as index 0 is always a special token (now equal to 1 due to +1):
-            rows = rows[torch.logical_and(cols != 1, cols < seq_len)]
-            cols = cols[torch.logical_and(cols != 1, cols < seq_len)]
-
-            # Indices to that correspond to the second sequence:
-            if rows.nelement() != 0:
-                ids = torch.stack([
-                    torch.stack([x, z]) for (x, y) in zip(rows, cols) for z in torch.arange(
-                        y, seq_len, device=token_ids.device,
-                    )
-                ])
-
-                token_type_ids[ids[:, 0], ids[:, 1]] = token_type_id_sections[i + 1]
-
-        return token_type_ids
-
-    def token_ids_to_token_type_ids_past(self, token_ids, special_token_ids, token_type_id_sections=None):
-        """
-        Extract token type identifiers from the token identifiers if past != None.
-
-        Argument/s:
-            token_ids - token identifiers.
-            special_token_ids - special token identifiers that indicate the separation between sections.
-
-        Returns:
-            token_type_ids - token type identifiers.
-        """
-
-        token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))
-        token_type_ids = torch.full([token_ids.shape[0], 1], token_type_id_sections[0], dtype=torch.long, device=token_ids.device)
-
-        # https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer.create_token_type_ids_from_sequences.example
-        token_ids = token_ids[:, :-1]
-
-        for i, j in enumerate(special_token_ids):
-
-            # Find first occurrence of special token, which indicates the boundary between sections:
-            exists = torch.any(token_ids == j, dim=1, keepdim=True)
-            token_type_ids[exists] = token_type_id_sections[i + 1]
-
-        return token_type_ids
-    
-    def tokenize_report_teacher_forcing(self, findings: str, impression: str, tokenizer: PreTrainedTokenizerFast, max_len: int):
-        """
-        Tokenize the reports and creates the inputs and targets for teacher forcing.
-
-        Argument/s:
-            findings - findings section.
-            impression - impression section.
-            return_token_type_ids - return the token type identifiers.
-            tokenizer - Hugging Face tokenizer.
-            max_len - maximum number of tokens.
-
-        Returns:
-            decoder_input_ids - the token identifiers for the input of the decoder.
-            decoder_attention_mask - the attention mask for the decoder_input_ids.
-            label_ids - the label token identifiers for the decoder.
-        """
-
-        # Prepare the sections for the tokenizer by placing special tokens between each section:
-        report = [f'{tokenizer.bos_token}{i}{tokenizer.sep_token}{j}{tokenizer.eos_token}' for i, j in
-                  zip(findings, impression)]
-
-        # Tokenize the report:
-        tokenized = tokenizer(
-            report,
-            padding='longest',
-            truncation=True,
-            max_length=max_len + 1,  # +1 to account for the bias between input and target.
-            return_tensors='pt',
-            return_token_type_ids=False,
-            add_special_tokens=False,
-        ).to(self.device)
-
-        # Modify for language modelling:
-        batch_dict = {
-
-            # Labels for the decoder (shifted right by one for autoregression):
-            'label_ids': tokenized['input_ids'][:, 1:].detach().clone(),
-
-            # Remove last token identifier to match the sequence length of the labels:
-            'decoder_input_ids': tokenized['input_ids'][:, :-1],
-
-            # Attention mask for the decoder_input_ids (remove first token so that the eos_token_id is not considered):
-            'decoder_attention_mask': tokenized['attention_mask'][:, 1:],
-        }
-
-        return batch_dict
-
-    def split_and_decode_sections(self, token_ids, special_token_ids, tokenizer: PreTrainedTokenizerFast):
-        """
-        Split the token identifiers into sections, then convert the token identifiers into strings.
-
-        Argument/s:
-            token_ids - token identifiers.
-            special_token_ids - special token identifiers that indicate the end of each section.
-            tokenizer - Hugging Face tokenizer.
-
-        Returns:
-            token_type_ids - token type identifiers.
-        """
-
-        _, seq_len = token_ids.shape
-
-        # The number of sections is the same as the number of special_token_ids:
-        num_sections = len(special_token_ids)
-
-        sections = {k: [] for k in range(num_sections)}
-
-        for i in token_ids:
-            prev_col = 0
-            for j, k in enumerate(special_token_ids):
-
-                # The maximum sequence length was exceeded, thus no more tokens:
-                if prev_col >= seq_len:
-                    sections[j].append('')
-                    continue
-
-                # Find first occurrence of special tokens that indicate the boundary between sections:
-                col = (i == k).int().argmax().item()
-
-                # If equal to 0, token was not found, set the column to the sequence length (as the decoder exceeded
-                # the maximum sequence length):
-                if col == 0:
-                    col = seq_len
-
-                # Extract section token identifiers:
-                section_token_ids = i[prev_col:col]
-                prev_col = col
-                section_string = tokenizer.decode(section_token_ids, skip_special_tokens=True)
-
-                sections[j].append(section_string)
-
-        return tuple(sections.values())