embeddinggemma-code-search

google/embeddinggemma-300m fine-tuned for code search over a real agent codebase, using graded retrieval traces from coding-agent sessions.

This is the "custom retrieval model from agent sessions" arm of a continual-learning course: instead of hand-labeled pairs, training signal comes from graded traces — for each natural-language query about the codebase, candidate code chunks carry relevance grades (0–3) derived from what the agent actually needed. The encoder is trained so that its softmax similarity distribution over the candidates matches the grade distribution (a listwise KL-divergence loss, in the spirit of Cursor's "align the embedding space to what sessions proved relevant" approach).

Training setup

• Base model: google/embeddinggemma-300m (sentence-transformers, Matryoshka-truncatable embeddings)
• Corpus: 240 code chunks extracted from a real Python agent-harness source tree (signatures, docstrings, AST body summaries, import/call/co-edit context)
• Traces: 800 graded retrieval traces (573 train / 227 held-out), 40 candidates per trace, indirect-intent natural-language queries
• Loss: listwise KL divergence between softmax(similarities) and the normalized grade distribution
• Schedule: deliberately gentle — 1 epoch, lr 5e-6, full-parameter. (Aggressive schedules overwrite the pretrained space and hurt held-out recall.)

Results (held-out queries, recall@5)

Embedding dim (MRL)	Base	Fine-tuned
768	0.753	0.793
256	0.643	0.753
128	0.568	0.722

The biggest lift is at truncated Matryoshka dimensions — the fine-tuned 128-d embeddings match the base model's 256-d quality, which is what you want for cheap, low-latency code search indexes.

Usage

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("jasperan/embeddinggemma-code-search")

query_emb = model.encode(["where do we retry failed tool calls?"])
doc_embs = model.encode(code_chunks)

# Matryoshka: truncate + re-normalize for a smaller index
import numpy as np
q128 = query_emb[:, :128]
q128 = q128 / np.linalg.norm(q128, axis=1, keepdims=True)

License

EmbeddingGemma is provided under and subject to the Gemma Terms of Use. This fine-tune inherits those terms.