"""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])