jasperan/embeddinggemma-code-search

1---2license: gemma3base_model: google/embeddinggemma-300m4pipeline_tag: sentence-similarity5library_name: sentence-transformers6tags:7- sentence-transformers8- sentence-similarity9- feature-extraction10- embeddinggemma11- code-search12- matryoshka13- fine-tuned14---15 16# embeddinggemma-code-search17 18[`google/embeddinggemma-300m`](https://huggingface.co/google/embeddinggemma-300m) fine-tuned for **code search over a real agent codebase**, using graded retrieval traces from coding-agent sessions.19 20This 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).21 22## Training setup23 24- **Base model**: `google/embeddinggemma-300m` (sentence-transformers, Matryoshka-truncatable embeddings)25- **Corpus**: 240 code chunks extracted from a real Python agent-harness source tree (signatures, docstrings, AST body summaries, import/call/co-edit context)26- **Traces**: 800 graded retrieval traces (573 train / 227 held-out), 40 candidates per trace, indirect-intent natural-language queries27- **Loss**: listwise KL divergence between softmax(similarities) and the normalized grade distribution28- **Schedule**: deliberately gentle — 1 epoch, lr 5e-6, full-parameter. (Aggressive schedules overwrite the pretrained space and hurt held-out recall.)29 30## Results (held-out queries, recall@5)31 32| Embedding dim (MRL) | Base | Fine-tuned |33|---|---|---|34| 768 | 0.753 | **0.793** |35| 256 | 0.643 | **0.753** |36| 128 | 0.568 | **0.722** |37 38The 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.39 40## Usage41 42```python43from sentence_transformers import SentenceTransformer44 45model = SentenceTransformer("jasperan/embeddinggemma-code-search")46 47query_emb = model.encode(["where do we retry failed tool calls?"])48doc_embs = model.encode(code_chunks)49 50# Matryoshka: truncate + re-normalize for a smaller index51import numpy as np52q128 = query_emb[:, :128]53q128 = q128 / np.linalg.norm(q128, axis=1, keepdims=True)54```55 56## License57 58EmbeddingGemma is provided under and subject to the [Gemma Terms of Use](https://ai.google.dev/gemma/terms). This fine-tune inherits those terms.59