ide-code-retrieval-qwen3-0.6b-cochange-mask-hard

A SentenceTransformer model fine-tuned from Qwen/Qwen3-Embedding-0.6B for code→code retrieval -- mapping a source file to files that co-change with it, via dense vector similarity.

Note: This is an intermediate checkpoint at step 6,760 / 6,760 (100.0% through 10 epochs). Training loss is still decreasing, so a later checkpoint may perform better.

Model Description

This model encodes source files into a shared embedding space so that files that historically co-change land near each other. Retrieval is performed by computing cosine similarity between a query-file embedding and candidate embeddings.

• Base model: Qwen/Qwen3-Embedding-0.6B (0.6B parameters)
• Max sequence length: 1024 tokens
• Output dimensionality: 1024 (normalized)
• Similarity function: Cosine similarity

Training Details

Dataset

• Source: aysinghal/code-retrieval-training-dataset
• Query rows: 173,093 train / 28,210 test (repo-level held-out split; no repo appears in both).
• Corpus: 477,864 files, referenced by ID from rows.
• Text handling: collator-side sliding window
• Negatives: Up to 128 mined hard + 128 mined easy negatives per query row, plus group-aware in-batch negatives.
• Pre-tokenized: No (text resolved from corpus in the collator)

Loss Function

MultipleNegativesRankingLoss (InfoNCE) with mined hard/easy negatives and group-aware in-batch negatives (batches are constructed so no two rows share a co-change group_key, preventing false negatives when two rows come from the same cluster).

Hyperparameters

Hardware

• GPUs: 4x NVIDIA L40S
• Total training steps: 6,760 (10 epochs)

Training Progress (at checkpoint step 6,760)

• Training loss: 3.1436 (step 50) → 0.1502 (step 6750)
• Best eval loss: 0.5257 (step 1,360)
• Progress: 6,760 / 6,760 steps (100.0%)

Evaluation Results

Usage

Loading the Model

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("aysinghal/ide-code-retrieval-qwen3-0.6b-cochange-mask-hard")

Computing Embeddings

query_file = open("src/auth/session.py").read()
candidates = [open(p).read() for p in candidate_paths]

query_emb = model.encode([query_file])
cand_embs = model.encode(candidates)

# Compute cosine similarities
from sentence_transformers.util import cos_sim
similarities = cos_sim(query_emb, cand_embs)
print(similarities)

Intended Use

• Primary use case: Given a file the developer is currently editing, retrieve related files across the same codebase that are likely to require changes.
• Search pipeline: Encode a corpus of code documents offline, then at query time encode the query and find nearest neighbors via cosine similarity

Limitations

• This is an early checkpoint (100.0% through training). The loss curve is still decreasing, so later checkpoints will likely perform better.
• Trained on a specific code retrieval dataset; may not generalize to all programming languages or query styles without further fine-tuning.
• Max context is 1024 tokens -- very long files are truncated.

Citation

If you use this model, please cite the base model:

@article{qwen3embedding,
  title={Qwen3-Embedding},
  author={Qwen Team},
  year={2025}
}