codebert-permissive-embed

Search code by describing what it does.

This model turns a piece of code, or a sentence about code, into a list of 768 numbers called a vector. Similar meanings get similar vectors. That lets you ask "where do we retry a failed request?" and get back the right function, even when the code never uses the word "retry".

Produced by ThinkingDBx Pvt. Ltd.

Quick start

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("thinkingdbx/codebert-permissive-embed")

query = model.encode("retry an http request with exponential backoff")
code = model.encode([
    "def fetch(url, tries=5):\n    for i in range(tries):\n        try:\n            return requests.get(url)\n        except Exception:\n            time.sleep(2 ** i)",
    "def parse_csv(path):\n    return list(csv.DictReader(open(path)))",
])

print(query @ code.T)     # [ 0.389, -0.134 ]

The first score is much higher, so the retry function wins. That is the whole idea.

If you prefer plain transformers, you have to average the token vectors yourself and then normalise them. The model was trained that way, and doing something else quietly gives worse results:

import torch
from transformers import AutoModel, AutoTokenizer

name = "thinkingdbx/codebert-permissive-embed"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModel.from_pretrained(name).eval()

def embed(texts, max_length=256):
    batch = tok(texts, padding=True, truncation=True,
                max_length=max_length, return_tensors="pt")
    with torch.no_grad():
        hidden = model(**batch).last_hidden_state
    mask = batch["attention_mask"].unsqueeze(-1).float()
    pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
    return torch.nn.functional.normalize(pooled, dim=-1)

What you can use it for

• Code search in plain English, across a repository or a whole company codebase.
• Finding duplicate or near duplicate functions.
• Grouping an unfamiliar codebase into related areas.
• The retrieval step in a code assistant, where you need to pull the right files before asking a language model about them.

It only produces vectors. It cannot write code or answer questions.

What it is not good at

It is a small model. It has 110 million parameters and it learned from about 857,000 examples. Popular general purpose models of the same size, such as E5, BGE and GTE, learned from hundreds of millions of examples covering both ordinary text and code, and they score better than this model on most tests.

It is also weak at matching code against other code, for example finding the Java version of a Python function. It was never trained to do that. It was trained on descriptions paired with code, so descriptions paired with code is what it does well.

If you want the best available code search quality and you do not care where the training data came from, use one of the larger models instead.

Where the training data came from

This is the part most models cannot tell you, and it is the main reason this one exists.

Every file used to train it carried a permissive open source licence such as MIT, Apache 2.0 or BSD. Files with unclear or restrictive licences were left out.

• 2,047,089 source files were examined. 1,980,241 were kept, which is 96.7 percent.
• 171 licences were accepted. 139 were rejected.
• No file with a missing licence was used.
• 1,384,479 pieces of personal data, such as email addresses and keys, were found and removed before any training started. The model read the cleaned copy.

Training happened in two steps. First the model learned the shape of code in general, by reading 3.97 billion words of code and guessing hidden pieces. Then it learned to match descriptions to code, using 857,000 pairs of a function and its own documentation.

The files lineage.json, benchmark.json and retrieval.json in this repository hold the raw numbers behind every claim on this page.

How well it works

Scores on CoIR, a public benchmark for code search. Higher is better. The measure is NDCG@10, which roughly means "how often the right answer appears near the top of the first ten results".

Test	This model	Earlier version	BM25 (keyword search)	UniXcoder	GTE-Base	E5-Base
CodeSearchNet, Python	87.84
CodeSearchNet, Go	68.71	53.43
CodeSearchNet, Ruby	56.52	40.68
Stack Overflow questions	55.35	58.40	56.80	44.67	62.71	86.86
CodeSearchNet, PHP	53.36
CodeSearchNet, JavaScript	50.41	38.73
Code feedback, single turn	48.58
Text to SQL	35.73
CoSQA	25.95	20.91	13.96	25.14	30.24	32.59
Code translation	24.12	28.14	50.13	41.82	33.81	62.50
Code feedback, multi turn	21.78
Programming problems	2.83	3.08

A few notes on reading this table honestly.

Blank cells mean the number was not measured, not that it was zero. The published paper reports CodeSearchNet as one combined score across six languages, while the scores here are per language, so those cells are left empty rather than compared to something different.

There is no average score on this page. Two of the CoIR tests were not run, and an average over some of the tests is not the same as the published average. Putting one next to the other would be misleading.

The comparison figures for BM25, UniXcoder, GTE-Base and E5-Base come from Table 3 of the CoIR paper, arXiv:2407.02883. The "earlier version" column is a previous build of this model, measured on the same machine with the same code.

On a separate test of 200 questions against 2,200 documents, this model scores 0.268 where plain keyword search scores 0.255. That is only slightly ahead. Keyword search is free and needs no hardware, so it is a fair thing to measure against, and on short questions over long documents it remains hard to beat.

Limitations

• Inputs longer than 256 tokens are cut off. Long files should be split into functions first.
• Trained mostly on Python, Java, JavaScript, Go, PHP and Ruby. Other languages will work less well.
• Not tested for bias, safety or licence detection. It is a search tool, not a judge.
• The scores above are the honest ones, including the tests where it loses.

Licence

Apache 2.0. The weights are free to use, including commercially.

The training data was permissively licensed throughout, and the record of what was included and excluded ships with the model.