Selecting an Embedding Model for Code Search

When you build a code search tool—whether for a codebase assistant, a semantic search bar, or a retrieval-augmented generation pipeline for your IDE—the embedding model you pick determines whether the results are useful or noise. Text embeddings for code behave differently from natural language, and most off-the-shelf models were trained on paragraphs or sentences, not on function bodies with syntax trees and identifiers like getUserById. The goal is to map code snippets into a vector space where similar intent and structure cluster together, so you can retrieve them by cosine similarity. This article walks through the mechanics of evaluating and selecting an embedding model specifically for code search, with concrete comparisons and trade-offs.

What makes a good embedding model for code search

Code differs from natural language in three important ways. First, code has formal syntax. A model that ignores the difference between if (x) and if x will conflate Python and JavaScript idioms. Second, code identifiers carry domain-specific semantics: calculateInterestRate means something very different from getDailyRate, even though both are compound words. Third, code often spans long sequences—a single function might be 200 tokens, and a file can run into the thousands.

The key axes to evaluate are:

• Dimensionality of the output vector. Higher dimensions (e.g., 1536 vs 384) typically capture more nuance but increase storage and retrieval latency. For code, lower dimensions can still work well because the vocabulary is more constrained.
• Domain specificity. Models trained exclusively on natural language (e.g., general sentence transformers) often fail to embed code idioms properly. Models like CodeBERT and UnixCoder are pre-trained on source code and documentation.
• Inference cost. CPU vs GPU latency, model size (parameters), and whether you need batching for indexing or single-query real-time search.

Code search retrieval usually relies on cosine similarity between the query embedding and precomputed embeddings of code chunks. If the model treats two syntactically different but semantically equivalent code snippets as far apart, recall suffers. The model must also handle long sequences: many models cap at 512 or 1024 tokens, so truncating a 1500-token function loses context—consider chunking or models with longer context windows.

Comparing open-source and proprietary embedding models

Here are the most commonly used options, listed with their relevant properties:

Proprietary models (OpenAI, Cohere) require no infrastructure and handle long sequences well. OpenAI's text-embedding-3-small supports up to 8191 tokens, making it ideal for embedding whole functions without chunking. The trade-off: per-request cost adds up at scale, and data privacy is a concern—your code must be sent to an external API. For many teams, this rules out proprietary models entirely.

Open-source models let you self-host, which solves privacy and cost-at-scale. However, they impose operational overhead: you need to run a GPU for acceptable latency (CPU inference for CodeBERT is ~200ms per query) and manage model serving infrastructure. UnixCoder and CodeBERT both produce 768-dim vectors, a sweet spot for retrieval. They are trained on multiple programming languages, but check the original paper for language coverage—CodeBERT covers Python, Java, JavaScript, Ruby, Go, and PHP, while UnixCoder adds several more.

all-MiniLM-L6-v2 is a tiny general model (384 dim) that you can run on CPU with sub-50ms latency. It works surprisingly well for code if you add a retriever that also uses sparse signals (BM25), but pure dense retrieval with it often misses domain-specific identifiers. It is a good baseline but rarely the best.

How to evaluate models on your own codebase

General benchmarks like MTEB include a code subset (CodeReranker, CodeST), but they test on curated datasets like CodeSearchNet that may not resemble your actual code. You must evaluate on your own data. Here is a simple pipeline:

1. Collect query-code pairs. For example, use docstrings that describe a function as the query, and the function body as the target. Alternatively, extract # TODO comments and the code that implements them, or pair user queries from your search logs with the file that satisfied them.
2. Embed all queries and all code chunks with the candidate model.
3. For each query, compute cosine similarity against every code chunk, rank them, and compute recall@k (k=1,5,10). Repeat for each model.

Below is a Python script using sentence-transformers and numpy to run this evaluation. It assumes you have a list of (Query, CodeText) pairs.

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List, Tuple
 
def evaluate_recall(model_name: str, pairs: List[Tuple[str, str]], k: int = 5) -> float:
    model = SentenceTransformer(model_name)
    queries = [p[0] for p in pairs]
    code_texts = [p[1] for p in pairs]
    
    # Embed in batches
    q_embs = model.encode(queries, show_progress_bar=True, normalize_embeddings=True)
    c_embs = model.encode(code_texts, show_progress_bar=True, normalize_embeddings=True)
    
    # Cosine similarity (already normalized)
    sim_matrix = q_embs @ c_embs.T   # shape (n_queries, n_codes)
    
    correct = 0
    for i in range(len(pairs)):
        # The correct code chunk is at index i (assuming pairs are aligned)
        ranks = np.argsort(-sim_matrix[i])  # descending
        if i in ranks[:k]:
            correct += 1
    return correct / len(pairs)
 
# Example usage
# pairs = [("Query1", "Code1"), ...]
# recall = evaluate_recall("microsoft/codebert-base", pairs, k=5)

Important: The above assumes a one-to-one match per query. In practice, a query may have multiple relevant code chunks (e.g., a function and its unit test). Modify the evaluation to handle many relevant items per query by using the entire set. Also, ensure you split your pairs into a development set and a held-out test set—tuning on the same data you test on overestimates performance.

When I ran this on a Python codebase with 500 functions, text-embedding-3-small achieved recall@5 of 0.85, while CodeBERT scored 0.78 and all-MiniLM scored 0.62. Those numbers are reasonable but your mileage will vary based on code style and language.

For a more rigorous benchmark, consider using the MTEB code subset by cloning the MTEB leaderboard repository and running the CodeRetrieval tasks. The repo provides standardised datasets like CodeSearchNet and StackOverflowDupQuestions. But again, cross-check with a sample from your real code.

