Upload modeling_chest2vec.py with huggingface_hub

modeling_chest2vec.py

@@ -1,301 +1,166 @@
-"""Chest2Vec — LoRA-tuned Qwen3-Embedding model for chest radiology reports.
+"""Chest2Vec — Qwen3-Embedding model (contrastive LoRA merged in) for chest radiology reports.
 
-Load with:
+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
-    model = AutoModel.from_pretrained("chest2vec/chest2vec_0.6B", trust_remote_code=True)
-    emb = model.embed_texts(["Frontal chest radiograph. No pneumothorax."])  # [N, H], L2-normalized
+    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
 
-Architecture:
-  1. Base   : Qwen/Qwen3-Embedding-{0.6B,4B}  (downloaded at runtime)
-  2. Adapter: frozen contrastive LoRA adapter  (./contrastive)
-
-Embeddings use last-token (EOS) pooling with left padding, matching Qwen3-Embedding
-and the Stage-2 training setup. FlashAttention-2 is used when CUDA + flash-attn>=2
-are available (matching training); otherwise it falls back to SDPA so the model
-also loads on CPU.
+Embedding = left-padding-aware last-token (EOS) pooling + L2-norm. Matryoshka: pass `dim=512`/`256`.
 """
-import os
-from typing import Dict, List, Optional
-
+from typing import List, Optional
 import torch
 import torch.nn.functional as F
-
-from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig, PreTrainedModel
-
+from transformers import PreTrainedModel, AutoConfig, AutoModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 from .configuration_chest2vec import Chest2VecConfig
 
-try:
-    from peft import PeftModel
-    _HAS_PEFT = True
-except Exception:
-    PeftModel = None
-    _HAS_PEFT = False
-
-try:
-    from huggingface_hub import snapshot_download
-    _HAS_HUB = True
-except Exception:
-    snapshot_download = None
-    _HAS_HUB = False
-
-
-# ----------------------------------------------------------------------------
-# Attention backend selection
-# ----------------------------------------------------------------------------
-def _flash_attn_available() -> bool:
-    if not torch.cuda.is_available():
-        return False
-    try:
-        import flash_attn  # noqa: F401
-        ver = getattr(flash_attn, "__version__", "0.0.0")
-        return int(str(ver).split(".")[0]) >= 2
-    except Exception:
-        return False
-
 
-def _pick_attn_impl(requested: Optional[str], want_flash: bool) -> str:
-    import warnings
-    if requested:
-        return requested
-    if want_flash and _flash_attn_available():
-        return "flash_attention_2"
-    if want_flash:
-        warnings.warn(
-            "Chest2Vec was trained with FlashAttention-2, but it is unavailable "
-            "(needs CUDA + flash-attn>=2). Falling back to 'sdpa'; embeddings may "
-            "differ very slightly from the reference implementation.",
-            RuntimeWarning,
-        )
-    return "sdpa"
-
-
-# ----------------------------------------------------------------------------
-# Tokenization / pooling helpers (match Qwen3-Embedding + training)
-# ----------------------------------------------------------------------------
 def build_qwen_query(instruction: str, query: str) -> str:
-    return f"Instruct: {str(instruction).strip()}\nQuery: {str(query).strip()}"
-
-
-def get_pool_token_id(tok) -> int:
-    eod_id = tok.convert_tokens_to_ids("<|endoftext|>")
-    if eod_id is None or eod_id < 0:
-        eod_id = tok.pad_token_id
-    return eod_id
+    instruction = str(instruction).strip()
+    return f"Instruct: {instruction}\nQuery: {str(query).strip()}" if instruction else str(query).strip()
 
 
-def encode_with_eos_ids(tok, texts: List[str], max_len: int) -> Dict[str, torch.Tensor]:
-    """add_special_tokens=False, truncate to max_len-1, append <|endoftext|>, left-pad."""
-    pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
-    eod_id = get_pool_token_id(tok)
-    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,
-    )
-    input_ids = [ids + [eod_id] for ids in enc["input_ids"]]
-    attn_mask = [[1] * len(ids) for ids in input_ids]
-    T = max((len(ids) for ids in input_ids), default=1)
-    input_ids = [[pad_id] * (T - len(ids)) + ids for ids in input_ids]
-    attn_mask = [[0] * (T - len(m)) + m for m in attn_mask]
-    return {
-        "input_ids": torch.tensor(input_ids, dtype=torch.long),
-        "attention_mask": torch.tensor(attn_mask, dtype=torch.long),
-    }
+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(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-    """Left-padding-aware last-token (EOS) pooling."""
-    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
-    if left_padding:
-        return last_hidden_states[:, -1]
+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 last_hidden_states[torch.arange(last_hidden_states.size(0), device=last_hidden_states.device), idx]
-
-
-def get_last_hidden_state(model, input_ids, attention_mask):
-    m = model.module if hasattr(model, "module") else model
-    position_ids = attention_mask.long().cumsum(-1) - 1
-    position_ids.masked_fill_(attention_mask == 0, 0)
-    out = m(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
-            use_cache=False, return_dict=True)
-    if getattr(out, "last_hidden_state", None) is not None:
-        return out.last_hidden_state
-    out = m(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
-            output_hidden_states=True, use_cache=False, return_dict=True)
-    return out.hidden_states[-1]
+    return h[torch.arange(h.size(0), device=h.device), idx]
 
 
 class Chest2VecModel(PreTrainedModel):
-    """LoRA-tuned Qwen3-Embedding model producing L2-normalized report embeddings."""
-
     config_class = Chest2VecConfig
-    base_model_prefix = "chest2vec"
-    # Attention is handled by the inner Qwen3 backbone; advertise support so the
-    # transformers attn-implementation validator on this wrapper passes.
-    _supports_sdpa = True
-    _supports_flash_attn_2 = True
-    _supports_flash_attn = True
-    _supports_attention_backend = True
+    base_model_prefix = "model"
 
     def __init__(self, config: Chest2VecConfig):
         super().__init__(config)
-        # The base+adapter are assembled in `from_pretrained` (base downloads at runtime).
-        self.backbone = None
-        self.tokenizer = None
-        self._device = torch.device("cpu")
-        self.register_buffer("_anchor", torch.zeros(1), persistent=False)
+        self.model = _build_encoder(config.encoder_config, getattr(config, "attn_implementation", "sdpa"))
+        self._tokenizer = None
+        self.post_init()
 
     def get_input_embeddings(self):
-        return None
+        return self.model.get_input_embeddings()
 
     def set_input_embeddings(self, value):
-        pass
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        config = kwargs.pop("config", None)
-        device = kwargs.pop("device", None)
-        use_4bit = kwargs.pop("use_4bit", False)
-        attn_implementation = kwargs.pop("attn_implementation", None)
-        torch_dtype = kwargs.pop("torch_dtype", None)
-        token = kwargs.pop("token", None) or kwargs.pop("use_auth_token", None)
-        cache_dir = kwargs.pop("cache_dir", None)
-        # remaining HF plumbing kwargs (state_dict, low_cpu_mem_usage, ...) are ignored
-
-        repo_path = pretrained_model_name_or_path
-        if not os.path.isdir(repo_path):
-            if not _HAS_HUB:
-                raise RuntimeError("huggingface_hub is required to load by repo_id.")
-            repo_path = snapshot_download(repo_path, token=token, cache_dir=cache_dir)
-
-        if config is None:
-            config = Chest2VecConfig.from_pretrained(repo_path)
-
-        if device is None:
-            device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        device_t = torch.device(device)
-        if torch_dtype is None:
-            torch_dtype = torch.bfloat16 if device_t.type == "cuda" else torch.float32
-
-        model = cls(config)
-        model._assemble(repo_path, device=device_t, use_4bit=use_4bit,
-                        attn_implementation=attn_implementation, torch_dtype=torch_dtype, token=token)
-        return model
-
-    def _assemble(self, repo_path, *, device, use_4bit, attn_implementation, torch_dtype, token=None):
-        cfg = self.config
-        if not _HAS_PEFT:
-            raise RuntimeError("peft is required. Install: pip install peft")
-
-        attn_impl = _pick_attn_impl(attn_implementation, bool(cfg.require_flash_attention_2))
-
-        tokenizer = AutoTokenizer.from_pretrained(
-            cfg.base_model, padding_side="left", trust_remote_code=True, token=token
-        )
-        if tokenizer.pad_token_id is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        base_kwargs = dict(trust_remote_code=True, attn_implementation=attn_impl, token=token)
-        if use_4bit:
-            base_kwargs["quantization_config"] = BitsAndBytesConfig(
-                load_in_4bit=True, bnb_4bit_quant_type="nf4",
-                bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
-            )
-            base_kwargs["device_map"] = {"": str(device)}
-        else:
-            base_kwargs["torch_dtype"] = torch_dtype
-            if device.type == "cuda":
-                base_kwargs["device_map"] = {"": str(device)}
-        try:
-            base = AutoModel.from_pretrained(cfg.base_model, **base_kwargs)
-        except TypeError as e:
-            raise RuntimeError("transformers too old for attn_implementation=...; please upgrade.") from e
-        if device.type != "cuda" and not use_4bit:
-            base = base.to(device)
-
-        adapter_dir = os.path.join(repo_path, cfg.adapter_subdir)
-        if not os.path.isfile(os.path.join(adapter_dir, "adapter_config.json")):
-            raise FileNotFoundError(f"adapter_config.json not found under: {adapter_dir}")
-        backbone = PeftModel.from_pretrained(base, adapter_dir)
-        backbone.eval()
-
-        self.backbone = backbone
-        self.tokenizer = tokenizer
-        self._device = device
-        self.eval()
+        self.model.set_input_embeddings(value)
 
     @property
     def device(self):
-        return self._device
-
-    @torch.inference_mode()
-    def embed_texts(self, texts: List[str], *, max_len: Optional[int] = None,
-                    batch_size: int = 16, return_cpu_float32: bool = True) -> torch.Tensor:
-        """Return L2-normalized report embeddings, shape [N, H]."""
-        if self.backbone is None:
-            raise RuntimeError("Model not assembled; load via from_pretrained(...).")
-        max_len = int(max_len or self.config.default_max_len)
-        device = self._device
-        if device.type == "cuda":
-            amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-            use_amp = True
-        else:
-            amp_dtype, use_amp = torch.float32, False
-
-        outs = []
+        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):
-            chunk = [str(t) for t in texts[i:i + batch_size]]
-            enc = encode_with_eos_ids(self.tokenizer, chunk, max_len)
-            input_ids = enc["input_ids"].to(device, non_blocking=True)
-            attention_mask = enc["attention_mask"].to(device, non_blocking=True)
-            with torch.autocast(device_type=("cuda" if device.type == "cuda" else "cpu"),
-                                dtype=amp_dtype, enabled=use_amp):
-                h = get_last_hidden_state(self.backbone, input_ids, attention_mask)
-                emb = F.normalize(last_token_pool(h, attention_mask).float(), p=2, dim=-1)
-            outs.append(emb.detach())
-        embeddings = torch.cat(outs, dim=0)
-        if return_cpu_float32:
-            embeddings = F.normalize(embeddings.float().cpu(), p=2, dim=-1)
-        return embeddings
-
-    @torch.inference_mode()
-    def embed_instruction_query(self, instructions: List[str], queries: List[str], **kw) -> torch.Tensor:
-        if len(instructions) != len(queries):
-            raise ValueError("instructions and queries must have the same length.")
-        return self.embed_texts([build_qwen_query(i, q) for i, q in zip(instructions, queries)], **kw)
-
-    def forward(self, texts: List[str], **kw) -> torch.Tensor:  # type: ignore[override]
-        return self.embed_texts(texts, **kw)
+            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, *, device="cuda",
-                    query_batch_size=256, doc_chunk_size=8192):
-        device_t = torch.device(device if torch.cuda.is_available() else "cpu")
-        q = F.normalize(query_emb.float(), p=2, dim=-1)
-        d = F.normalize(cand_emb.float(), p=2, dim=-1)
-        Nq, _ = q.shape
-        Nd = d.shape[0]
-        k = min(int(k), Nd)
-        top_scores_all = torch.empty((Nq, k), dtype=torch.float32)
-        top_indices_all = torch.empty((Nq, k), dtype=torch.long)
-        for qs in range(0, Nq, query_batch_size):
-            qe = q[qs:qs + query_batch_size].to(device_t, non_blocking=True)
-            bq = qe.size(0)
-            top_scores = torch.full((bq, k), -1e9, device=device_t, dtype=torch.float32)
-            top_indices = torch.full((bq, k), -1, device=device_t, dtype=torch.long)
-            for ds in range(0, Nd, doc_chunk_size):
-                de = d[ds:ds + doc_chunk_size].to(device_t, non_blocking=True)
-                scores = (qe @ de.T).float()
-                chunk = scores.size(1)
-                idx_chunk = torch.arange(ds, ds + chunk, device=device_t, dtype=torch.long).unsqueeze(0).expand(bq, -1)
-                comb_scores = torch.cat([top_scores, scores], dim=1)
-                comb_idx = torch.cat([top_indices, idx_chunk], dim=1)
-                new_scores, new_pos = torch.topk(comb_scores, k, dim=1)
-                top_scores, top_indices = new_scores, comb_idx.gather(1, new_pos)
-            top_scores_all[qs:qs + bq] = top_scores.cpu()
-            top_indices_all[qs:qs + bq] = top_indices.cpu()
-        return top_scores_all, top_indices_all
+    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