🔍 Semantic Code Search Engine

A Test Assignment for JetBrains Internship

This project implements an embeddings-based code search engine, evaluates its performance on the CoSQA dataset using standard retrieval metrics (Recall@10, MRR@10, nDCG@10), and demonstrates improvement through fine-tuning on code-query pairs.

📁 Project Structure

├── config.py                   # Project configuration: paths, models, and hyperparameters  
├── src/                              
│   ├── search_engine.py        # SearchEngine class    
│   ├── data_loader.py          # Loads and preprocesses CoSQA dataset    
│   ├── evaluation.py           # Metrics and model evaluation    
│   ├── train.py                # Training logic     
│   └── bonus_experiments.py    # Bonus analyses (function names vs. bodies)    
│    
├── scripts/                        
│   ├── index_and_search_demo.py    # Demo on code snippets    
│   ├── evaluate_baseline.py        # Evaluate base model on CoSQ    
│   ├── fine_tune_and_evaluate.py   # Fine-tune + compare metrics    
│   └── run_bonus_experiments.py    # Run bonus analyses    
│    
├── report.ipynb                # Complete implementation with results, plots, and analysis    
├── requirements.txt            # Dependencies    
└── README.md                   # This file: setup, structure, and usage instructions  

🚀 Quick Start

1. Install dependencies

pip install -r requirements.txt

2. Run the semantic search demo

python scripts/index_and_search_demo.py

3. Evaluate the baseline model on CoSQA

python scripts/evaluate_baseline.py --k 10

4. Fine-tune the model and compare metrics

python scripts/fine_tune_and_evaluate.py --epochs 3

5. Run bonus experiments

python scripts/run_bonus_experiments.py

📊 Evaluation Metrics

The system is evaluated using standard information retrieval metrics on the CoSQA test set:

Recall@10: Proportion of queries with at least one relevant result in the top 10.
MRR@10 (Mean Reciprocal Rank): Average of the reciprocal rank of the first relevant result (capped at 10).
nDCG@10 (Normalized Discounted Cumulative Gain): Measures ranking quality considering graded relevance.

🔧 Fine-Tuning

I fine-tune the base embedding model on the CoSQA training set to specialize it for code search.

Loss Function: MultipleNegativesRankingLoss — chosen because it efficiently leverages in-batch negative sampling: for each natural language query, all non-matching code snippets in the same batch are treated as negatives. This directly optimizes the model to rank the correct code snippet higher than irrelevant ones, which aligns perfectly with semantic code search objectives. The loss is widely used in state-of-the-art retrieval systems and requires no manual negative creation.
Training Progress: Training loss vs. step is logged and plotted in the notebook.
Key Result: Fine-tuning consistently improves all three metrics (see report).

🌟 Bonus Experiments

The notebook also includes additional studies:

1. Function Names vs. Full Bodies:
Using only function names (instead of full bodies) as documents reduces performance, confirming that semantic context in code bodies is critical.

📓 Report Notebook

Open report.ipynb to see:

• End-to-end demo of indexing and querying
• Baseline model performace
• Fine-tuning results
• Training loss curve
• Bonus experiment findings
• Summary

🛠️ Technologies Used

• Embedding Model: BAAI/bge-base-en-v1.5 — a strong general-purpose encoder for semantic retrieval
• Embedding Framework: sentence-transformers
• Vector Index: FAISS — efficient similarity search over dense vectors
• Loss: MultipleNegativesRankingLoss for contrastive fine-tuning
• Data: CoSQA dataset — query–code pairs for code search evaluation

✅ What’s Implemented

Part	Status
Part 1: Search API	✅ Complete
Part 2: Evaluation	✅ All metrics
Part 3: Fine-tuning	✅ With results
Bonus	✅ Bonus 1 explored

🙏 Acknowledgments

This project builds on the following open-source resources:

• CoSQA dataset: https://huggingface.co/datasets/CoIR-Retrieval/cosqa
• BAAI/bge-base-en-v1.5 embedding model: https://huggingface.co/BAAI/bge-base-en-v1.5
• Sentence-Transformers library: https://www.sbert.net
• FAISS (Facebook AI Similarity Search): https://github.com/facebookresearch/faiss