File size: 8,319 Bytes
f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef c036088 f115fef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """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
|