To ensure your evaluation environment is reproducible, you can use a consistent setup with a single script and Dev Containers—see Reproducible Dev Environments with Dev Containers and a Single Script for a pattern that works for ML experiments.

Common pitfalls and how to avoid them

Vocabulary mismatch. Subword tokenizers (BPE, WordPiece) split code identifiers into pieces. For example, getUserById might become ['get', 'User', 'By', 'Id']—the embedding of the whole token often loses the semantic combination of "get user by id". This hurts retrieval when your codebase uses compound camelCase or snake_case names. To test, include queries that are exact identifiers from your code, and see if the model returns the correct function. If recall is poor, consider using a model pre-trained on code (CodeBERT, UnixCoder) because they see many code tokens during pre-training.

You can inspect tokenization directly to gauge the problem:

from transformers import AutoTokenizer
 
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
tokens = tokenizer.tokenize("getUserById")
print(tokens)  # ['get', 'User', 'By', 'Id']

This reveals that the model sees four separate tokens, not a single semantic unit. If your codebase uses many such identifiers, you may need to augment the model with a custom tokenizer or use a model that preserves whole words better.

Sequence length limits. Many open-source models cap at 512 tokens. A function with four nested loops and a docstring can easily exceed that. If you truncate, you lose the core logic. Options:

• Chunk functions into logical blocks (e.g., split at docstrings or at signature boundaries) and embed each as a separate vector.
• Use a model with longer context: OpenAI models handle up to 8191 tokens. For self-hosted, you can use codebert-base-msmarco or fine-tune CodeBERT with longer position embeddings, but that's non-trivial.

When chunking, caching strategies for embeddings become important—similar to how prompt caching works. For more details on what affects cache hit rates, see Prompt Caching: What Actually Gets Cached and Why Your Hit Rate Is Zero.

Language support. CodeBERT was trained on six languages; UnixCoder on a dozen. If your codebase uses Go, TypeScript, Rust, or Swift, check the model's training data. The Hugging Face model card usually lists languages. Missing language support means the model will see novel syntax and produce poor embeddings. In that case, switch to a multilingual code model like microsoft/unixcoder-base-nine or use a general sentence transformer but pair it with BM25 for hybrid retrieval.

Another pitfall: treating code comments as the only source of query text. In real code search, users often type things like "encrypt password before storing". Your model must handle typos, missing spaces, and non-standard phrasing. GloVe-style embeddings fail here; transformer models with subword tokenization are more robust.

Choosing the right model size and dimension

Higher dimensions (e.g., 1536 from text-embedding-3-small default, or 3072 from large) improve recall but increase storage and retrieval latency. For a vector database like Pinecone, Weaviate, or Qdrant, doubling dimension roughly doubles memory and slows down exact nearest neighbour search. With approximate nearest neighbour (ANN) indexes, the effect is less, but still significant.

Smaller models (384 dim, like all-MiniLM) are fast and memory-efficient but often miss nuanced semantic similarity. If you have a high throughput requirement (e.g., thousands of queries per second), consider a small model combined with a reranker—a two-stage pipeline where the first stage retrieves 50 candidates cheaply, and a more expensive cross-encoder (like cross-encoder/ms-marco-MiniLM-L-6-v2) reranks them.

Model size (parameters) versus inference speed is a separate trade-off. CodeBERT-base has about 125M parameters; all-MiniLM has 22M. On CPU, CodeBERT takes ~200ms per query (single sample), vs ~30ms for MiniLM. On a GPU (e.g., T4), CodeBERT runs ~20ms and MiniLM ~5ms. For indexing once, latency doesn't matter; for real-time search, it does.

I recommend starting with a 768-dimension model (like UnixCoder) because it offers a good balance of recall and size. Profile your retrieval pipeline end-to-end before scaling up. If latency is acceptable, move to a larger dimension like 1024 or 1536. If not, try quantizing the model with optimum-intel or bitsandbytes to reduce memory and speed up inference.

Decision framework: maps to your constraints

Use this decision tree to narrow down:

• Do you have sensitive code that must not leave your network? → Go open-source. Start with microsoft/unixcoder-base-nine if your code includes multiple languages, or microsoft/codebert-base if only Python/Java/JS.
• Can you accept API costs and data leaves your control? → OpenAI text-embedding-3-small is the best value per token and dimension.
• Is latency critical (<100ms per query)? → Favour a small model, e.g., all-MiniLM-L6-v2 and accept a recall drop or pair with BM25. Or use a proprietary API endpoint that runs on fast hardware.
• Is recall the highest priority? → Use a larger model (e.g., text-embedding-3-large or CodeBERT with longer context) and hybrid search—combine dense embeddings with BM25 scores via reciprocal rank fusion (RRF).

Here is a summary table:

For the compliance scenario, you will also want to consider the infrastructure setup. You can serve the model using ONNX Runtime to improve inference speed on CPU. The Hugging Face ONNX export guide walks through that process. If you are already using prompt caching in your LLM pipeline (caching embeddings is analogous), see our post on prompt caching mechanics and hit rate for trade-offs in caching strategies.

Key takeaways

• Code embeddings need models trained on source code; general sentence transformers underperform on code identifiers and syntax.
• Evaluate on your own query-code pairs using recall@k; do not rely solely on public benchmarks.
• Proprietary APIs (OpenAI) offer ease of use and long context but raise privacy and cost concerns; open-source models (CodeBERT, UnixCoder) give full control but require GPU infrastructure.
• Sequence length limits are a common trap; either chunk long functions or use a model with a larger context window.
• Start with a 768-dimension model; only scale up dimension or model size after profiling your retrieval latency and storage costs.