""" Chest2Vec CT Report Labeler — HuggingFace `AutoModel` wrapper. A weakly-supervised multi-label classifier that maps a free-text chest-CT report to a 137-leaf chest-imaging taxonomy with a ternary status per label (negative / uncertain / positive). Architecture: `Qwen/Qwen3-Embedding-0.6B` encoder (LoRA merged in) → left-padding-aware last-token (EOS) pooling → L2-normalization → a single linear ternary head (`hidden=1024 → 137 × 3`). Usage: from transformers import AutoModel, AutoTokenizer model = AutoModel.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True).eval() tok = AutoTokenizer.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True) reports = ["Bibasilar atelectasis with small bilateral pleural effusions. Cardiomegaly."] print(model.label_reports(reports, tokenizer=tok)) # -> [{'Pleural effusion': 'positive', ...}] # CheXbert / SRR-BERT-style report comparison (label both, compare): res = model.score_reports(gt_reports, pred_reports, tokenizer=tok) print(res["micro"]["f1"], res["macro"]["f1"], res["weighted"]["f1"]) """ from typing import Dict, List, Optional, Any import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel from transformers.modeling_outputs import ModelOutput from dataclasses import dataclass # class index ordering produced by the head's softmax (axis=-1) NEGATIVE, UNCERTAIN, POSITIVE = 0, 1, 2 _CLASS_TO_VALUE = {NEGATIVE: 0, UNCERTAIN: -1, POSITIVE: 1} _CLASS_TO_NAME = {NEGATIVE: "negative", UNCERTAIN: "uncertain", POSITIVE: "positive"} class Chest2VecLabelerConfig(PretrainedConfig): model_type = "chest2vec_labeler" def __init__( self, encoder_config: Optional[dict] = None, base_model: str = "Qwen/Qwen3-Embedding-0.6B", hidden_size: int = 1024, n_labels: int = 137, num_classes_per_label: int = 3, labels: Optional[List[str]] = None, instruction: str = "Given the following chest CT report, extract the presence/absence of entities", max_len: int = 512, default_threshold: float = 0.5, **kwargs, ): super().__init__(**kwargs) self.encoder_config = encoder_config or {} self.base_model = base_model self.hidden_size = hidden_size self.n_labels = n_labels self.num_classes_per_label = num_classes_per_label self.labels = labels or [] self.instruction = instruction self.max_len = max_len self.default_threshold = default_threshold @dataclass class LabelerOutput(ModelOutput): logits: torch.FloatTensor = None # [B, num_labels, 3] embedding: torch.FloatTensor = None # [B, hidden] L2-normalized pooled 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: cfg._attn_implementation = attn_implementation except Exception: pass 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-aware last-token (EOS) pooling — matches the training pipeline.""" 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 Chest2VecLabelerModel(PreTrainedModel): config_class = Chest2VecLabelerConfig base_model_prefix = "model" def __init__(self, config: Chest2VecLabelerConfig): super().__init__(config) self.model = _build_encoder(config.encoder_config, getattr(config, "attn_implementation", "sdpa")) self.head = nn.Linear(config.hidden_size, config.n_labels * config.num_classes_per_label) self.num_labels = config.n_labels self.num_classes_per_label = config.num_classes_per_label self._tokenizer = None self.post_init() # ---- core forward (token tensors in, logits out) ---- def forward(self, input_ids=None, attention_mask=None, position_ids=None, **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) emb = F.normalize(emb.float(), p=2, dim=-1) logits = self.head(emb).view(emb.size(0), self.num_labels, self.num_classes_per_label) return LabelerOutput(logits=logits, embedding=emb) # ---- tokenization (matches training: Instruct/Query + reserved EOS + left pad) ---- 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, reports: List[str], max_len: int): instr = self.config.instruction.strip() texts = [(f"Instruct: {instr}\nQuery: {str(r).strip()}" if instr else str(r).strip()) for r in reports] 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(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)) # ---- high-level prediction API ---- @torch.no_grad() def predict_proba(self, reports: List[str], tokenizer=None, batch_size: int = 16, max_len: Optional[int] = None, device=None) -> torch.Tensor: """Return [N, num_labels] probability of the POSITIVE class for each label.""" if isinstance(reports, str): reports = [reports] tok = self._get_tokenizer(tokenizer) max_len = max_len or self.config.max_len device = device or next(self.parameters()).device self.eval() out = [] for i in range(0, len(reports), batch_size): ii, am = self._encode(tok, reports[i:i + batch_size], max_len) logits = self(input_ids=ii.to(device), attention_mask=am.to(device)).logits out.append(torch.softmax(logits.float(), dim=-1)[:, :, POSITIVE].cpu()) return torch.cat(out, dim=0) @torch.no_grad() def predict(self, reports: List[str], tokenizer=None, threshold: Optional[float] = None, batch_size: int = 16, max_len: Optional[int] = None, device=None, return_ternary: bool = False) -> Dict[str, Any]: """Return {'labels': names, 'positive': [N,L] 0/1, 'proba': [N,L], ('ternary': [N,L] in {-1,0,1})}.""" if isinstance(reports, str): reports = [reports] thr = self.config.default_threshold if threshold is None else threshold tok = self._get_tokenizer(tokenizer) max_len = max_len or self.config.max_len device = device or next(self.parameters()).device self.eval() proba, ternary = [], [] for i in range(0, len(reports), batch_size): ii, am = self._encode(tok, reports[i:i + batch_size], max_len) logits = self(input_ids=ii.to(device), attention_mask=am.to(device)).logits.float().cpu() proba.append(torch.softmax(logits, dim=-1)[:, :, POSITIVE]) if return_ternary: cls = logits.argmax(-1) ternary.append(torch.tensor([[_CLASS_TO_VALUE[int(c)] for c in row] for row in cls])) proba = torch.cat(proba, dim=0) res = {"labels": list(self.config.labels), "proba": proba.numpy(), "positive": (proba >= thr).int().numpy(), "threshold": thr} if return_ternary: res["ternary"] = torch.cat(ternary, dim=0).numpy() return res def label_reports(self, reports: List[str], tokenizer=None, threshold: Optional[float] = None, **kw) -> List[Dict[str, str]]: """Return, per report, a dict {label_name: 'positive'} for labels above threshold.""" out = self.predict(reports, tokenizer=tokenizer, threshold=threshold, **kw) names = out["labels"] return [{names[j]: "positive" for j in range(len(names)) if row[j]} for row in out["positive"]] # ---- CheXbert / SRR-BERT-style report-comparison F1 ---- @torch.no_grad() def score_reports(self, gt_reports: List[str], pred_reports: List[str], tokenizer=None, threshold: Optional[float] = None, batch_size: int = 16, max_len: Optional[int] = None, device=None) -> Dict[str, Any]: """ Label both GT and predicted reports, then compute label-level F1 (CheXbert-style). Treats the labels extracted from `gt_reports` as ground truth and those from `pred_reports` as the prediction. Returns micro / macro / weighted P,R,F1 over the 137 labels, plus per-label scores. """ from sklearn.metrics import precision_recall_fscore_support if len(gt_reports) != len(pred_reports): raise ValueError("gt_reports and pred_reports must have the same length") kw = dict(tokenizer=tokenizer, threshold=threshold, batch_size=batch_size, max_len=max_len, device=device) y_true = self.predict(gt_reports, **kw)["positive"] y_pred = self.predict(pred_reports, **kw)["positive"] names = list(self.config.labels) res: Dict[str, Any] = {"n_reports": len(gt_reports), "threshold": self.config.default_threshold if threshold is None else threshold} for avg in ("micro", "macro", "weighted"): p, r, f, _ = precision_recall_fscore_support(y_true, y_pred, average=avg, zero_division=0) res[avg] = {"precision": float(p), "recall": float(r), "f1": float(f)} p, r, f, s = precision_recall_fscore_support(y_true, y_pred, average=None, labels=list(range(len(names))), zero_division=0) res["per_label"] = {names[j]: {"precision": float(p[j]), "recall": float(r[j]), "f1": float(f[j]), "support_gt": int(s[j])} for j in range(len(names))} return res def report_f1(gt_reports: List[str], pred_reports: List[str], model=None, tokenizer=None, model_id: str = "chest2vec/chest2vec_labeler", **kw) -> Dict[str, Any]: """Convenience wrapper: load the labeler (if not supplied) and score GT vs predicted reports.""" if model is None: model = Chest2VecLabelerModel.from_pretrained(model_id).eval() return model.score_reports(gt_reports, pred_reports, tokenizer=tokenizer, **kw)