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 — documents split
• Eval: same questions split, filtered to those with at least one expected_doc_id in the docs corpus (470 questions), using the first expected_doc_id as gold.

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.