faisalmumtaz/codecompass-embed

CodeCompass-Embed

CodeCompass-Embed is a 494M-parameter embedding model for semantic code search and retrieval, trained on 86B tokens total. It produces 896-dimensional embeddings optimized for matching natural language queries to code across Python, Java, JavaScript, Go, Ruby, and PHP, achieving state-of-the-art results on the CoIR code retrieval benchmark.

Model Highlights

• —Code search from natural language — find relevant code snippets across Python, Java, JavaScript, Go, Ruby, PHP
• —Competitive with models 3× smaller and larger — 494M params, 896-dim embeddings
• —Bidirectional attention — all 24 layers converted from causal for better embedding quality
• —Lightweight — runs on consumer GPUs, trained at 512 tokens with RoPE extrapolation for longer inputs
• —Versatile — supports NL→Code, Code→Code, Q&A, and Text→SQL retrieval via instruction templates

Model Details

Benchmark Results (CoIR)

Evaluated on the CoIR Benchmark (ACL 2025). All scores are NDCG@10. Sorted by CSN-Python.

Multi-Language Code Search (CodeSearchNet)

Full Results (All 12 Tasks)

Usage

With Transformers

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

# Load model
model = AutoModel.from_pretrained("faisalmumtaz/codecompass-embed", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("faisalmumtaz/codecompass-embed")

# CRITICAL: Enable bidirectional attention for embeddings
for layer in model.model.layers:
    layer.self_attn.is_causal = False

model.eval()

def encode(texts, is_query=False):
    # Add instruction prefix for queries
    if is_query:
        texts = [f"Instruct: Find the most relevant code snippet given the following query:\nQuery: {{t}}" for t in texts]
    
    inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)
        hidden = outputs.hidden_states[-1]
        
        # Mean pooling
        mask = inputs["attention_mask"].unsqueeze(-1).float()
        embeddings = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
        
        # L2 normalize
        embeddings = F.normalize(embeddings, p=2, dim=-1)
    
    return embeddings

# Example: Code Search
query = "How to sort a list in Python"
code_snippets = [
    "def sort_list(lst):\n    return sorted(lst)",
    "def add_numbers(a, b):\n    return a + b",
    "def reverse_string(s):\n    return s[::-1]",
]

query_emb = encode([query], is_query=True)
code_embs = encode(code_snippets, is_query=False)

# Compute similarities
similarities = (query_emb @ code_embs.T).squeeze()
print(f"Query: {{query}}")
for i, (code, sim) in enumerate(zip(code_snippets, similarities)):
    print(f"  [{{sim:.4f}}] {{code[:50]}}...")

Instruction Templates

For optimal performance, use these instruction prefixes for queries:

Note: Document/corpus texts do NOT need instruction prefixes.

Training Details

Training followed a two-stage approach:

Stage 1 — Embedding Conversion (8.8M samples): Converted Qwen2.5-Coder-0.5B from a causal language model to a bidirectional embedding model. Trained on 8.8M samples spanning CoRNStack (Python, Java, JavaScript, Go, Ruby, PHP), CoderPile, StackOverflow, and synthetic data with mined hard negatives.

Stage 2 — Hard Negative Refinement (100K samples): Continued fine-tuning on a curated 100K-sample subset with hard negatives.

• —Base Model: Qwen2.5-Coder-0.5B
• —Architecture: Bidirectional attention across all 24 layers, mean pooling, L2 normalization
• —Loss: InfoNCE with temperature τ=0.05
• —Effective Batch Size: 1024 (via GradCache)
• —Hardware: NVIDIA H100 (95GB)

Limitations

• —Strongest on Python; other languages show lower but competitive performance
• —Weaker on competitive programming tasks (APPS) due to long solution lengths vs. 512 training context
• —May not generalize to low-resource programming languages not seen in training

Citation

@misc{{codecompass2026,
  author = {{Faisal Mumtaz}},
  title = {{CodeCompass-Embed: A Code Embedding Model for Semantic Code Search}},
  year = {{2026}},
  publisher = {{Hugging Face}},
  url = {{https://huggingface.co/faisalmumtaz/codecompass-embed}}
}}

License

Apache 2.0