aehrc/cxrmate-tf
0116
1import os2import warnings3from dataclasses import dataclass4from typing import Any, Optional, Tuple, Union5 6import torch7import transformers8from peft import LoraConfig, TaskType, get_peft_config, get_peft_model9from torch.nn import CrossEntropyLoss10from transformers import AutoModel, PreTrainedTokenizerFast, VisionEncoderDecoderModel11from transformers.configuration_utils import PretrainedConfig12from transformers.modeling_outputs import BaseModelOutput, ModelOutput, Seq2SeqLMOutput13from transformers.modeling_utils import PreTrainedModel14from transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder import (15 VisionEncoderDecoderConfig,16)17from transformers.utils import logging18 19logger = logging.get_logger(__name__)20 21 22class CvtWithProjectionHeadConfig(transformers.CvtConfig):23 def __init__(self, projection_size: int = None, **kwargs: Any) -> None:24 super().__init__(**kwargs)25 self.projection_size = projection_size26 27 28class CvtProjectionHead(torch.nn.Module):29 30 def __init__(self, config) -> None:31 super().__init__()32 33 # https://github.com/huggingface/transformers/blob/68287689f2f0d8b7063c400230b3766987abf18d/src/transformers/models/cvt/modeling_cvt.py#L65734 self.layer_norm = torch.nn.LayerNorm(config.embed_dim[-1], eps=config.layer_norm_eps)35 36 # No bias as following layer normalisation with bias:37 self.projection = torch.nn.Linear(config.embed_dim[-1], config.projection_size, bias=False)38 39 40 def forward(self, x: torch.Tensor) -> torch.Tensor:41 x = self.layer_norm(x)42 x = self.projection(x)43 return x44 45 46class MultiCvtWithProjectionHead(transformers.CvtPreTrainedModel):47 def __init__(self, config):48 super().__init__(config)49 50 self.cvt = transformers.CvtModel(config, add_pooling_layer=False)51 self.projection_head = CvtProjectionHead(config)52 53 # Initialize weights and apply final processing:54 self.post_init()55 56 def forward(57 self,58 pixel_values: Optional[torch.Tensor] = None,59 output_hidden_states: Optional[bool] = None,60 return_dict: Optional[bool] = None,61 output_attentions: Optional[bool] = None,62 ) -> Union[Tuple, ModelOutput]:63 64 return_dict = return_dict if return_dict is not None else self.config.use_return_dict65 66 # Flatten the batch and study_id dimensions:67 outputs = self.cvt(68 pixel_values.view(-1, *pixel_values.shape[2:]),69 output_hidden_states=output_hidden_states,70 return_dict=return_dict,71 )72 73 # Flatten h x w:74 last_hidden_state = torch.flatten(outputs.last_hidden_state, 2)75 76 # Project the features for each spatial position to the decoder's hidden size:77 projection = self.projection_head(torch.permute(last_hidden_state, [0, 2, 1]))78 79 # Concatenate the features for each chest X-ray:80 projection = projection.view(pixel_values.shape[0], -1, projection.shape[-1])81 82 # Derive the attention mask from the pixel values:83 attention_mask = (pixel_values[:, :, 0, 0, 0] != 0.0).repeat_interleave(last_hidden_state.shape[-1], dim=1)84 85 if not return_dict:86 return projection87 88 return ModelOutput(89 last_hidden_state=projection, attention_mask=attention_mask,90 )91 92 93class LongitudinalPromptMultiCXREncoderDecoderModel(VisionEncoderDecoderModel):94 95 config_class = VisionEncoderDecoderConfig96 base_model_prefix = "vision_encoder_decoder"97 main_input_name = "pixel_values"98 supports_gradient_checkpointing = True99 100 def __init__( 101 self,102 config: Optional[PretrainedConfig] = None,103 encoder: Optional[PreTrainedModel] = None,104 decoder: Optional[PreTrainedModel] = None,105 encoder_decoder_ckpt_name: Optional[str] = None,106 ):107 108 if decoder:109 assert decoder.config.add_cross_attention, '"add_cross_attention" must be True for the given decoder'110 assert decoder.config.is_decoder, '"is_decoder" must be True for the given decoder'111 112 if config is None and (encoder is None or decoder is None):113 raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")114 if config is None:115 config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)116 else:117 if not isinstance(config, self.config_class):118 raise ValueError(f"Config: {config} has to be of type {self.config_class}")119 120 config.tie_word_embeddings = False121 122 # initialize with config123 PreTrainedModel.__init__(self, config)124 125 # Encoder:126 if encoder is None:127 encoder = MultiCvtWithProjectionHead(config=config.encoder)128 129 # Decoder:130 config.decoder._attn_implementation = 'eager'131 if decoder is None:132 decoder = transformers.BertLMHeadModel(config=config.decoder)133 134 self.encoder = encoder135 self.decoder = decoder136 137 if self.encoder.config.to_dict() != self.config.encoder.to_dict():138 logger.warning(139 f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"140 f" {self.config.encoder}"141 )142 if self.decoder.config.to_dict() != self.config.decoder.to_dict():143 logger.warning(144 f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"145 f" {self.config.decoder}"146 )147 148 self.encoder.config = self.config.encoder149 self.decoder.config = self.config.decoder150 151 # Load multi checkpoint:152 if encoder_decoder_ckpt_name:153 encoder_decoder = AutoModel.from_pretrained(encoder_decoder_ckpt_name, trust_remote_code=True)154 self.load_state_dict(encoder_decoder.state_dict())155 else:156 warnings.warn('The encoder-to-decoder model was not warm-started before applying low-rank approximation.')157 158 # Freeze the encoder:159 for p in self.encoder.parameters():160 p.requires_grad = False161 162 # Freeze the decoder and add LoRA:163 peft_config = LoraConfig(164 inference_mode=False, 165 r=8, 166 lora_alpha=32, 167 lora_dropout=0.1, 168 target_modules='bert.encoder.layer.[0-9]+.attention.self.(query|key)',169 )170 self.decoder = get_peft_model(self.decoder, peft_config)171 self.decoder.print_trainable_parameters()172 173 def forward(174 self,175 pixel_values: Optional[torch.FloatTensor] = None,176 decoder_input_ids: Optional[torch.LongTensor] = None,177 decoder_attention_mask: Optional[torch.BoolTensor] = None,178 encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,179 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,180 decoder_inputs_embeds: Optional[torch.FloatTensor] = None,181 labels: Optional[torch.LongTensor] = None,182 use_cache: Optional[bool] = None,183 output_attentions: Optional[bool] = None,184 output_hidden_states: Optional[bool] = None,185 return_dict: Optional[bool] = None,186 **kwargs,187 ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:188 189 return_dict = return_dict if return_dict is not None else self.config.use_return_dict190 191 kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}192 193 kwargs_decoder = {194 argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")195 }196 197 if encoder_outputs is None:198 if pixel_values is None:199 raise ValueError("You have to specify pixel_values")200 201 encoder_outputs = self.encoder(202 pixel_values,203 output_hidden_states=output_hidden_states,204 return_dict=return_dict,205 **kwargs_encoder,206 ) # CvT does not support output_attentions.207 elif isinstance(encoder_outputs, tuple):208 encoder_outputs = BaseModelOutput(*encoder_outputs)209 210 encoder_hidden_states = encoder_outputs[0]211 212 decoder_outputs = self.decoder(213 input_ids=decoder_input_ids,214 attention_mask=decoder_attention_mask,215 encoder_hidden_states=encoder_hidden_states,216 encoder_attention_mask=encoder_outputs.attention_mask,217 inputs_embeds=decoder_inputs_embeds,218 output_attentions=output_attentions,219 output_hidden_states=output_hidden_states,220 use_cache=use_cache,221 past_key_values=past_key_values,222 return_dict=return_dict,223 **kwargs_decoder,224 )225 226 # Loss:227 loss = None228 if labels is not None:229 logits = decoder_outputs.logits if return_dict else decoder_outputs[0]230 loss_fct = CrossEntropyLoss()231 loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.reshape(-1))232 233 if not return_dict:234 if loss is not None:235 return (loss,) + decoder_outputs + encoder_outputs236 else:237 return decoder_outputs + encoder_outputs238 239 return Seq2SeqLMOutput(240 loss=loss,241 logits=decoder_outputs.logits,242 past_key_values=decoder_outputs.past_key_values,243 decoder_hidden_states=decoder_outputs.hidden_states,244 decoder_attentions=decoder_outputs.attentions,245 cross_attentions=decoder_outputs.cross_attentions,246 encoder_last_hidden_state=encoder_outputs.last_hidden_state,247 # encoder_hidden_states=encoder_outputs.hidden_states,248 # encoder_attentions=encoder_outputs.attentions,249 )250 251 def prepare_inputs_for_generation(252 self,253 input_ids,254 special_token_ids,255 mask_token_id,256 past_key_values=None,257 attention_mask=None,258 use_cache=None,259 encoder_outputs=None,260 **kwargs,261 ):262 """263 Modification of: 264 https://github.com/huggingface/transformers/blob/main/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py#L660265 """266 267 # An update to generate() now prepends bos_token_id to each sequence if it does not exist at the start of the input: 268 # https://github.com/huggingface/transformers/blob/d533465150532b0c5de167b574e59f64c68b1154/src/transformers/generation/utils.py#L699C13-L699C30269 # Hence, we remove the prepended bos_token_id from each sequence if it is there:270 if torch.all(input_ids[:, 0] == 1):271 input_ids = input_ids[:, 1:]272 273 decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)274 decoder_attention_mask = (input_ids != mask_token_id).int()275 decoder_position_ids = torch.nn.functional.relu(276 torch.cumsum(decoder_attention_mask, dim=1, dtype=torch.int64) - 1277 )278 279 if not past_key_values:280 token_type_ids = self.token_ids_to_token_type_ids(input_ids, special_token_ids, [0, 1, 0, 1])281 else:282 token_type_ids = self.token_ids_to_token_type_ids_past(input_ids, special_token_ids, [0, 1, 0, 1])283 decoder_position_ids = decoder_position_ids[:, -1:]284 285 input_dict = {286 'attention_mask': attention_mask,287 'decoder_attention_mask': decoder_attention_mask,288 'decoder_input_ids': decoder_inputs['input_ids'],289 'decoder_token_type_ids': token_type_ids,290 'decoder_position_ids': decoder_position_ids,291 'encoder_outputs': encoder_outputs,292 'past_key_values': past_key_values,293 'use_cache': use_cache,294 }295 return input_dict296 297 def token_ids_to_token_type_ids(self, token_ids, special_token_ids, token_type_id_sections=None):298 """299 Extract token type identifiers from the token identifiers.300 301 Argument/s:302 token_ids - token identifiers.303 special_token_ids - special token identifiers that indicate the separation between sections.304 token_type_id_section - token type identifier for each section.305 306 Returns:307 token_type_ids - token type identifiers.308 """309 310 token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))311 312 mbatch_size, seq_len = token_ids.shape313 token_type_ids = torch.full_like(token_ids, token_type_id_sections[0], dtype=torch.long, device=token_ids.device)314 315 for i, j in enumerate(special_token_ids):316 # Find first occurrence of special tokens that indicate the boundary between sections:317 cols = (token_ids == j).int().argmax(dim=1)318 rows = torch.arange(mbatch_size, device=token_ids.device)319 320 # https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer.create_token_type_ids_from_sequences.example321 cols += 1322 323 # Ensure that the column index is not out of bounds. If 0, then token_id not present.324 # This is safe as index 0 is always a special token (now equal to 1 due to +1):325 rows = rows[torch.logical_and(cols != 1, cols < seq_len)]326 cols = cols[torch.logical_and(cols != 1, cols < seq_len)]327 328 # Indices to that correspond to the second sequence:329 if rows.nelement() != 0:330 ids = torch.stack([331 torch.stack([x, z]) for (x, y) in zip(rows, cols) for z in torch.arange(332 y, seq_len, device=token_ids.device,333 )334 ])335 336 token_type_ids[ids[:, 0], ids[:, 1]] = token_type_id_sections[i + 1]337 338 return token_type_ids339 340 def token_ids_to_token_type_ids_past(self, token_ids, special_token_ids, token_type_id_sections=None):341 """342 Extract token type identifiers from the token identifiers if past != None.343 344 Argument/s:345 token_ids - token identifiers.346 special_token_ids - special token identifiers that indicate the separation between sections.347 348 Returns:349 token_type_ids - token type identifiers.350 """351 352 token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))353 token_type_ids = torch.full([token_ids.shape[0], 1], token_type_id_sections[0], dtype=torch.long, device=token_ids.device)354 355 # https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer.create_token_type_ids_from_sequences.example356 token_ids = token_ids[:, :-1]357 358 for i, j in enumerate(special_token_ids):359 360 # Find first occurrence of special token, which indicates the boundary between sections:361 exists = torch.any(token_ids == j, dim=1, keepdim=True)362 token_type_ids[exists] = token_type_id_sections[i + 1]363 364 return token_type_ids365 366 def tokenize_report_teacher_forcing(self, findings: str, impression: str, tokenizer: PreTrainedTokenizerFast, max_len: int):367 """368 Tokenize the reports and creates the inputs and targets for teacher forcing.369 370 Argument/s:371 findings - findings section.372 impression - impression section.373 return_token_type_ids - return the token type identifiers.374 tokenizer - Hugging Face tokenizer.375 max_len - maximum number of tokens.376 377 Returns:378 decoder_input_ids - the token identifiers for the input of the decoder.379 decoder_attention_mask - the attention mask for the decoder_input_ids.380 label_ids - the label token identifiers for the decoder.381 """382 383 # Prepare the sections for the tokenizer by placing special tokens between each section:384 report = [f'{tokenizer.bos_token}{i}{tokenizer.sep_token}{j}{tokenizer.eos_token}' for i, j in385 zip(findings, impression)]386 387 # Tokenize the report:388 tokenized = tokenizer(389 report,390 padding='longest',391 truncation=True,392 max_length=max_len + 1, # +1 to account for the bias between input and target.393 return_tensors='pt',394 return_token_type_ids=False,395 add_special_tokens=False,396 ).to(self.device)397 398 # Modify for language modelling:399 batch_dict = {400 401 # Labels for the decoder (shifted right by one for autoregression):402 'label_ids': tokenized['input_ids'][:, 1:].detach().clone(),403 404 # Remove last token identifier to match the sequence length of the labels:405 'decoder_input_ids': tokenized['input_ids'][:, :-1],406 407 # Attention mask for the decoder_input_ids (remove first token so that the eos_token_id is not considered):408 'decoder_attention_mask': tokenized['attention_mask'][:, 1:],409 }410 411 return batch_dict412 413 def split_and_decode_sections(self, token_ids, special_token_ids, tokenizer: PreTrainedTokenizerFast):414 """415 Split the token identifiers into sections, then convert the token identifiers into strings.416 417 Argument/s:418 token_ids - token identifiers.419 special_token_ids - special token identifiers that indicate the end of each section.420 tokenizer - Hugging Face tokenizer.421 422 Returns:423 token_type_ids - token type identifiers.424 """425 426 _, seq_len = token_ids.shape427 428 # The number of sections is the same as the number of special_token_ids:429 num_sections = len(special_token_ids)430 431 sections = {k: [] for k in range(num_sections)}432 433 for i in token_ids:434 prev_col = 0435 for j, k in enumerate(special_token_ids):436 437 # The maximum sequence length was exceeded, thus no more tokens:438 if prev_col >= seq_len:439 sections[j].append('')440 continue441 442 # Find first occurrence of special tokens that indicate the boundary between sections:443 col = (i == k).int().argmax().item()444 445 # If equal to 0, token was not found, set the column to the sequence length (as the decoder exceeded446 # the maximum sequence length):447 if col == 0:448 col = seq_len449 450 # Extract section token identifiers:451 section_token_ids = i[prev_col:col]452 prev_col = col453 section_string = tokenizer.decode(section_token_ids, skip_special_tokens=True)454 455 sections[j].append(section_string)456 457 return tuple(sections.values())458 459 def tokenize_prompt(460 self, 461 previous_findings: str, 462 previous_impression: str, 463 tokenizer: PreTrainedTokenizerFast, 464 max_len: int,465 add_bos_token_id: bool = False,466 ):467 """468 Tokenize the sections of the previous report to be used as a prompt.469 470 Argument/s:471 previous_findings - previous findings section.472 previous_impression - previous impression section.473 tokenizer - Hugging Face tokenizer.474 max_len - maximum number of tokens.475 add_bos_token_id - whether to add the BOS token identifier to the prompt.476 477 Returns:478 input_ids - the input identifiers for the previous impression.479 attention_mask - the attention mask for the previous impression480 """481 482 # Use [NPF]/[NPI] special token if no previous findings/impression:483 previous_findings = ['[NPF]' if not i else i for i in previous_findings]484 previous_impression = ['[NPI]' if not i else i for i in previous_impression]485 486 # Prepare the sections for the tokenizer by placing special tokens:487 previous_sections = [488 f'[PMT]{i}[PMT-SEP]{j}{tokenizer.bos_token}' if add_bos_token_id else f'[PMT]{i}[PMT-SEP]{j}' \489 for i, j in zip(previous_findings, previous_impression)490 ]491 492 # Tokenize:493 previous_sections = tokenizer(494 previous_sections,495 padding='longest',496 truncation=True,497 max_length=max_len,498 return_tensors='pt',499 return_token_type_ids=False,500 add_special_tokens=False,501 ).to(self.device)502 503 # Ensure BOS token identifier is at the end of the input_ids:504 if previous_sections.input_ids.shape[1] == max_len:505 previous_sections.input_ids[:, -1] = torch.where(506 previous_sections.attention_mask[:, -1] == 1,507 tokenizer.bos_token_id,508 previous_sections.input_ids[:, -1],509 ) 510 511 assert previous_sections.input_ids.shape[1] <= max_len512 513 return {'input_ids': previous_sections.input_ids, 'attention_mask': previous_sections.attention_mask}514