Upload folder using huggingface_hub

.gitattributes

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +34,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+=======
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+>>>>>>> 2dc7409 (Release chest2vec_0.6b_cxr (Stage2 LoRA + Stage3 pooler + API))

.gitignore

@@ -0,0 +1,24 @@
+# Build artifacts
+dist/
+build/
+*.egg-info/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Jupyter
+.ipynb_checkpoints/
+

PUBLISH.md

@@ -0,0 +1,135 @@
+# Publishing chest2vec_0.6b_chest
+
+This guide covers publishing to both **Hugging Face Hub** (model repository) and **PyPI** (Python package).
+
+## Option 1: Publish to Hugging Face Hub
+
+### Prerequisites
+
+1. Install Hugging Face Hub CLI:
+   ```bash
+   pip install huggingface_hub
+   ```
+
+2. Login to Hugging Face:
+   ```bash
+   huggingface-cli login
+   # Enter your HF token when prompted
+   ```
+
+### Initialize Git Repository (if not already)
+
+```bash
+cd /opt/project/chest2vec/export_chest2vec_0.6b_chest
+git init
+git add .
+git commit -m "Initial commit: chest2vec_0.6b_chest"
+```
+
+### Create Repository on Hugging Face Hub
+
+1. Go to https://huggingface.co/new
+2. Create a new repository named `chest2vec_0.6b_chest` (or `chest2vec/chest2vec_0.6b_chest` if under an organization)
+3. Choose "Model" as the repository type
+
+### Push to Hugging Face Hub
+
+```bash
+# Add HF remote (replace with your username/org)
+git remote add origin https://huggingface.co/YOUR_USERNAME/chest2vec_0.6b_chest
+
+# Push to HF Hub
+git push origin main
+```
+
+Or use the HF Hub CLI:
+
+```bash
+huggingface-cli upload YOUR_USERNAME/chest2vec_0.6b_chest . --repo-type model
+```
+
+### Verify Upload
+
+Visit `https://huggingface.co/YOUR_USERNAME/chest2vec_0.6b_chest` to verify all files are uploaded.
+
+---
+
+## Option 2: Publish to PyPI
+
+### Prerequisites
+
+1. Create a PyPI account at https://pypi.org/account/register/
+2. Create an API token at https://pypi.org/manage/account/token/
+3. Install build tools:
+   ```bash
+   pip install build twine
+   ```
+
+### Build the Package
+
+```bash
+cd /opt/project/chest2vec/export_chest2vec_0.6b_chest
+python3 -m build
+```
+
+This creates:
+- `dist/chest2vec-0.6.0-py3-none-any.whl`
+- `dist/chest2vec-0.6.0.tar.gz`
+
+### Upload to PyPI
+
+#### Test first on TestPyPI
+
+```bash
+# Upload to TestPyPI first to test
+twine upload --repository testpypi dist/*
+# You'll be prompted for:
+#   username: __token__
+#   password: your API token
+```
+
+Test installation:
+```bash
+pip install --index-url https://test.pypi.org/simple/ chest2vec
+```
+
+#### Then upload to PyPI
+
+```bash
+twine upload dist/*
+# You'll be prompted for:
+#   username: __token__
+#   password: your API token
+```
+
+### After Publishing
+
+Once published, users can install with:
+
+```bash
+pip install chest2vec
+```
+
+**Note:** Users will still need to install PyTorch and flash-attention separately as documented in the README, since PyPI doesn't support custom index URLs in dependencies.
+
+---
+
+## Option 3: Publish to Both
+
+You can publish to both platforms:
+
+1. **First publish to Hugging Face Hub** (for model weights and repository)
+2. **Then publish to PyPI** (for easy Python package installation)
+
+Users can then either:
+- Load from HF Hub: `Chest2Vec.from_pretrained("YOUR_USERNAME/chest2vec_0.6b_chest")`
+- Install from PyPI: `pip install chest2vec` (then load from HF Hub or local path)
+
+---
+
+## Important Notes
+
+- **Model weights** (section_pooler.pt, contrastive/) should be on Hugging Face Hub
+- **Python package** can be on PyPI for easier installation
+- The package name on PyPI is `chest2vec` (shared across all variants)
+- Make sure to update version numbers in `setup.py` and `pyproject.toml` if publishing a new version

