Code Search Engine

Embeddings-based semantic code search engine using sentence transformers, evaluated on the CoSQA benchmark dataset.

Architecture

• Embedding Model: sentence-transformers/all-MiniLM-L6-v2 (384-dimensional)
• Vector Store: USearch with cosine similarity
• Fine-tuning: MultipleNegativesRankingLoss

Usage

from src.search_engine import CodeSearchEngine

# Initialize search engine
engine = CodeSearchEngine(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    metric="cos"
)

# Index code snippets
documents = [
    "def binary_search(arr, target): ...",
    "def bubble_sort(arr): ...",
    # ... more code snippets
]
engine.index_documents(documents)

# Search
results = engine.search("search in sorted array", top_k=5)
for result in results:
    print(f"Distance: {result['distance']:.4f}")
    print(result['document'])

Evaluation

Evaluated on CoSQA dataset using standard information retrieval metrics:

Results (CoSQA Test Set)

Metric	Baseline	Fine-tuned (3 epochs)	Improvement
Recall@10	0.560	0.624	+11.4%
MRR@10	0.307	0.333	+8.4%
NDCG@10	0.366	0.402	+9.6%

Project Structure

src/
    search_engine.py      # Core search engine implementation
    evaluation.py          # Evaluation metrics (Recall, MRR, NDCG)
    finetuning.py          # Fine-tuning module
    utils.py               # Dataset loading utilities
notebooks/
    code_search_demo.ipynb # Demo and evaluation notebook
models/                    # Saved model checkpoints
data/                      # Dataset cache
requirements.txt

Demo

See notebooks/code_search_demo.ipynb for usage examples, evaluation, and fine-tuning.