modeling_chest2vec_embedding.py

"""
chest2vec — chest-radiology text embedding model (HuggingFace `AutoModel` wrapper).

A `Qwen3-Embedding` encoder LoRA-adapted for chest CT/CXR report retrieval, with the LoRA
**merged into the weights** so the repo is fully self-contained: loading needs neither the
`chest2vec` package nor a download of the base Qwen3-Embedding weights.

Embedding = left-padding-aware last-token (EOS) pooling of the final hidden state, L2-normalized.

Usage:

    from transformers import AutoModel, AutoTokenizer
    model = AutoModel.from_pretrained("lukeingawesome/chest2vec_0.6b_chest", trust_remote_code=True).eval()
    tok   = AutoTokenizer.from_pretrained("lukeingawesome/chest2vec_0.6b_chest", trust_remote_code=True)

    docs = ["Bibasilar atelectasis with small bilateral pleural effusions."]
    emb  = model.embed(docs, tokenizer=tok)                 # [N, H] float32, L2-normalized
    # instruction-conditioned queries (Qwen3-Embedding convention):
    q    = model.embed(["pleural effusion"], tokenizer=tok,
                       instruction="Retrieve chest CT reports relevant to the query")
"""
from typing import List, Optional
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel
from transformers.modeling_outputs import BaseModelOutputWithPooling


class Chest2VecEmbeddingConfig(PretrainedConfig):
    model_type = "chest2vec_embedding"

    def __init__(self, encoder_config: Optional[dict] = None,
                 base_model: str = "Qwen/Qwen3-Embedding-0.6B",
                 hidden_size: int = 1024, max_len: int = 512,
                 pooling: str = "last_token", matryoshka_dims: Optional[list] = None, **kwargs):
        super().__init__(**kwargs)
        self.encoder_config = encoder_config or {}
        self.base_model = base_model
        self.hidden_size = hidden_size
        self.max_len = max_len
        self.pooling = pooling
        self.matryoshka_dims = matryoshka_dims or []


def _build_encoder(encoder_config: dict, attn_implementation: str = "sdpa", dtype=None):
    ecfg = dict(encoder_config)
    for k in ("architectures", "auto_map", "transformers_version", "_name_or_path", "torch_dtype"):
        ecfg.pop(k, None)
    model_type = ecfg.pop("model_type", "qwen3")
    cfg = AutoConfig.for_model(model_type, **ecfg)
    try:
        return AutoModel.from_config(cfg, attn_implementation=attn_implementation)
    except TypeError:
        return AutoModel.from_config(cfg)


def _last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    idx = attention_mask.sum(dim=1) - 1
    return last_hidden_states[torch.arange(last_hidden_states.size(0), device=last_hidden_states.device), idx]


class Chest2VecEmbeddingModel(PreTrainedModel):
    config_class = Chest2VecEmbeddingConfig
    base_model_prefix = "model"

    def __init__(self, config: Chest2VecEmbeddingConfig):
        super().__init__(config)
        self.model = _build_encoder(config.encoder_config, getattr(config, "attn_implementation", "sdpa"))
        self._tokenizer = None
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, position_ids=None, normalize=True, **kwargs):
        if position_ids is None and attention_mask is not None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 0)
        out = self.model(input_ids=input_ids, attention_mask=attention_mask,
                         position_ids=position_ids, use_cache=False, return_dict=True)
        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out.hidden_states[-1]
        emb = _last_token_pool(h, attention_mask).float()
        if normalize:
            emb = F.normalize(emb, p=2, dim=-1)
        return BaseModelOutputWithPooling(last_hidden_state=h, pooler_output=emb)

    def _get_tokenizer(self, tokenizer=None):
        if tokenizer is not None:
            return tokenizer
        if self._tokenizer is None:
            from transformers import AutoTokenizer
            src = self.config._name_or_path or self.config.base_model
            self._tokenizer = AutoTokenizer.from_pretrained(src, padding_side="left", trust_remote_code=True)
            if self._tokenizer.pad_token_id is None:
                self._tokenizer.pad_token = self._tokenizer.eos_token
        return self._tokenizer

    def _encode(self, tok, texts: List[str], max_len: int):
        pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
        eod_id = tok.convert_tokens_to_ids("<|endoftext|>")
        if eod_id is None or eod_id < 0:
            eod_id = pad_id
        enc = tok([str(t) for t in texts], add_special_tokens=False, truncation=True,
                  max_length=max_len - 1, padding=False, return_attention_mask=False)
        ids = [x + [eod_id] for x in enc["input_ids"]]
        T = max((len(x) for x in ids), default=1)
        input_ids = [[pad_id] * (T - len(x)) + x for x in ids]
        attn = [[0] * (T - len(x)) + [1] * len(x) for x in ids]
        return torch.tensor(input_ids, dtype=torch.long), torch.tensor(attn, dtype=torch.long)

    @torch.no_grad()
    def embed(self, texts, tokenizer=None, instruction: Optional[str] = None, batch_size: int = 16,
              max_len: Optional[int] = None, device=None, normalize: bool = True,
              dim: Optional[int] = None) -> torch.Tensor:
        """Embed a list of texts -> [N, dim] L2-normalized.

        If `instruction` is given, each text is formatted as `Instruct: {instruction}\\nQuery: {text}`
        (Qwen3-Embedding query convention) — apply it to queries, embed the corpus without it.

        `dim` enables **Matryoshka** truncation: the first `dim` dimensions are kept and
        re-normalized. This model was MRL-trained, so dim in {256, 512, full} retains quality."""
        if isinstance(texts, str):
            texts = [texts]
        if dim is not None and dim > self.config.hidden_size:
            raise ValueError(f"dim {dim} > embedding dim {self.config.hidden_size}")
        tok = self._get_tokenizer(tokenizer)
        max_len = max_len or self.config.max_len
        device = device or next(self.parameters()).device
        self.eval()
        if instruction:
            instruction = str(instruction).strip()
            texts = [f"Instruct: {instruction}\nQuery: {str(t).strip()}" for t in texts]
        out = []
        for i in range(0, len(texts), batch_size):
            ii, am = self._encode(tok, texts[i:i + batch_size], max_len)
            emb = self(input_ids=ii.to(device), attention_mask=am.to(device), normalize=False).pooler_output
            if dim is not None:
                emb = emb[:, :dim]
            if normalize:
                emb = F.normalize(emb, p=2, dim=-1)
            out.append(emb.cpu())
        return torch.cat(out, dim=0)