Qwen3-Embedding-4B fine-tuned on Disease-Drug pairs
This is a fine-tuned version of unsloth/Qwen3-Embedding-4B, adapted to better capture disease-to-drug therapeutic relationships. It was trained using Unsloth with LoRA adapters merged back into the base weights, so this repo contains a complete, standalone model, no separate adapter loading required.
What this model does
Given a disease description as a query, this model produces embeddings where drugs that treat that disease are pulled closer in vector space, while unrelated drugs are pushed further away. It was built as a small, educational fine-tuning exercise, not an industry-grade clinical tool, to demonstrate how domain-specific embedding fine-tuning works end to end.
Training data
The training set was built from the DISEASE and DRUG databases, using each disease's DESCRIPTION field (combined with its name) as the anchor, and its listed therapeutic drugs as positives.
Column breakdown:
- anchor: Disease name plus description
- positive: Name of a drug indicated for that disease
The dataset contains 2,950 anchor-positive pairs, deduplicated, and split by disease (not by row) to prevent the same disease appearing in both train and eval sets.
Training details
- Base model: unsloth/Qwen3-Embedding-4B
- Method: LoRA fine-tuning via Unsloth (r=32, alpha=32, target modules: q/k/v/o/gate/up/down projections), merged to 16-bit for this release
- Loss function: MultipleNegativesRankingLoss, treats all other positives in a training batch as in-batch negatives for a given anchor
- Batch sampling: NO_DUPLICATES, to prevent the same disease anchor (which can have multiple associated drugs) from appearing twice in one batch and being wrongly treated as a negative for itself
- Epochs: 4
- Effective batch size: 32 (per-device batch size 8 times gradient accumulation 4)
- Learning rate: 3e-5, constant with warmup
Before vs after fine-tuning
Example query: "lung cancer"
| Candidate | Before fine-tuning | After fine-tuning |
|---|---|---|
| Erlotinib hydrochloride (treats lung cancer) | 0.5983 | 0.5973 |
| Gemcitabine hydrochloride (treats lung cancer) | 0.5787 | 0.4176 |
| Water (unrelated) | 0.4716 | 0.1344 |
Fine-tuning clearly pushed the unrelated candidate (Water) much further away, confirming the model learned to separate irrelevant drugs from disease-relevant ones. Note that not all correct positives moved in the same direction (Gemcitabine's similarity dropped), this is an honest artifact of training on a small dataset for a limited number of epochs, and is discussed further in the accompanying Medium article as a teaching point about the practical realities of small-scale contrastive fine-tuning.
Usage
from sentence_transformers import SentenceTransformer
from sentence_transformers import util
model = SentenceTransformer("umakantcurateai/qwen_embedding_4b_disease_drug_relationship_similarity")
query = "lung cancer"
candidates = ["Erlotinib hydrochloride", "Gemcitabine hydrochloride", "Water (JP19/USP)"]
query_emb = model.encode(query)
candidate_embs = model.encode(candidates)
scores = util.cos_sim(query_emb, candidate_embs)
print(scores)
Limitations
- Trained on a small, curated subset of disease-drug pairs (2,950 pairs), not comprehensive across all diseases or drugs in .
- Built for educational and demonstration purposes; not validated for clinical or diagnostic use.
- Some positive pairs may show reduced similarity post-training due to limited training data and epochs; see the before/after table above.
Acknowledgements
- Base model: Qwen3-Embedding-4B by Unsloth and the Qwen team
- Training data sourced from DISEASE DRUG DATABASE
- Fine-tuned using Unsloth and Sentence Transformers