Upload 5 files

config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": [
+    "CheXficientModel"
+  ],
+  "image_size": 378,
+  "model_type": "chexficient_clip",
+  "projection_dim": 512,
+  "text_model_name": "emilyalsentzer/Bio_ClinicalBERT",
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "vision_model_name": "dinov2_vitb14",
+  "auto_map": {
+  "AutoConfig": "configuration_chexficient.CheXficientConfig",
+  "AutoModel": "modeling_chexficient.CheXficientModel"
+  }
+}

configuration_chexficient.py

@@ -0,0 +1,18 @@
+from transformers import PretrainedConfig
+
+class CheXficientConfig(PretrainedConfig):
+    model_type = "chexficient_clip"
+    def __init__(
+        self,
+        vision_model_name="dinov2_vitb14",
+        text_model_name="emilyalsentzer/Bio_ClinicalBERT",
+        projection_dim=512,
+        image_size=378,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.vision_model_name = vision_model_name
+        self.text_model_name = text_model_name
+        self.projection_dim = projection_dim
+        self.image_size = image_size

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f691262c3f77c3c850bebff420180602cf4ca5d5214449e377366f3205548336
+size 780793036

modeling_chexficient.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import (
+    PreTrainedModel,
+    AutoTokenizer,
+    AutoModel
+)
+
+from dinov2.models.vision_transformer import vit_base
+from projection import load_projection_head
+from configuration_chexficient import CheXficientConfig
+
+
+URL_DICT = {
+    "dinov2_vits14": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_reg4_pretrain.pth",
+    "dinov2_vitb14": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_reg4_pretrain.pth",
+    "dinov2_vitl14": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth",
+}
+
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, model_name='emilyalsentzer/Bio_ClinicalBERT'):
+        super().__init__()
+        # self.model = AutoModel.from_pretrained(model_name, ignore_mismatched_sizes=False, cache_dir='./huggingface',)
+        # self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='./huggingface/tokenizers')
+        self.model = AutoModel.from_pretrained(model_name, use_safetensors=True, ignore_mismatched_sizes=False, )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, )
+        if self.tokenizer.bos_token_id is None:
+            self.tokenizer.bos_token_id = self.tokenizer.cls_token_id
+        self.out_dim = self.model.config.hidden_size
+
+    def forward(self, inputs):
+        outputs = self.model(**inputs)
+        return outputs["last_hidden_state"]  # (batch, seq_len, hidden_size)
+
+
+class ImageEncoder(nn.Module):
+    def __init__(self, model_name='dinov2_vitb14', image_size=224):
+        super().__init__()
+        self.model = vit_base(patch_size=14, img_size=image_size, init_values=1.0, block_chunks=0)
+        stact_dict = torch.hub.load_state_dict_from_url(URL_DICT[model_name], map_location="cpu")
+        ##########################################################
+        if self.model.pos_embed.shape[1] != stact_dict['pos_embed'].shape[1]:
+            cls_pos_embed = stact_dict['pos_embed'][:, 0:1, :]  # [1, hidden_dim]
+            patch_pos_embed = stact_dict['pos_embed'][:, 1:, :]  # [1369, hidden_dim]
+            # raw patch grid size
+            orig_size = int(patch_pos_embed.shape[1] ** 0.5)  # 37
+            new_size = image_size // self.model.patch_size  # 512 // 16 = 32
+            patch_pos_embed = patch_pos_embed.reshape(1, orig_size, orig_size, -1).permute(0, 3, 1, 2)  # [1, dim, 37, 37]
+            patch_pos_embed = F.interpolate(patch_pos_embed, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).reshape(1, new_size * new_size, -1)
+            stact_dict['pos_embed'] = torch.cat((cls_pos_embed, patch_pos_embed), dim=1)  # [1, 1+new_size*new_size, dim]
+        ##########################################################
+        res = self.model.load_state_dict(stact_dict, strict=False)
+        print('load dinov2 pretrained model:', res)
+        self.out_dim = self.model.embed_dim
+
+    def forward(self, x):
+        feats = self.model(x)  # Shape: (b, d)
+        return feats
+
+
+
+class CheXficientModel(PreTrainedModel):
+    config_class = CheXficientConfig
+    base_model_prefix = "chexficient"
+
+    def __init__(self, config: CheXficientConfig):
+        super().__init__(config)
+
+        # ===== Encoders =====
+        # self.image_encoder = AutoModel.from_pretrained(
+        #     config.vision_model_name,
+        #     use_safetensors=True
+        # )
+        # self.text_encoder = AutoModel.from_pretrained(
+        #     config.text_model_name,
+        #     use_safetensors=True
+        # )
+        self.image_encoder = ImageEncoder(model_name=config.vision_model_name, image_size=config.image_size)
+        self.text_encoder = TextEncoder(model_name=config.text_model_name)
+
+        # ===== Projection heads =====
+        self.image_projection = load_projection_head(
+            embedding_dim=self.image_encoder.out_dim,
+            config_projection_head={'name': 'linear', 'dropout': 0.1, 'proj_dim': config.projection_dim}
+        )
+        self.text_projection = load_projection_head(
+            embedding_dim=self.text_encoder.out_dim,
+            config_projection_head={'name': 'linear', 'dropout': 0.1, 'proj_dim': config.projection_dim}
+        )
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * 0.01)
+
+        self.post_init()
+
+    def get_image_features(self, pixel_values):
+        vision_outputs = self.image_encoder(pixel_values=pixel_values)
+        pooled = vision_outputs.last_hidden_state[:, 0]
+        projected = self.image_projection(pooled)
+        return F.normalize(projected, dim=-1)
+
+    def get_text_features(self, input_ids, attention_mask):
+        text_outputs = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+        pooled = text_outputs.last_hidden_state[:, 0]
+        projected = self.text_projection(pooled)
+        return F.normalize(projected, dim=-1)
+
+    def forward(
+        self,
+        pixel_values=None,
+        input_ids=None,
+        attention_mask=None,
+        return_loss=False
+    ):
+        image_features = self.get_image_features(pixel_values)
+        text_features = self.get_text_features(input_ids, attention_mask)
+
+        logit_scale = self.logit_scale.exp()
+
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        loss = None
+        if return_loss:
+            labels = torch.arange(len(logits_per_image)).to(logits_per_image.device)
+            loss_i = F.cross_entropy(logits_per_image, labels)
+            loss_t = F.cross_entropy(logits_per_text, labels)
+            loss = (loss_i + loss_t) / 2
+
+        return {
+            "loss": loss,
+            "logits_per_image": logits_per_image,
+            "logits_per_text": logits_per_text,
+            "image_embeds": image_features,
+            "text_embeds": text_features,
+        }
+

