Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Code Understanding › Code Implementation: Semantic Code Search

Michael BrenndoerferFebruary 2, 20266 min read

Language AI HandbookAll chapters

Part 47 of 67Code Generation

Chapter 358Code Understanding › Code Implementation: Semantic Code Search

Reading section 10 of 13 · Code Implementation: Semantic Code Search · 6 minSection 10 of 13 · 6 min

Section 10 of 13 · 6 min read

Code Understanding › Code Implementation: Semantic Code SearchOne section · Sections in this chapter10 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's implement semantic code search using the sentence-transformers library with a code-aware embedding model.

In[18]:

Code

from sentence_transformers import SentenceTransformer

search_model = SentenceTransformer("all-MiniLM-L6-v2")

In[19]:

Code

code_repository = {
    "binary_search": """\
def binary_search(arr, target):
    lo, hi = 0, len(arr) - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return -1""",
    "merge_sort": """\
def merge_sort(arr):
    if len(arr) <= 1:
        return arr
    mid = len(arr) // 2
    left = merge_sort(arr[:mid])
    right = merge_sort(arr[mid:])
    return merge(left, right)

def merge(left, right):
    result = []
    i = j = 0
    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            result.append(left[i]); i += 1
        else:
            result.append(right[j]); j += 1
    return result + left[i:] + right[j:]""",
    "find_duplicates": """\
def find_duplicates(lst):
    seen = set()
    duplicates = set()
    for item in lst:
        if item in seen:
            duplicates.add(item)
        else:
            seen.add(item)
    return sorted(duplicates)""",
    "word_count": """\
def word_count(text):
    words = text.lower().split()
    counts = {}
    for word in words:
        counts[word] = counts.get(word, 0) + 1
    return counts""",
    "flatten_list": """\
def flatten(nested):
    result = []
    for item in nested:
        if isinstance(item, list):
            result.extend(flatten(item))
        else:
            result.append(item)
    return result""",
}

In[20]:

Code

code_names = list(code_repository.keys())
code_snippets = list(code_repository.values())
code_embeddings = search_model.encode(code_snippets, normalize_embeddings=True)


def semantic_search(query: str, top_k: int = 3) -> list:
    """Retrieve top_k most relevant code snippets for a natural language query."""
    query_embedding = search_model.encode([query], normalize_embeddings=True)[0]
    scores = code_embeddings @ query_embedding
    ranked_indices = np.argsort(scores)[::-1][:top_k]
    results = []
    for idx in ranked_indices:
        results.append(
            {
                "name": code_names[idx],
                "score": float(scores[idx]),
            }
        )
    return results

In[21]:

Code

test_queries = [
    "find an element in a sorted array",
    "count word frequency in text",
    "remove duplicates from a list",
    "sort a list recursively",
    "recursively flatten a nested list structure",
]

search_results = {q: semantic_search(q, top_k=2) for q in test_queries}

Out[22]:

Console

Query: 'find an element in a sorted array'
  Rank 1: find_duplicates (score=0.3704)
  Rank 2: binary_search (score=0.3511)

Query: 'count word frequency in text'
  Rank 1: word_count (score=0.7495)
  Rank 2: find_duplicates (score=0.1245)

Query: 'remove duplicates from a list'
  Rank 1: find_duplicates (score=0.6154)
  Rank 2: merge_sort (score=0.2775)

Query: 'sort a list recursively'
  Rank 1: merge_sort (score=0.4107)
  Rank 2: find_duplicates (score=0.3764)

Query: 'recursively flatten a nested list structure'
  Rank 1: flatten_list (score=0.6671)
  Rank 2: merge_sort (score=0.2401)

The similarity scores reflect semantic closeness, but not perfectly. The word-frequency query retrieves word_count, while the sorted-array query ranks find_duplicates just above binary_search. That miss is a useful reminder that this general sentence encoder was not trained specifically for code retrieval. A code-specific embedding model such as UniXcoder or CodeBERT fine-tuned on code search generally improves the ranking.

Out[23]:

Visualization

Annotated five-by-five cosine-similarity heatmap with four correct row maxima and one error where find-duplicates narrowly outranks binary-search for the sorted-array query. — Cosine similarities between five natural-language queries and five code snippets. Darker blue cells indicate stronger alignment. Four intended matches are row maxima; for 'find in sorted array,' the general sentence encoder scores `find_duplicates` at 0.37 and `binary_search` at 0.35, exposing a small but consequential retrieval error.

The heatmap reveals both the intended structure and the model's mistakes. Four rows peak on the expected function, while the first row's off-diagonal maximum shows the sorted-array retrieval error. Other off-diagonal cells capture semantic overlap: the model associates both merge_sort and find_duplicates with list operations even when only one is the intended target.

Evaluation Metrics for Code SearchLink Copied

Code search systems are evaluated using metrics adapted from information retrieval. The two most common are Mean Reciprocal Rank and Recall at K.

MRR (Mean Reciprocal Rank): For each query, the reciprocal of the rank of the first correct result, averaged over queries. A model that consistently puts the correct answer first achieves MRR = 1.0. If the correct answer appears at rank 2, the reciprocal rank is 0.5. If it appears at rank 5, the reciprocal rank is 0.2.

\text{MRR} = \frac{1}{|Q|} \sum_{q \in Q} \frac{1}{\text{rank}(q)}

where $\text{rank}(q)$ is the rank position of the correct code snippet for query $q$ .

R@K (Recall at K): The fraction of queries for which the correct answer appears in the top $K$ results. R@1 is the strictest, requiring the correct answer to be the top result. R@10 is more lenient but reflects realistic usage: developers will often scan the top 10 results to find what they need.

On the CodeSearchNet benchmark, current models achieve MRR around 0.70-0.85 depending on programming language, with Python and JavaScript showing stronger performance than Go and Ruby due to training data volume. This reflects an important practical consideration: the quality of code search degrades for languages with less public code, independent of the model architecture.

Practical Deployment ConsiderationsLink Copied

Deploying semantic code search in a real engineering organization introduces several challenges beyond model quality.

Index freshness. Code repositories change continuously. The ANN index must be updated as functions are added, modified, and deleted. Incremental index updates are more efficient than full rebuilds but require careful management to avoid stale results.

Granularity. Functions are the natural unit for code search, but they are not always the right granularity. A developer searching for "how does the authentication system work" may want a higher-level answer that spans many functions. A developer searching for "the specific line that writes to the audit log" wants sub-function granularity. Systems that index at multiple granularities and blend results across levels handle a wider range of queries.

Private code and confidentiality. Enterprise code search requires that the model not leak proprietary code through its embeddings. This is less of a concern for retrieval systems, where the model only produces vectors, than for generation systems that might reproduce training code verbatim. Nonetheless, access control is important: a search query from one team should not return results from another team's private repository.