chest2vec_0.6b_chest

chest2vec_0.6b_chest

A chest-radiology text embedding model: Qwen/Qwen3-Embedding-0.6B contrastively LoRA-adapted for chest CT / CXR report retrieval.

The embedding is the left-padding-aware last-token (EOS) pooled final hidden state, L2-normalized.

Self-contained AutoModel (recommended)

The LoRA adapter is merged into the weights (model.safetensors, bf16) and the tokenizer is bundled, so the model loads with no chest2vec package and no download of the base Qwen3-Embedding weights:

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("lukeingawesome/chest2vec_0.6b_chest", trust_remote_code=True).eval()
tok   = AutoTokenizer.from_pretrained("lukeingawesome/chest2vec_0.6b_chest", trust_remote_code=True)

# documents
docs = ["Bibasilar atelectasis with small bilateral pleural effusions. Cardiomegaly.",
        "Several pulmonary nodules, largest 6 mm in the right lower lobe. Mild emphysema."]
doc_emb = model.embed(docs, tokenizer=tok)            # [N, 1024] float32, L2-normalized

# queries (Qwen3-Embedding instruction convention)
q_emb = model.embed(["pleural effusion"], tokenizer=tok,
                    instruction="Retrieve chest CT reports relevant to the query")

# cosine similarity = dot product (already normalized)
sims = q_emb @ doc_emb.T

model(input_ids, attention_mask) returns pooler_output (the normalized embedding); the model.embed(...) helper handles tokenization, the EOS/last-token pooling, and batching.

Details

• Base: Qwen/Qwen3-Embedding-0.6B (Apache-2.0) — architecture rebuilt from the bundled config; merged weights loaded from this repo.
• Embedding dim: 1024 · max length: 512 · pooling: last-token (EOS) + L2-norm.
• Precision: bf16 weights. Default attention sdpa (set attn_implementation="flash_attention_2" for speed on Ampere+).
• Merged weights reproduce the original adapter-based embeddings to cosine ≥ 0.9996 (bf16-rounding only).

Matryoshka embeddings

This model was Matryoshka (MRL)-trained (mrl_dims = 512, 256), so the 1024-d embedding can be truncated to 512 or 256 dimensions with minimal quality loss — just keep the first N dims and re-normalize:

emb_512 = model.embed(docs, tokenizer=tok, dim=512)   # [N, 512], L2-normalized
emb_256 = model.embed(docs, tokenizer=tok, dim=256)   # [N, 256]

Use the same dim for queries and corpus. Recommended dims: 1024 (full) · 512 · 256 (config.matryoshka_dims).

Recommended instructions

The model is instruction-conditioned: prepend a task instruction as Instruct: {instruction}\nQuery: {report} (handled by model.embed(..., instruction=...)). It was trained on chest CT and CXR reports across the task families below — use the matching string. Convention: apply the instruction to the query side; embed the corpus/documents without it.

1. Retrieval (report → similar report)

2. Summarization-oriented embedding

3. Entity extraction / classification (full leaf taxonomy)

4. Entity extraction — upper / coarse class

5. Anatomy-specific extraction — template: 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

# task-conditioned embedding
q = model.embed(["pleural effusion and cardiomegaly"], tokenizer=tok,
                instruction="Retrieve the chest CT report that is similar to the given report.")
lungs = model.embed([report], tokenizer=tok,
                    instruction="From the following chest CT report, extract and return only the findings related to the lungs, ignoring all information about other structures.")

For supervised classification, embed reports under the relevant instruction (e.g. the extraction one) and train a linear head on the embeddings.

Legacy loader (still available)

The original delta-weights layout is retained for backward compatibility: chest2vec.py + chest2vec_config.json + contrastive/ (LoRA adapter) + section_pooler.pt. That path uses the chest2vec package (Chest2Vec.from_pretrained) and downloads the base Qwen3-Embedding weights at runtime. New users should prefer the self-contained AutoModel path above.