modeling_chest2vec.py

"""Chest2Vec — Qwen3-Embedding model (contrastive LoRA merged in) for chest radiology reports.

Self-contained: load with `AutoModel` — no `chest2vec` package, and no download of the base
Qwen3-Embedding weights (the merged encoder ships in this repo).

    from transformers import AutoModel, AutoTokenizer
    model = AutoModel.from_pretrained("chest2vec/chest2vec_0.6B", trust_remote_code=True).eval()
    tok   = AutoTokenizer.from_pretrained("chest2vec/chest2vec_0.6B", trust_remote_code=True)
    emb   = model.embed_texts(["Frontal chest radiograph. No pneumothorax."], tokenizer=tok)  # [N,H], L2-normalized

Embedding = left-padding-aware last-token (EOS) pooling + L2-norm. Matryoshka: pass `dim=512`/`256`.
"""
from typing import List, Optional
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel, AutoConfig, AutoModel
from transformers.modeling_outputs import BaseModelOutputWithPooling
from .configuration_chest2vec import Chest2VecConfig


def build_qwen_query(instruction: str, query: str) -> str:
    instruction = str(instruction).strip()
    return f"Instruct: {instruction}\nQuery: {str(query).strip()}" if instruction else str(query).strip()


def _build_encoder(encoder_config: dict, attn_implementation: str = "sdpa"):
    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)
    cfg.torch_dtype = "float32"
    try:
        return AutoModel.from_config(cfg, attn_implementation=attn_implementation)
    except TypeError:
        return AutoModel.from_config(cfg)


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


class Chest2VecModel(PreTrainedModel):
    config_class = Chest2VecConfig
    base_model_prefix = "model"

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

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    @property
    def device(self):
        return next(self.parameters()).device

    # ---- low-level encoder forward (token tensors -> pooled, L2-normalized embedding) ----
    def encode(self, input_ids, attention_mask, position_ids=None, normalize=True):
        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_ids(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_formatted(self, texts, tokenizer, max_len, batch_size, return_cpu, dim):
        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.default_max_len
        dev = self.device
        self.eval()
        out = []
        for i in range(0, len(texts), batch_size):
            ii, am = self._encode_ids(tok, texts[i:i + batch_size], max_len)
            emb = self.encode(ii.to(dev), am.to(dev), normalize=False).pooler_output
            if dim is not None:
                emb = emb[:, :dim]
            emb = F.normalize(emb, p=2, dim=-1)
            out.append(emb.cpu() if return_cpu else emb)
        return torch.cat(out, dim=0)

    # ---- public API ----
    def embed_texts(self, texts, *, tokenizer=None, max_len: Optional[int] = None,
                    batch_size: int = 16, return_cpu: bool = True, dim: Optional[int] = None):
        """Embed reports/documents (no instruction). Returns [N, dim] L2-normalized."""
        return self._embed_formatted(texts, tokenizer, max_len, batch_size, return_cpu, dim)

    def embed_instruction_query(self, instructions, queries, *, tokenizer=None,
                                max_len: Optional[int] = None, batch_size: int = 16,
                                return_cpu: bool = True, dim: Optional[int] = None):
        """Embed instruction-conditioned queries. `instructions` may be one string or a list."""
        if isinstance(queries, str):
            queries = [queries]
        if isinstance(instructions, str):
            instructions = [instructions] * len(queries)
        texts = [build_qwen_query(i, q) for i, q in zip(instructions, queries)]
        return self._embed_formatted(texts, tokenizer, max_len, batch_size, return_cpu, dim)

    def embed(self, texts, *, instruction: Optional[str] = None, tokenizer=None,
              max_len: Optional[int] = None, batch_size: int = 16, return_cpu: bool = True,
              dim: Optional[int] = None):
        """Convenience: with `instruction`, embed as instruction-conditioned queries; else plain."""
        if instruction:
            return self.embed_instruction_query(instruction, texts, tokenizer=tokenizer,
                                                max_len=max_len, batch_size=batch_size,
                                                return_cpu=return_cpu, dim=dim)
        return self.embed_texts(texts, tokenizer=tokenizer, max_len=max_len,
                                batch_size=batch_size, return_cpu=return_cpu, dim=dim)

    def forward(self, texts=None, *, input_ids=None, attention_mask=None, position_ids=None,
                normalize=True, **kwargs):
        if input_ids is not None:
            return self.encode(input_ids, attention_mask, position_ids, normalize=normalize)
        if texts is not None:
            return BaseModelOutputWithPooling(pooler_output=self.embed_texts(texts, return_cpu=False))
        raise ValueError("Provide either `texts` or (`input_ids`, `attention_mask`).")

    @staticmethod
    def cosine_topk(query_emb, cand_emb, k=10):
        """Top-k most similar candidates per query (embeddings assumed L2-normalized)."""
        sims = query_emb @ cand_emb.T
        k = min(k, cand_emb.shape[0])
        vals, idx = torch.topk(sims, k, dim=-1)
        return vals, idx