README.md

---
license: apache-2.0
language:
- en
- code
library_name: PyLate
tags:
- ColBERT
- PyLate
- sentence-transformers
- code-search
- code-retrieval
- late-interaction
- reasoning
base_model: lightonai/GTE-ModernColBERT-v1
datasets:
- nomic-ai/cornstack-python-v1
- nomic-ai/cornstack-java-v1
- nomic-ai/cornstack-javascript-v1
- nomic-ai/cornstack-php-v1
- nomic-ai/cornstack-go-v1
- nomic-ai/cornstack-ruby-v1
pipeline_tag: sentence-similarity
---

# Reason-Code-ModernColBERT

The **first reasoning-enhanced ColBERT model for code search and retrieval**.

Extends the [ReasonIR methodology](https://arxiv.org/abs/2504.20595) to the code domain — generating reasoning-intensive queries that require understanding algorithms, edge cases, and design patterns, not just keyword matching. Built on research from [LightOn AI](https://huggingface.co/lightonai) (ColBERT for code) and [Facebook Research](https://github.com/facebookresearch/ReasonIR) (reasoning-enhanced retrieval).

## Why Reasoning-Enhanced Training for Code?

Standard code search training uses docstring→code pairs. Our approach generates **reasoning-intensive queries** that require understanding the code's algorithm, behavior, and edge cases — not just surface-level keyword matching. This is the same methodology that enabled [Reason-ModernColBERT](https://huggingface.co/lightonai/Reason-ModernColBERT) to outperform 7B dense models on reasoning tasks at only 150M parameters.

## Model Details

| Property | Value |
|---|---|
| **Base model** | [lightonai/GTE-ModernColBERT-v1](https://huggingface.co/lightonai/GTE-ModernColBERT-v1) |
| **Architecture** | ColBERT (late-interaction, multi-vector) |
| **Parameters** | 150M |
| **Embedding dim** | 128 per token |
| **Document length** | 512 tokens |
| **Query length** | 128 tokens |
| **Similarity** | MaxSim |
| **Languages** | Python, Java, JavaScript, PHP, Go, Ruby |
| **License** | Apache 2.0 |

## Training

### Two-Stage Training Pipeline

**Stage 1: CoRNStack Base (1 epoch)**
- 100,000 high-quality code search pairs from [CoRNStack](https://huggingface.co/collections/nomic-ai/cornstack-67c60fda17322ce742fe9dac) (Apache 2.0)
- 6 languages: Python (25K), Java (20K), JavaScript (15K), PHP (15K), Go (15K), Ruby (10K)
- Loss: 2.42 → 0.63

**Stage 2: Reasoning-Enhanced Fine-Tuning (3 epochs)**
- 9,959 reasoning-intensive code search queries generated from CoRNStack code samples
- Queries require understanding algorithms, edge cases, design patterns, and complexity
- Each query includes a chain-of-thought reasoning prefix (ReasonIR methodology)
- Loss: 2.36 → 0.54

### Training Configuration

```python
# Both stages
model = ColBERT(document_length=512, query_length=128)
loss = CachedContrastive(temperature=1.0, mini_batch_size=32)
batch_size = 256
optim = "adamw_torch"
bf16 = True

# Stage 1: lr=1e-5, 1 epoch, warmup=5%
# Stage 2: lr=5e-6, 3 epochs, warmup=5%
```

### Hardware

Trained on a single NVIDIA DGX Spark (GB10 Blackwell, 128GB unified memory).
- Stage 1: ~130 min (391 steps)
- Stage 2: ~37 min (117 steps)

## Benchmark Results

### CodeSearchNet MRR (500 queries per language, 500 candidates)

| Language   | GTE-ModernColBERT (base) | **Reason-Code-ModernColBERT (ours)** | Δ |
|------------|:---:|:---:|:---:|
| Python     | 0.991 | 0.989 | -0.002 |
| Java       | 0.829 | **0.866** | +0.037 |
| JavaScript | 0.802 | **0.839** | +0.037 |
| PHP        | 0.841 | **0.862** | +0.021 |
| Go         | 0.879 | **0.887** | +0.008 |
| Ruby       | 0.773 | **0.831** | +0.058 |
| **Average** | 0.853 | **0.879** | **+0.026** |

Improves on the base model in 5 of 6 languages. Largest gains in Ruby (+5.8pp), Java (+3.7pp), and JavaScript (+3.7pp) — languages that benefited most from reasoning-enhanced training data. Python is near-ceiling at 0.99.

## Usage

```python
from pylate import models

model = models.ColBERT(model_name_or_path="ctrltokyo/Reason-Code-ModernColBERT")

queries = ["function that sorts an array in descending order using a comparison-based algorithm"]
code_docs = ["def sort_desc(arr):\n    return sorted(arr, reverse=True)"]

query_embeddings = model.encode(queries, is_query=True)
doc_embeddings = model.encode(code_docs, is_query=False)
```

## Citation

This model extends the methodology from:

```bibtex
@article{shao2025reasonir,
  title={ReasonIR: Training Retrievers for Reasoning Tasks},
  author={Shao, Rulin and Jiang, Rui and Yu, Tao and Hashimoto, Tatsunori},
  journal={arXiv preprint arXiv:2504.20595},
  year={2025}
}

@misc{Reason-ModernColBERT,
  title={Reason-ModernColBERT},
  author={LightOn AI},
  year={2025},
  url={https://huggingface.co/lightonai/Reason-ModernColBERT}
}

@inproceedings{cornstack2025,
  title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking},
  author={Gangisetty, Zach and others},
  booktitle={ICLR},
  year={2025}
}
```

Built with [PyLate](https://github.com/lightonai/pylate) and [Sentence Transformers](https://www.sbert.net/).