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