nomic-ai / CodeRankEmbed

CodeRankEmbed

CodeRankEmbed is a 137M bi-encoder supporting 8192 context length for code retrieval. It significantly outperforms various open-source and proprietary code embedding models on various code retrieval tasks.

Check out our blog post and paper for more details!

Combine CodeRankEmbed with our re-ranker CodeRankLLM for even higher quality code retrieval.

Performance Benchmarks

We release the scripts to evaluate our model's performance here .

Usage

Important : the query prompt must include the following task instruction prefix : "Represent this query for searching relevant code"

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("nomic-ai/CodeRankEmbed", trust_remote_code=True)
queries = ['Represent this query for searching relevant code: Calculate the n-th factorial']
codes = ['def fact(n):\n if n < 0:\n  raise ValueError\n return 1 if n == 0 else n * fact(n - 1)']
query_embeddings = model.encode(queries)
print(query_embeddings)
code_embeddings = model.encode(codes)
print(code_embeddings)

Training

We use a bi-encoder architecture for CodeRankEmbed , with weights shared between the text and code encoder. The retriever is contrastively fine-tuned with InfoNCE loss on a 21 million example high-quality dataset we curated called CoRNStack . Our encoder is initialized with Arctic-Embed-M-Long , a 137M parameter text encoder supporting an extended context length of 8,192 tokens.