minishlab / potion-code-16M-v2

potion-code-16M-v2 Model Card

Overview

potion-code-16M-v2 is a fast static code embedding model optimized for code retrieval tasks. It powers Semble , a code search library for agents. It is distilled from nomic-ai/CodeRankEmbed and trained on the CornStack code corpus using Tokenlearn and contrastive fine-tuning. It is the successor to potion-code-16M . It uses static embeddings, allowing text and code embeddings to be computed orders of magnitude faster than transformer-based models on both GPU and CPU.

Installation

pip install model2vec

Usage

from model2vec import StaticModel

model = StaticModel.from_pretrained("minishlab/potion-code-16M-v2")

# Embed natural language queries
query_embeddings = model.encode(["How to read a file in Python?"])

# Embed code documents
code_embeddings = model.encode(["def read_file(path):\n    with open(path) as f:\n        return f.read()"])

How it works

potion-code-16M-v2 is created using the following pipeline:

1. Vocabulary mining : code-specific tokens are mined from CornStack and added to the base CodeRankEmbed tokenizer (43k extra tokens → ~63.5k total)
2. Distillation : the extended vocabulary is distilled from CodeRankEmbed using Model2Vec (256-dimensional embeddings, PCA)
3. Tokenlearn : the distilled model is fine-tuned on 1.2 million (query, document) pairs from CornStack using cosine similarity loss
4. Contrastive fine-tuning : the model is further fine-tuned using MultipleNegativesRankingLoss on 1.2 million CornStack query-document pairs

Results

Results on the CoIR benchmark on MTEB (NDCG@10, mteb>=2.10 ):

CoIR covers a broad range of code retrieval scenarios. For the use case of finding code given a natural language query, CosQA and CodeFeedback (ST/MT) are the most relevant tasks. Others are less so: COIRCodeSearchNetRetrieval retrieves text given a code query (the reverse direction), and the CodeTransOcean tasks target cross-language code translation. The hybrid row combines dense retrieval with BM25 using min-max score normalization and equal weighting (alpha=0.5).

Model Details

Additional Resources

• Semble repository
• Model2Vec repository
• Tokenlearn repository
• Tokenlearn document dataset
• Tokenlearn query dataset
• CornStack dataset
• CoIR benchmark

Citation

@software{minishlab2024model2vec,
  author       = {Stephan Tulkens and {van Dongen}, Thomas},
  title        = {Model2Vec: Fast State-of-the-Art Static Embeddings},
  year         = {2024},
  publisher    = {Zenodo},
  doi          = {10.5281/zenodo.17270888},
  url          = {https://github.com/MinishLab/model2vec},
  license      = {MIT}
}