ctrltokyo/Reason-Code-ModernColBERT

1---2license: apache-2.03language:4- en5- code6library_name: PyLate7tags:8- ColBERT9- PyLate10- sentence-transformers11- code-search12- code-retrieval13- late-interaction14- reasoning15base_model: lightonai/GTE-ModernColBERT-v116datasets:17- nomic-ai/cornstack-python-v118- nomic-ai/cornstack-java-v119- nomic-ai/cornstack-javascript-v120- nomic-ai/cornstack-php-v121- nomic-ai/cornstack-go-v122- nomic-ai/cornstack-ruby-v123pipeline_tag: sentence-similarity24---25 26# Reason-Code-ModernColBERT27 28The **first reasoning-enhanced ColBERT model for code search and retrieval**.29 30Extends 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).31 32## Why Reasoning-Enhanced Training for Code?33 34Standard 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.35 36## Model Details37 38| Property | Value |39|---|---|40| **Base model** | [lightonai/GTE-ModernColBERT-v1](https://huggingface.co/lightonai/GTE-ModernColBERT-v1) |41| **Architecture** | ColBERT (late-interaction, multi-vector) |42| **Parameters** | 150M |43| **Embedding dim** | 128 per token |44| **Document length** | 512 tokens |45| **Query length** | 128 tokens |46| **Similarity** | MaxSim |47| **Languages** | Python, Java, JavaScript, PHP, Go, Ruby |48| **License** | Apache 2.0 |49 50## Training51 52### Two-Stage Training Pipeline53 54**Stage 1: CoRNStack Base (1 epoch)**55- 100,000 high-quality code search pairs from [CoRNStack](https://huggingface.co/collections/nomic-ai/cornstack-67c60fda17322ce742fe9dac) (Apache 2.0)56- 6 languages: Python (25K), Java (20K), JavaScript (15K), PHP (15K), Go (15K), Ruby (10K)57- Loss: 2.42 → 0.6358 59**Stage 2: Reasoning-Enhanced Fine-Tuning (3 epochs)**60- 9,959 reasoning-intensive code search queries generated from CoRNStack code samples61- Queries require understanding algorithms, edge cases, design patterns, and complexity62- Each query includes a chain-of-thought reasoning prefix (ReasonIR methodology)63- Loss: 2.36 → 0.5464 65### Training Configuration66 67```python68# Both stages69model = ColBERT(document_length=512, query_length=128)70loss = CachedContrastive(temperature=1.0, mini_batch_size=32)71batch_size = 25672optim = "adamw_torch"73bf16 = True74 75# Stage 1: lr=1e-5, 1 epoch, warmup=5%76# Stage 2: lr=5e-6, 3 epochs, warmup=5%77```78 79### Hardware80 81Trained on a single NVIDIA DGX Spark (GB10 Blackwell, 128GB unified memory).82- Stage 1: ~130 min (391 steps)83- Stage 2: ~37 min (117 steps)84 85## Benchmark Results86 87### CodeSearchNet MRR (500 queries per language, 500 candidates)88 89| Language   | GTE-ModernColBERT (base) | **Reason-Code-ModernColBERT (ours)** | Δ |90|------------|:---:|:---:|:---:|91| Python     | 0.991 | 0.989 | -0.002 |92| Java       | 0.829 | **0.866** | +0.037 |93| JavaScript | 0.802 | **0.839** | +0.037 |94| PHP        | 0.841 | **0.862** | +0.021 |95| Go         | 0.879 | **0.887** | +0.008 |96| Ruby       | 0.773 | **0.831** | +0.058 |97| **Average** | 0.853 | **0.879** | **+0.026** |98 99Improves 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.100 101## Usage102 103```python104from pylate import models105 106model = models.ColBERT(model_name_or_path="ctrltokyo/Reason-Code-ModernColBERT")107 108queries = ["function that sorts an array in descending order using a comparison-based algorithm"]109code_docs = ["def sort_desc(arr):\n    return sorted(arr, reverse=True)"]110 111query_embeddings = model.encode(queries, is_query=True)112doc_embeddings = model.encode(code_docs, is_query=False)113```114 115## Citation116 117This model extends the methodology from:118 119```bibtex120@article{shao2025reasonir,121  title={ReasonIR: Training Retrievers for Reasoning Tasks},122  author={Shao, Rulin and Jiang, Rui and Yu, Tao and Hashimoto, Tatsunori},123  journal={arXiv preprint arXiv:2504.20595},124  year={2025}125}126 127@misc{Reason-ModernColBERT,128  title={Reason-ModernColBERT},129  author={LightOn AI},130  year={2025},131  url={https://huggingface.co/lightonai/Reason-ModernColBERT}132}133 134@inproceedings{cornstack2025,135  title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking},136  author={Gangisetty, Zach and others},137  booktitle={ICLR},138  year={2025}139}140```141 142Built with [PyLate](https://github.com/lightonai/pylate) and [Sentence Transformers](https://www.sbert.net/).143