README.md

@@ -1,3 +1,220 @@
----
-license: apache-2.0
----
+---
+tags:
+- text-embeddings
+- retrieval
+- radiology
+- chest
+- qwen
+library_name: transformers
+---
+
+# chest2vec_0.6b_chest
+
+This repository contains the *delta weights and pooling head* for a section-aware embedding model on top of **Qwen/Qwen3-Embedding-0.6B**:
+
+- **Stage-2**: Frozen LoRA adapter (contrastive) under `./contrastive/`
+- **Stage-3**: Section pooler `section_pooler.pt` producing **9 section embeddings**
+- **Inference helper**: `chest2vec.py`
+
+Base model weights are **not** included; they are downloaded from Hugging Face at runtime.
+
+## Model Architecture
+
+Chest2Vec is a three-stage model:
+1. **Base**: Qwen/Qwen3-Embedding-0.6B (downloaded at runtime)
+2. **Stage-2**: Contrastive LoRA adapter trained with multi-positive sigmoid loss
+3. **Stage-3**: Section-aware query-attention pooler producing embeddings for 9 radiology report sections
+
+## Sections
+
+The model produces embeddings for 9 distinct sections:
+
+1. Lungs and Airways
+2. Pleura
+3. Cardiovascular
+4. Hila and Mediastinum
+5. Tubes & Devices
+6. Musculoskeletal and Chest Wall
+7. Abdominal
+8. impression
+9. Other
+
+## Installation
+
+Install the package and all dependencies:
+
+```bash
+# Install PyTorch with CUDA 12.6 support
+pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
+
+# Install transformers and trl
+pip install transformers==4.57.3 trl==0.9.3
+
+# Install deepspeed
+pip install deepspeed==0.16.9
+
+# Install flash-attention
+pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+
+# Install chest2vec package
+pip install chest2vec
+```
+
+Or use the installation script:
+
+```bash
+bash install_deps.sh
+```
+
+## Requirements
+
+This model **requires FlashAttention-2** (CUDA) by default, which is automatically installed with the package.
+
+## Quickstart
+
+### Installation + Loading
+
+```python
+from chest2vec import Chest2Vec
+
+# Load model from Hugging Face Hub
+m = Chest2Vec.from_pretrained("chest2vec/chest2vec_0.6b_chest", device="cuda:0")
+```
+
+### Instruction + Query Embeddings
+
+```python
+instructions = ["Find findings about the lungs."]
+queries = ["Consolidation in the right lower lobe."]
+
+out = m.embed_instruction_query(instructions, queries, max_len=512, batch_size=8)
+
+# Global embedding (derived): mean of 9 section vectors then L2-normalized
+g = out.global_embedding                 # [N, H]
+
+# Per-section embeddings (by full name)
+lung = out.by_section_name["Lungs and Airways"]  # [N, H]
+imp  = out.by_section_name["impression"]          # [N, H]
+
+# Or use aliases (case-insensitive)
+lung = out.by_alias["lungs"]   # [N, H]
+cardio = out.by_alias["cardio"] # [N, H]
+```
+
+### Candidate Embeddings (Retrieval Bank)
+
+```python
+candidates = [
+    "Lungs are clear. No focal consolidation.",
+    "Pleural effusion on the left.",
+    "Cardiomediastinal silhouette is normal."
+]
+
+cand_out = m.embed_texts(candidates, max_len=512, batch_size=16)
+
+cand_global = cand_out.global_embedding  # [N, H]
+cand_lung   = cand_out.by_alias["lungs"]  # [N, H]
+```
+
+### Retrieval Example (Cosine Top-K)
+
+```python
+# Query embeddings for "Lungs and Airways" section
+q = out.by_alias["lungs"]       # [Nq, H]
+
+# Document embeddings for "Lungs and Airways" section
+d = cand_out.by_alias["lungs"]  # [Nd, H]
+
+# Compute top-k cosine similarities
+scores, idx = Chest2Vec.cosine_topk(q, d, k=5, device="cuda")
+# scores: [Nq, k] - similarity scores
+# idx: [Nq, k] - indices of top-k candidates
+
+print(f"Top-5 scores: {scores[0]}")
+print(f"Top-5 indices: {idx[0]}")
+```
+
+## API Reference
+
+### `Chest2Vec.from_pretrained()`
+
+Load the model from Hugging Face Hub or local path.
+
+```python
+m = Chest2Vec.from_pretrained(
+    repo_id_or_path: str,      # Hugging Face repo ID or local path
+    device: str = "cuda:0",    # Device to load model on
+    use_4bit: bool = False,    # Use 4-bit quantization
+    force_flash_attention_2: bool = True
+)
+```
+
+### `embed_instruction_query()`
+
+Embed instruction-query pairs. Returns `EmbedOutput` with:
+- `section_matrix`: `[N, 9, H]` - embeddings for all 9 sections
+- `global_embedding`: `[N, H]` - global embedding (mean of sections, L2-normalized)
+- `by_section_name`: Dict mapping full section names to `[N, H]` tensors
+- `by_alias`: Dict mapping aliases to `[N, H]` tensors
+
+```python
+out = m.embed_instruction_query(
+    instructions: List[str],
+    queries: List[str],
+    max_len: int = 512,
+    batch_size: int = 16
+)
+```
+
+### `embed_texts()`
+
+Embed plain texts (for document/candidate encoding).
+
+```python
+out = m.embed_texts(
+    texts: List[str],
+    max_len: int = 512,
+    batch_size: int = 16
+)
+```
+
+### `cosine_topk()`
+
+Static method for efficient top-k cosine similarity search.
+
+```python
+scores, idx = Chest2Vec.cosine_topk(
+    query_emb: torch.Tensor,  # [Nq, H]
+    cand_emb: torch.Tensor,   # [Nd, H]
+    k: int = 10,
+    device: str = "cuda"
+)
+```
+
+## Model Files
+
+- `chest2vec.py` - Model class and inference utilities
+- `chest2vec_config.json` - Model configuration
+- `section_pooler.pt` - Stage-3 pooler weights
+- `section_pooler_config.json` - Pooler configuration
+- `contrastive/` - Stage-2 LoRA adapter directory
+  - `adapter_config.json` - LoRA adapter configuration
+  - `adapter_model.safetensors` - LoRA adapter weights
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{chest2vec_0.6b_chest,
+  title={Chest2Vec: Section-Aware Embeddings for Chest X-Ray Reports},
+  author={Your Name},
+  year={2024},
+  howpublished={\url{https://huggingface.co/chest2vec/chest2vec_0.6b_chest}}
+}
+```
+
+## License
+
+[Specify your license here]
+

__init__.py

@@ -0,0 +1,5 @@
+from chest2vec import Chest2Vec, EmbedOutput
+
+__all__ = ["Chest2Vec", "EmbedOutput"]
+__version__ = "0.6.0"
+

chest2vec.py

@@ -0,0 +1,553 @@
+import os
+import json
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
+
+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
+
+
+# -----------------------------
+# Sections (must match training)
+# -----------------------------
+SECTION_NAMES = [
+    "Lungs and Airways",
+    "Pleura",
+    "Cardiovascular",
+    "Hila and Mediastinum",
+    "Tubes & Devices",
+    "Musculoskeletal and Chest Wall",
+    "Abdominal",
+    "impression",
+    "Other",
+]
+
+SECTION_ALIASES = {
+    "global": "global",
+    "lungs": "Lungs and Airways",
+    "lung": "Lungs and Airways",
+    "pleura": "Pleura",
+    "cardio": "Cardiovascular",
+    "cardiovascular": "Cardiovascular",
+    "hila": "Hila and Mediastinum",
+    "mediastinum": "Hila and Mediastinum",
+    "tubes": "Tubes & Devices",
+    "devices": "Tubes & Devices",
+    "msk": "Musculoskeletal and Chest Wall",
+    "musculoskeletal": "Musculoskeletal and Chest Wall",
+    "abd": "Abdominal",
+    "abdominal": "Abdominal",
+    "impression": "impression",
+    "other": "Other",
+}
+
+
+def require_flash_attention_2() -> str:
+    if not torch.cuda.is_available():
+        raise RuntimeError("FlashAttention-2 requires CUDA, but torch.cuda.is_available() is False.")
+    try:
+        import flash_attn  # noqa: F401
+        ver = getattr(flash_attn, "__version__", "0.0.0")
+        major = int(str(ver).split(".")[0])
+        if major < 2:
+            raise RuntimeError(f"flash-attn version {ver} < 2.0.0")
+    except Exception as e:
+        raise RuntimeError(
+            "FlashAttention-2 is REQUIRED but not available/importable.\n"
+            "Install flash-attn>=2 and ensure it matches your torch/CUDA.\n"
+            f"Import/Version error: {repr(e)}"
+        )
+    return "flash_attention_2"
+
+
+def build_qwen_query(instruction: str, query: str) -> str:
+    instruction = str(instruction).strip()
+    query = str(query).strip()
+    return f"Instruct: {instruction}\nQuery: {query}"
+
+
+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
+
+
+def encode_with_eos_ids(tok, texts: List[str], max_len: int) -> Dict[str, torch.Tensor]:
+    """
+    Must match Stage-3 training:
+      - add_special_tokens=False
+      - truncation 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) if input_ids else 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 last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    """
+    Left-padding aware last-token pooling (extracts EOS token embedding).
+    """
+    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]
+
+
+def get_last_hidden_state(model, input_ids, attention_mask):
+    """
+    Provide position_ids for left padding (FlashAttention-2).
+    """
+    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 hasattr(out, "last_hidden_state"):
+        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]
+
+
+# -----------------------------
+# Stage-3 pooler (query_attn)
+# -----------------------------
+class SectionQueryAttnPooler(nn.Module):
+    """
+    Match your Stage-3 training pooler.
+    """
+    def __init__(
+        self,
+        hidden_size: int,
+        num_sections: int,
+        mlp_hidden: int,
+        use_layernorm: bool = True,
+        pool_dropout: float = 0.1,
+        pool_scale: float = 0.0,  # 0 => 1/sqrt(H)
+    ):
+        super().__init__()
+        self.hidden_size = int(hidden_size)
+        self.num_sections = int(num_sections)
+
+        self.ln = nn.LayerNorm(self.hidden_size) if use_layernorm else nn.Identity()
+
+        self.pool_queries = nn.Parameter(torch.empty(self.num_sections, self.hidden_size))
+        nn.init.normal_(self.pool_queries, mean=0.0, std=0.02)
+
+        self.pool_scale = float(pool_scale) if (pool_scale and pool_scale > 0) else (1.0 / math.sqrt(self.hidden_size))
+        self.pool_dropout = nn.Dropout(pool_dropout) if pool_dropout and pool_dropout > 0 else nn.Identity()
+
+        # Bias-free MLP
+        self.mlp = nn.Sequential(
+            nn.Linear(self.hidden_size, int(mlp_hidden), bias=False),
+            nn.GELU(),
+            nn.Linear(int(mlp_hidden), self.hidden_size, bias=False),
+        )
+
+    def forward_all(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        # hidden_states: [B,T,H] -> [B,S,H]
+        if isinstance(self.ln, nn.LayerNorm):
+            x = F.layer_norm(
+                hidden_states.float(),
+                self.ln.normalized_shape,
+                self.ln.weight.float() if self.ln.weight is not None else None,
+                self.ln.bias.float() if self.ln.bias is not None else None,
+                self.ln.eps,
+            ).to(dtype=hidden_states.dtype)
+        else:
+            x = hidden_states
+
+        scores = torch.einsum("bth,sh->bts", x.float(), self.pool_queries.float()) * self.pool_scale
+        scores = scores.masked_fill(attention_mask.unsqueeze(-1) == 0, -1e4)
+
+        attn = torch.softmax(scores, dim=1).to(dtype=x.dtype)  # [B,T,S]
+        attn = self.pool_dropout(attn)
+
+        pooled = torch.einsum("bth,bts->bsh", x, attn)  # [B,S,H]
+        pooled = pooled.to(dtype=next(self.mlp.parameters()).dtype)
+        pooled = self.mlp(pooled)
+
+        return F.normalize(pooled, p=2, dim=-1)
+
+
+def _ensure_pooler_device_dtype(pooler: nn.Module, device: torch.device, dtype: torch.dtype) -> None:
+    p = next(pooler.parameters(), None)
+    if p is None:
+        return
+    if p.device != device or p.dtype != dtype:
+        pooler.to(device=device, dtype=dtype)
+
+
+def _read_json(path: str) -> Dict[str, Any]:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _resolve_repo_path(repo_id_or_path: str) -> str:
+    # If it's a local directory, use it as-is.
+    if os.path.isdir(repo_id_or_path):
+        return repo_id_or_path
+    # Otherwise treat as HF repo_id and download snapshot.
+    if not _HAS_HUB:
+        raise RuntimeError(
+            "huggingface_hub is required to load by repo_id. "
+            "Install it: pip install huggingface_hub"
+        )
+    return snapshot_download(repo_id_or_path)
+
+
+@dataclass
+class EmbedOutput:
+    # Always available:
+    section_matrix: torch.Tensor  # [N,S,H], float32 on CPU by default
+    global_embedding: torch.Tensor  # [N,H], float32 on CPU by default
+    # Convenience dicts:
+    by_section_name: Dict[str, torch.Tensor]  # each [N,H]
+    by_alias: Dict[str, torch.Tensor]         # alias -> [N,H]
+
+
+class Chest2Vec:
+    """
+    Lightweight wrapper:
+      - loads base Qwen3-Embedding
+      - applies LoRA adapter
+      - attaches Stage-3 section pooler
+    """
+    def __init__(self, tokenizer, model, pooler, sections: List[str], device: torch.device):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.pooler = pooler
+        self.sections = list(sections)
+        self.device = device
+
+        self.model.eval()
+        self.pooler.eval()
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id_or_path: str,
+        *,
+        device: str = "cuda:0",
+        use_4bit: bool = False,
+        force_flash_attention_2: bool = True,
+    ) -> "Chest2Vec":
+        repo_path = _resolve_repo_path(repo_id_or_path)
+
+        cfg_path = os.path.join(repo_path, "chest2vec_config.json")
+        if not os.path.isfile(cfg_path):
+            raise FileNotFoundError(f"Missing chest2vec_config.json in {repo_path}")
+        cfg = _read_json(cfg_path)
+
+        base_model = str(cfg["base_model"])
+        adapter_subdir = str(cfg.get("adapter_subdir", "contrastive"))
+        pooler_pt = str(cfg.get("pooler_pt", "section_pooler.pt"))
+        pooler_cfg = str(cfg.get("pooler_cfg", "section_pooler_config.json"))
+        sections = cfg.get("sections", SECTION_NAMES)
+
+        if force_flash_attention_2 or bool(cfg.get("require_flash_attention_2", False)):
+            attn_impl = require_flash_attention_2()
+        else:
+            attn_impl = "sdpa"
+
+        if not _HAS_PEFT:
+            raise RuntimeError("peft is required. Install: pip install peft")
+
+        device_t = torch.device(device)
+
+        tokenizer = AutoTokenizer.from_pretrained(base_model, padding_side="left", trust_remote_code=True)
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        device_map = {"": str(device_t)}
+
+        # Load base model with FlashAttention-2
+        if use_4bit:
+            qconf = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+            )
+            try:
+                base = AutoModel.from_pretrained(
+                    base_model,
+                    trust_remote_code=True,
+                    attn_implementation=attn_impl,
+                    quantization_config=qconf,
+                    device_map=device_map,
+                )
+            except TypeError as e:
+                raise RuntimeError(
+                    "Your transformers version does not support attn_implementation=... "
+                    "Upgrade transformers to use FlashAttention-2."
+                ) from e
+        else:
+            try:
+                base = AutoModel.from_pretrained(
+                    base_model,
+                    trust_remote_code=True,
+                    attn_implementation=attn_impl,
+                    torch_dtype=torch.bfloat16,
+                    device_map=device_map,
+                )
+            except TypeError as e:
+                raise RuntimeError(
+                    "Your transformers version does not support attn_implementation=... "
+                    "Upgrade transformers to use FlashAttention-2."
+                ) from e
+
+        # Load adapter from this repo folder
+        adapter_dir = os.path.join(repo_path, 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}")
+
+        model = PeftModel.from_pretrained(base, adapter_dir)
+        model.eval()
+
+        # Attach section pooler
+        pooler_cfg_path = os.path.join(repo_path, pooler_cfg)
+        pooler_pt_path = os.path.join(repo_path, pooler_pt)
+        if not os.path.isfile(pooler_cfg_path):
+            raise FileNotFoundError(f"Missing pooler config: {pooler_cfg_path}")
+        if not os.path.isfile(pooler_pt_path):
+            raise FileNotFoundError(f"Missing pooler weights: {pooler_pt_path}")
+
+        pcfg = _read_json(pooler_cfg_path)
+
+        hidden_size = int(getattr(model.module if hasattr(model, "module") else model, "config").hidden_size)
+        mlp_hidden = int(pcfg.get("mlp_hidden", hidden_size))
+        use_layernorm = bool(pcfg.get("use_layernorm", True))
+        pool_dropout = float(pcfg.get("pool_dropout", 0.1))
+        pool_scale = float(pcfg.get("pool_scale", 0.0))
+
+        pooler = SectionQueryAttnPooler(
+            hidden_size=hidden_size,
+            num_sections=len(sections),
+            mlp_hidden=mlp_hidden,
+            use_layernorm=use_layernorm,
+            pool_dropout=pool_dropout,
+            pool_scale=pool_scale,
+        )
+        sd = torch.load(pooler_pt_path, map_location="cpu")
+        pooler.load_state_dict(sd, strict=True)
+        pooler.eval()
+
+        # Move pooler to same device/dtype as hidden states
+        # (we keep inference in autocast)
+        pooler.to(device=device_t, dtype=torch.bfloat16 if device_t.type == "cuda" else torch.float32)
+
+        return cls(tokenizer=tokenizer, model=model, pooler=pooler, sections=sections, device=device_t)
+
+    @torch.inference_mode()
+    def embed_texts(
+        self,
+        texts: List[str],
+        *,
+        max_len: int = 512,
+        batch_size: int = 16,
+        return_cpu_float32: bool = True,
+    ) -> EmbedOutput:
+        """
+        Encodes arbitrary texts (candidates, section strings, etc.)
+        
+        NOTE: This uses Stage-3 section pooling:
+        - Section embeddings: section_pooler → [B,S,H] (9 section-specific embeddings)
+        - Global embedding: EOS token embedding extracted BEFORE pooler → [B,H] (matches Stage-3 training)
+        
+        Returns:
+          - section_matrix: [N,9,H] - section-specific embeddings
+          - global_embedding: [N,H] - EOS token embedding (extracted before pooler)
+          - by_section_name: dict[name] -> [N,H]
+          - by_alias: dict['lungs'/'impression'/...] -> [N,H]
+        """
+        # Determine AMP
+        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 = torch.float32
+            use_amp = False
+
+        outs_sec = []
+        outs_global = []
+        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.model, input_ids, attention_mask)  # [B,T,H]
+                
+                # Global embedding: extract EOS token embedding BEFORE pooler (matches Stage-3 training)
+                global_eos = last_token_pool(h, attention_mask)  # [B,H]
+                global_eos = F.normalize(global_eos.float(), p=2, dim=-1)
+                
+                # Section embeddings: pass through pooler
+                _ensure_pooler_device_dtype(self.pooler, device=h.device, dtype=h.dtype)
+                sec = self.pooler.forward_all(h, attention_mask)  # [B,S,H] normalized
+
+            outs_sec.append(sec.detach())
+            outs_global.append(global_eos.detach())
+
+        section_matrix = torch.cat(outs_sec, dim=0)  # on device, dtype ~ bf16
+        global_emb = torch.cat(outs_global, dim=0)  # on device, dtype ~ bf16
+
+        # Move to CPU float32 if requested (recommended for retrieval stability)
+        if return_cpu_float32:
+            section_matrix_cpu = section_matrix.float().cpu()
+            # re-normalize to fix any numerical drift
+            section_matrix_cpu = F.normalize(section_matrix_cpu, p=2, dim=-1)
+            global_cpu = global_emb.float().cpu()
+            global_cpu = F.normalize(global_cpu, p=2, dim=-1)
+        else:
+            section_matrix_cpu = section_matrix
+            global_cpu = global_emb
+
+        by_section_name = {name: section_matrix_cpu[:, idx, :] for idx, name in enumerate(self.sections)}
+
+        # Helpful aliases for quick access
+        by_alias: Dict[str, torch.Tensor] = {}
+        by_alias["global"] = global_cpu
+        for alias, real in SECTION_ALIASES.items():
+            if real == "global":
+                continue
+            if real in by_section_name:
+                by_alias[alias] = by_section_name[real]
+
+        return EmbedOutput(
+            section_matrix=section_matrix_cpu,
+            global_embedding=global_cpu,
+            by_section_name=by_section_name,
+            by_alias=by_alias,
+        )
+
+    @torch.inference_mode()
+    def embed_instruction_query(
+        self,
+        instructions: List[str],
+        queries: List[str],
+        *,
+        max_len: int = 512,
+        batch_size: int = 16,
+        return_cpu_float32: bool = True,
+    ) -> EmbedOutput:
+        if len(instructions) != len(queries):
+            raise ValueError("instructions and queries must have the same length.")
+        q_texts = [build_qwen_query(i, q) for i, q in zip(instructions, queries)]
+        return self.embed_texts(
+            q_texts,
+            max_len=max_len,
+            batch_size=batch_size,
+            return_cpu_float32=return_cpu_float32,
+        )
+
+    @staticmethod
+    def cosine_topk(
+        query_emb: torch.Tensor,     # [Nq,H] CPU float32 recommended
+        cand_emb: torch.Tensor,      # [Nd,H] CPU float32 recommended
+        k: int = 10,
+        *,
+        device: str = "cuda",
+        query_batch_size: int = 256,
+        doc_chunk_size: int = 8192,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Chunked cosine top-k, stable in float32.
+        Returns (top_scores [Nq,k], top_indices [Nq,k]) on CPU.
+        """
+        device_t = torch.device(device)
+        q = F.normalize(query_emb.float(), p=2, dim=-1)
+        d = F.normalize(cand_emb.float(), p=2, dim=-1)
+        Nq, H = 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)
+                new_idx = comb_idx.gather(1, new_pos)
+
+                top_scores, top_indices = new_scores, new_idx
+
+            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

chest2vec_config.json

@@ -0,0 +1,22 @@
+{
+  "name": "chest2vec_0.6b_chest",
+  "base_model": "Qwen/Qwen3-Embedding-0.6B",
+  "adapter_subdir": "contrastive",
+  "pooler_pt": "section_pooler.pt",
+  "pooler_cfg": "section_pooler_config.json",
+  "require_flash_attention_2": true,
+  "default_max_len": 512,
+  "sections": [
+    "Lungs and Airways",
+    "Pleura",
+    "Cardiovascular",
+    "Hila and Mediastinum",
+    "Tubes & Devices",
+    "Musculoskeletal and Chest Wall",
+    "Abdominal",
+    "impression",
+    "Other"
+  ],
+  "global_pool": "mean_of_sections_then_l2"
+}
+

contrastive/adapter_config.json

@@ -0,0 +1,42 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "Qwen3Model",
+    "parent_library": "transformers.models.qwen3.modeling_qwen3"
+  },
+  "base_model_name_or_path": "Qwen/Qwen3-Embedding-0.6B",
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "k_proj",
+    "v_proj",
+    "down_proj",
+    "gate_proj",
+    "up_proj",
+    "q_proj"
+  ],
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_rslora": false
+}

contrastive/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00c52073fa4941ac01672b5961fdeccb24f2e46240f8066415dc14910f8600ed
+size 40419816

install_deps.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+# Installation script for chest2vec dependencies
+# This script installs PyTorch and flash-attention with the correct versions
+
+set -e
+
+echo "Installing PyTorch packages from custom index..."
+pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
+
+echo "Installing flash-attention from GitHub release..."
+pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+
+echo "Installing chest2vec package..."
+pip install chest2vec
+
+echo "Installation complete!"
+

pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "chest2vec"
+version = "0.6.0"
+description = "Section-aware embeddings for chest X-ray reports"
+readme = "README.md"
+requires-python = ">=3.8"
+dependencies = [
+    "transformers==4.57.3",
+    "trl==0.9.3",
+    "deepspeed==0.16.9",
+    "peft",
+    "huggingface_hub",
+    "bitsandbytes",
+    "accelerate",
+    "numpy",
+]
+
+[project.urls]
+Homepage = "https://github.com/chest2vec/chest2vec"
+

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.6.0
+torchvision==0.21.0
+torchaudio==2.6.0
+transformers==4.57.3
+trl==0.9.3
+deepspeed==0.16.9
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+peft
+huggingface_hub
+bitsandbytes
+accelerate
+numpy

section_pooler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79b316d23166dc0b4dfef9cd9b7e25c9ecf5aa9ed7e800b8e1080be10a85db7c
+size 4219403

section_pooler_config.json

@@ -0,0 +1,20 @@
+{
+  "pooler_type": "query_attn",
+  "sections": [
+    "Lungs and Airways",
+    "Pleura",
+    "Cardiovascular",
+    "Hila and Mediastinum",
+    "Tubes & Devices",
+    "Musculoskeletal and Chest Wall",
+    "Abdominal",
+    "impression",
+    "Other"
+  ],
+  "hidden_size": 1024,
+  "mlp_hidden": 1024,
+  "use_layernorm": true,
+  "pool_dropout": 0.1,
+  "pool_scale": 0.0,
+  "loss": "multipos_sigmoid"
+}

setup.py

@@ -0,0 +1,44 @@
+from setuptools import setup, find_packages
+from pathlib import Path
+
+# Read README for long description
+readme_file = Path(__file__).parent / "README.md"
+long_description = readme_file.read_text(encoding="utf-8") if readme_file.exists() else ""
+
+setup(
+    name="chest2vec",
+    version="0.6.0",
+    description="Section-aware embeddings for chest X-ray reports",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Chest2Vec Team",
+    url="https://github.com/chest2vec/chest2vec",
+    packages=find_packages(),
+    py_modules=["chest2vec"],
+    include_package_data=True,
+    package_data={"": ["__init__.py"]},
+    install_requires=[
+        "transformers==4.57.3",
+        "trl==0.9.3",
+        "deepspeed==0.16.9",
+        "peft",
+        "huggingface_hub",
+        "bitsandbytes",
+        "accelerate",
+        "numpy",
+    ],
+    python_requires=">=3.8",
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)
+