auto model bug fix

modeling_radzero.py

@@ -1,8 +1,7 @@
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import AutoTokenizer, BertModel
+from transformers import BertModel
 from transformers.models.clip.modeling_clip import CLIPTextModel
 from transformers.models.mpnet.modeling_mpnet import MPNetModel
 from transformers.trainer import logger
@@ -11,8 +10,8 @@ from .align_transformers import build_align_transformer
 from .common_layers import BasePreTrainedModel
 from .configuration_radzero import CxrAlignConfig
 from .losses import KeyPhraseAlignmentLoss
-from .text_encoders import aggregate_tokens, build_text_encoder
-from .vision_encoders import MRM, Dinov2Model, build_vision_encoder
+from .text_encoders import build_text_encoder
+from .vision_encoders import Dinov2Model, build_vision_encoder
 
 
 class CxrAlignModel(BasePreTrainedModel):
@@ -28,13 +27,6 @@ class CxrAlignModel(BasePreTrainedModel):
     def build_text_model(self, config: CxrAlignConfig):
         text_config = config.text_config
         text_model = build_text_encoder(text_config)
-
-        if text_config.model_type == "bioclinicalmpbert":
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                text_config.pretrained_tokenizer_name_or_path
-            )
-            self.idxtoword = {v: k for k, v in self.tokenizer.get_vocab().items()}
-
         return text_model
 
     def build_align_transformer_model(self, config: CxrAlignConfig):
@@ -94,13 +86,7 @@ class CxrAlignModel(BasePreTrainedModel):
 
         if isinstance(self.vision_model, Dinov2Model):
             vision_tokens = self.vision_model(pixel_values)["last_hidden_state"]
-        elif isinstance(self.vision_model, MRM):
-            img_emb_g, img_emb_l = self.vision_model(pixel_values)
-            img_emb_g = img_emb_g.unsqueeze(1)
-            img_emb_l = img_emb_l.view(img_emb_l.size(0), img_emb_l.size(1), -1)
-            img_emb_l = img_emb_l.permute(0, 2, 1)
 
-            vision_tokens = torch.cat([img_emb_g, img_emb_l], dim=1)
         else:
             raise NotImplementedError
 
@@ -152,53 +138,6 @@ class CxrAlignModel(BasePreTrainedModel):
                     token_embeddings * input_mask_expanded, 1
                 ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
-        elif isinstance(self.text_model, BertModel):
-            # BioClinicalMPBERT
-
-            model_output = self.text_model(
-                input_ids=encoded_input["input_ids"],
-                attention_mask=encoded_input["attention_mask"],
-                token_type_ids=encoded_input.get("token_type_ids", None),
-            )
-
-            if self.config.text_config.use_cls_token:
-                text_features = model_output.last_hidden_state[:, 0, :]
-
-            elif self.config.text_config.use_aggregate_tokens:
-
-                all_embeddings = model_output[2]
-                embeddings = torch.stack(
-                    all_embeddings[-self.config.text_config.last_n_layers :]
-                )
-                embeddings = embeddings.permute(1, 0, 2, 3)
-
-                embeddings, sents = aggregate_tokens(
-                    embeddings, encoded_input["input_ids"], self.idxtoword
-                )
-                sent_embeddings = embeddings.mean(axis=2)
-
-                if self.config.text_config.aggregate_method == "sum":
-                    word_embeddings = embeddings.sum(axis=1)
-                    sent_embeddings = sent_embeddings.sum(axis=1)
-                elif self.config.text_config.aggregate_method == "mean":
-                    word_embeddings = embeddings.mean(axis=1)
-                    sent_embeddings = sent_embeddings.mean(axis=1)
-
-                word_embeddings = word_embeddings.permute(0, 2, 1)
-
-                text_features = sent_embeddings
-                text_outputs["word_embeddings"] = word_embeddings
-
-            else:
-                text_features = model_output.last_hidden_state
-                mask = encoded_input["attention_mask"].unsqueeze(-1).float()
-                text_features = torch.sum(text_features * mask, dim=1) / torch.clamp(
-                    mask.sum(dim=1), min=1e-9
-                )
-
-            if self.text_projector is not None:
-                text_features = self.text_projector(text_features)
-
         else:
             raise NotImplementedError