Upload README.md with huggingface_hub

README.md

@@ -5,211 +5,69 @@ tags:
 - radiology
 - chest
 - qwen
+base_model:
+- Qwen/Qwen3-Embedding-4B
 library_name: transformers
+pipeline_tag: feature-extraction
 ---
 
 # chest2vec_4B
 
-This repository contains the *delta weights* for a global embedding model on top of **Qwen/Qwen3-Embedding-4B**:
+Chest-radiology **text embedding** model: [`Qwen/Qwen3-Embedding-4B`](https://huggingface.co/Qwen/Qwen3-Embedding-4B)
+contrastively LoRA-adapted for chest CT / CXR report retrieval. Embedding = left-padding-aware
+last-token (EOS) pooling + L2-norm. **Embedding dim: 2560.**
 
-- **LoRA Adapter**: Contrastive LoRA adapter trained with multi-positive sigmoid loss under `./contrastive/`
-- **Inference helper**: `chest2vec.py`
+## Self-contained `AutoModel`
 
-Base model weights are **not** included; they are downloaded from Hugging Face at runtime.
-
-## Model Architecture
-
-Chest2Vec is a two-stage model:
-1. **Base**: Qwen/Qwen3-Embedding-4B (downloaded at runtime)
-2. **LoRA Adapter**: Contrastive LoRA adapter trained with multi-positive sigmoid loss
-3. **Pooling**: Last-token pooling (EOS token) for global embeddings
-
-The model produces **global embeddings only** (no section-specific embeddings).
-
-## 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
+The LoRA adapter is **merged into the weights** (`model.safetensors`) and the tokenizer is bundled,
+so loading needs **no `chest2vec` package and no download of the base Qwen3-Embedding weights**:
 
 ```python
-from chest2vec import Chest2Vec
-
-# Load model from Hugging Face Hub
-m = Chest2Vec.from_pretrained("lukeingawesome/chest2vec_4b_chest", device="cuda:0")
+from transformers import AutoModel, AutoTokenizer
+model = AutoModel.from_pretrained("chest2vec/chest2vec_4B", trust_remote_code=True).eval()
+tok   = AutoTokenizer.from_pretrained("chest2vec/chest2vec_4B", trust_remote_code=True)
+
+docs = ["Bibasilar atelectasis with small bilateral pleural effusions. Cardiomegaly."]
+doc_emb = model.embed_texts(docs, tokenizer=tok)                      # [N, 2560], L2-normalized
+
+# instruction-conditioned query
+q_emb = model.embed_instruction_query(
+    "Retrieve the chest CT report that is similar to the given report.",
+    ["pleural effusion and cardiomegaly"], tokenizer=tok)
+vals, idx = model.cosine_topk(q_emb, doc_emb, k=5)
 ```
 
-### 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 (last-token pooling)
-emb = out.embedding  # [N, H]
-```
+## Matryoshka embeddings
 
-### Candidate Embeddings (Retrieval Bank)
+Matryoshka (MRL)-trained — truncate to **512** or **256** dims (keep first *N* dims, re-normalize):
 
 ```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_emb = cand_out.embedding  # [N, H]
+emb512 = model.embed_texts(docs, tokenizer=tok, dim=512)
+emb256 = model.embed_texts(docs, tokenizer=tok, dim=256)
 ```
 
-### Retrieval Example (Cosine Top-K)
+Recommended dims: **2560 (full) · 512 · 256** (`config.matryoshka_dims`). Use the same `dim` for query and corpus.
 
-```python
-# Query embeddings
-q = out.embedding  # [Nq, H]
+## Recommended instructions
 
-# Document embeddings
-d = cand_out.embedding  # [Nd, H]
+Instruction-conditioned (`Instruct: {instruction}\nQuery: {report}`). Apply to the **query** side;
+embed the corpus without an instruction. Trained on chest **CT and CXR** across these families:
 
-# 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
+**Retrieval** — `Retrieve the chest CT report that is similar to the given report.` ·
+`Retrieve the CXR report that is similar to the given report.` ·
+`Retrieve the CXR report that is similar to the given report with prior reference omitted.`
 
-print(f"Top-5 scores: {scores[0]}")
-print(f"Top-5 indices: {idx[0]}")
-```
+**Summarization** — `Summarize the following chest CT report` · `Summarize the following CXR report` · `Summarize the given report.`
 
-## API Reference
+**Entity extraction (leaf)** — `Given the following chest CT report, extract the presence/absence of entities` · `Given the following CXR report, extract the presence/absence of entities`
 
-### `Chest2Vec.from_pretrained()`
+**Entity extraction (upper/coarse)** — `Given the following chest CT report, extract the presence/absence of upper-level entities` · `Given the following CXR report, extract the presence/absence of upper class entities`
 
-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
-)
-```
+**Anatomy-specific** — `From the following chest {CT report | X-ray report}, extract and return only the findings related to {REGION}, ignoring all information about other structures.`
+- CT regions: lungs · airways and trachea · pleura · mediastinum and hilum · cardiovascular system · chest wall · bones and spine · upper abdomen · lower neck
+- CXR regions: lungs and airways · pleura · hila and mediastinum · cardiovascular system · musculoskeletal structures and chest wall · tubes, catheters, and support devices · abdomen
 
-### `embed_instruction_query()`
-
-Embed instruction-query pairs. Returns `EmbedOutput` with:
-- `embedding`: `[N, H]` - global embeddings (L2-normalized, last-token pooling)
-
-```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
-)
-```
-
-Returns `EmbedOutput` with:
-- `embedding`: `[N, H]` - global embeddings (L2-normalized, last-token pooling)
-
-### `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
-- `contrastive/` - 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_4b_chest,
-  title={Chest2Vec: Global Embeddings for Chest X-Ray Reports},
-  author={Your Name},
-  year={2024},
-  howpublished={\url{https://huggingface.co/lukeingawesome/chest2vec_4b_chest}}
-}
-```
-
-## License
-
-[Specify your license here]
-
-## Usage (🤗 transformers)
-
-```python
-from transformers import AutoModel
-
-# base Qwen3-Embedding weights download automatically; needs trust_remote_code
-model = AutoModel.from_pretrained("chest2vec/chest2vec_4B", trust_remote_code=True)
-
-emb = model.embed_texts([
-    "Frontal chest radiograph. No focal consolidation. No pneumothorax. Heart size normal.",
-])
-emb            # [N, H] L2-normalized report embedding (last-token / EOS pooling)
-
-# similarity
-(emb[0] @ emb[1])    # cosine similarity (rows are unit-norm)
-```
+## Details
 
-FlashAttention-2 is used automatically on CUDA when `flash-attn>=2` is installed
-(matching training); otherwise it falls back to SDPA so the model also loads on CPU.
+- **Base:** Qwen/Qwen3-Embedding-4B (Apache-2.0) — architecture rebuilt from the bundled config; merged weights loaded from this repo. Default attention `sdpa` (use `flash_attention_2` on Ampere+ for speed).
+- Merged weights reproduce the original adapter-based embeddings to **cosine ≥ 0.999**.