CoolFace
Modelpublic

shanaka95/embeddinggemma-300m-rag-stage1

sourceHugging Facegemmaupdated 2mo agoView on Hugging Face
0likes90downloads
Model Card

embeddinggemma-300m-rag-stage1

Fine-tuned version of google/embeddinggemma-300m on the EnterpriseRAG-Bench question-document pairs for first-place retrieval.

Training

Base modelgoogle/embeddinggemma-300m
Training pairs200,000 (one question per unique doc_id, joined from shanaka95/enterprise-rag-questions + onyx-dot-app/EnterpriseRAG-Bench)
LossMultipleNegativesRankingLoss (scale=20.0) wrapped in MatryoshkaLoss(dims=[768,512,256,128])
Optimizeradamwtorchfused, lr=2e-5, cosine, warmup_ratio=0.1
Batch4, bf16, gradient checkpointing
Steps8000 (~2.5h on RTX 5060 Ti 16GB)
Context2048 tokens (EmbeddingGemma native)
Prompt format`task: search resultquery: {q} / title: nonetext: {d}`

Evaluation (89 questions, 5k corpus subset, 2048 ctx)

Compared against google/embeddinggemma-300m baseline.

MetricStage 1 (this model)BaselineImprovement
Recall@10.39330.1750+21.83pp (2.25×)
Recall@50.60670.2350+37.17pp
Recall@100.68540.2600+42.54pp
Recall@1000.87640.3050+57.14pp
MRR0.50460.2063+0.298

Eval setup: 89 of 470 EnterpriseRAG-Bench questions whose first expecteddocid falls in the first 5k corpus docs, scored against those 5k. Re-encoded corpus with the fine-tuned model (NEVER re-use base-model corpus embeddings — different spaces).

Usage

python
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("shanaka95/embeddinggemma-300m-rag-stage1")
model.max_seq_length = 2048  # EmbeddingGemma native

query_emb = model.encode("How does X work?", prompt_name="query")
doc_emb   = model.encode("X is a method that ...", prompt_name="document")
score = (query_emb @ doc_emb.T)  # cosine (already normalized)

Training data

Notes

  • Use prompt names query and document (matches model.prompts).
  • For best retrieval results, re-encode the corpus with this model (do not reuse base-model embeddings).
  • Stage 2 (hard negative mining) is in progress for further improvements.