File size: 1,604 Bytes
c2ede1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """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])
|