CoolFace
Modelpublic

umakantcurateai/qwen_embedding_4b_disease_drug_relationship_similarity

sourceHugging Faceapache-2.0updated 29d agoView on Hugging Face
0likes71downloads
Model Card

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"

CandidateBefore fine-tuningAfter fine-tuning
Erlotinib hydrochloride (treats lung cancer)0.59830.5973
Gemcitabine hydrochloride (treats lung cancer)0.57870.4176
Water (unrelated)0.47160.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

python
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