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