projection.py

@@ -0,0 +1,42 @@
+from torch import nn
+from typing import Dict
+
+
+class MLPProjectionHead(nn.Module):
+    def __init__(self, embedding_dim, projection_dim, dropout):
+        super().__init__()
+        self.projection = nn.Linear(embedding_dim, projection_dim)
+        self.gelu = nn.GELU()
+        self.fc = nn.Linear(projection_dim, projection_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(projection_dim)
+
+    def forward(self, x):
+        projected = self.projection(x)
+        x = self.gelu(projected)
+        x = self.fc(x)
+        x = self.dropout(x)
+        x = x + projected
+        x = self.layer_norm(x)
+        return x
+
+
+class LinearProjectionHead(nn.Module):
+    def __init__(self, embedding_dim, projection_dim):
+        super().__init__()
+        self.projection = nn.Linear(embedding_dim, projection_dim)
+
+    def forward(self, x):
+        return self.projection(x)
+
+
+def load_projection_head(embedding_dim: int, config_projection_head: Dict):
+    if config_projection_head["name"].lower() == "mlp":
+        projection_head = MLPProjectionHead(
+            embedding_dim=embedding_dim, projection_dim=config_projection_head["proj_dim"], dropout=config_projection_head["dropout"]
+        )
+    elif config_projection_head["name"].lower() == "linear":
+        projection_head = LinearProjectionHead(embedding_dim=embedding_dim, projection_dim=config_projection_head["proj_dim"])
+    else:
+        raise KeyError(f"Not supported text encoder: {config_projection_head}")
+    return projection_head