| """BiomedCLIP text tower without open_clip (not installed in this env): BiomedBERT + 2-layer MLP projection, CLS pooling, |
| 256 tokens, attention mask = token != pad (= open_clip HFTextEncoder with hf_proj_type mlp, cls_last_hidden_state_pooler). |
| Weights: text.* of microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224/open_clip_pytorch_model.bin, strict load.""" |
| import torch |
| import torch.nn as nn |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoConfig, AutoTokenizer, BertModel |
|
|
| REPO = "microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224" |
| BERT = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract" |
|
|
|
|
| class BiomedCLIPText(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.transformer = BertModel(AutoConfig.from_pretrained(BERT), add_pooling_layer=False) |
| self.proj = nn.Sequential(nn.Linear(768, 640, bias=False), nn.GELU(), nn.Linear(640, 512, bias=False)) |
| sd = torch.load(hf_hub_download(REPO, "open_clip_pytorch_model.bin"), map_location="cpu", weights_only=True) |
| sd = {k[5:]: v for k, v in sd.items() if k.startswith("text.") and not k.endswith("position_ids")} |
| self.load_state_dict(sd, strict=True) |
| self.tok = AutoTokenizer.from_pretrained(BERT) |
|
|
| def forward(self, texts): |
| dev = next(self.parameters()).device |
| ids = self.tok(texts, return_tensors="pt", max_length=256, padding="max_length", truncation=True).input_ids.to(dev) |
| h = self.transformer(input_ids=ids, attention_mask=(ids != self.tok.pad_token_id).long()).last_hidden_state |
| return self.proj(h[:, 0]) |
|
|
