embeddinggemma-300m-rag-stage1
Fine-tuned version of google/embeddinggemma-300m on the EnterpriseRAG-Bench question-document pairs for first-place retrieval.
Training
| Base model | google/embeddinggemma-300m |
| Training pairs | 200,000 (one question per unique doc_id, joined from shanaka95/enterprise-rag-questions + onyx-dot-app/EnterpriseRAG-Bench) |
| Loss | MultipleNegativesRankingLoss (scale=20.0) wrapped in MatryoshkaLoss(dims=[768,512,256,128]) |
| Optimizer | adamw_torch_fused, lr=2e-5, cosine, warmup_ratio=0.1 |
| Batch | 4, bf16, gradient checkpointing |
| Steps | 8000 (~2.5h on RTX 5060 Ti 16GB) |
| Context | 2048 tokens (EmbeddingGemma native) |
| Prompt format | `task: search result |
Evaluation (89 questions, 5k corpus subset, 2048 ctx)
Compared against google/embeddinggemma-300m baseline.
| Metric | Stage 1 (this model) | Baseline | Improvement |
|---|---|---|---|
| Recall@1 | 0.3933 | 0.1750 | +21.83pp (2.25×) |
| Recall@5 | 0.6067 | 0.2350 | +37.17pp |
| Recall@10 | 0.6854 | 0.2600 | +42.54pp |
| Recall@100 | 0.8764 | 0.3050 | +57.14pp |
| MRR | 0.5046 | 0.2063 | +0.298 |
Eval setup: 89 of 470 EnterpriseRAG-Bench questions whose first expected_doc_id 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
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
- Queries: shanaka95/enterprise-rag-questions
- Documents: onyx-dot-app/EnterpriseRAG-Bench —
documentssplit - Eval: same
questionssplit, filtered to those with at least oneexpected_doc_idin the docs corpus (470 questions), using the first expected_doc_id as gold.
Notes
- Use prompt names
queryanddocument(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.