Code Embedding
Mistral Code Embedding and Retrieval Evaluation
This notebook demonstrates a pipeline for code embedding, chunking, indexing, retrieval, and evaluation using the Mistral API, FAISS, and the SWE-bench Lite dataset.
It uses the Mistral code embedding model codestral-embed to generate code embeddings and FAISS for fast similarity search. The workflow includes flattening repository structures, chunking code files into smaller segments, and generating embeddings for each chunk. These embeddings are indexed to enable efficient retrieval of relevant code snippets in response to user queries. The notebook evaluates retrieval performance on the SWE-bench Lite dataset, using recall metrics to measure effectiveness. This methodology is especially valuable for applications such as code search, code comprehension, and automated software maintenance.
Environment Setup
Install required packages for code embedding, retrieval, and dataset handling.
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 173.4/173.4 kB 3.7 MB/s eta 0:00:00 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 31.3/31.3 MB 31.5 MB/s eta 0:00:00 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 372.3/372.3 kB 12.3 MB/s eta 0:00:00 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.5/6.5 MB 32.7 MB/s eta 0:00:00 ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts. torch 2.6.0+cu124 requires nvidia-cublas-cu12==12.4.5.8; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cublas-cu12 12.5.3.2 which is incompatible. torch 2.6.0+cu124 requires nvidia-cuda-cupti-cu12==12.4.127; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-cupti-cu12 12.5.82 which is incompatible. torch 2.6.0+cu124 requires nvidia-cuda-nvrtc-cu12==12.4.127; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-nvrtc-cu12 12.5.82 which is incompatible. torch 2.6.0+cu124 requires nvidia-cuda-runtime-cu12==12.4.127; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-runtime-cu12 12.5.82 which is incompatible. torch 2.6.0+cu124 requires nvidia-cudnn-cu12==9.1.0.70; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cudnn-cu12 9.3.0.75 which is incompatible. torch 2.6.0+cu124 requires nvidia-cufft-cu12==11.2.1.3; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cufft-cu12 11.2.3.61 which is incompatible. torch 2.6.0+cu124 requires nvidia-curand-cu12==10.3.5.147; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-curand-cu12 10.3.6.82 which is incompatible. torch 2.6.0+cu124 requires nvidia-cusolver-cu12==11.6.1.9; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusolver-cu12 11.6.3.83 which is incompatible. torch 2.6.0+cu124 requires nvidia-cusparse-cu12==12.3.1.170; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusparse-cu12 12.5.1.3 which is incompatible. torch 2.6.0+cu124 requires nvidia-nvjitlink-cu12==12.4.127; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-nvjitlink-cu12 12.5.82 which is incompatible. gcsfs 2025.3.2 requires fsspec==2025.3.2, but you have fsspec 2023.9.2 which is incompatible.
Imports and Tokenizer Initialization
Import necessary libraries and initialize the tokenizer for code embedding.
/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:104: UserWarning: Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'. You are not authenticated with the Hugging Face Hub in this notebook. If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new). warnings.warn(
tekken.json: 0%| | 0.00/14.8M [00:00<?, ?B/s]
/usr/local/lib/python3.11/dist-packages/mistral_common/tokens/tokenizers/tekken.py:184: FutureWarning: Special tokens not found in /root/.cache/huggingface/hub/models--mistralai--Mistral-Small-3.1-24B-Base-2503/snapshots/db7c968753c07380364d963090b5cf8cc131a0c3/tekken.json and default to ({'rank': 0, 'token_str': <SpecialTokens.unk: '<unk>'>, 'is_control': True}, {'rank': 1, 'token_str': <SpecialTokens.bos: '<s>'>, 'is_control': True}, {'rank': 2, 'token_str': <SpecialTokens.eos: '</s>'>, 'is_control': True}, {'rank': 3, 'token_str': <SpecialTokens.begin_inst: '[INST]'>, 'is_control': True}, {'rank': 4, 'token_str': <SpecialTokens.end_inst: '[/INST]'>, 'is_control': True}, {'rank': 5, 'token_str': <SpecialTokens.begin_tools: '[AVAILABLE_TOOLS]'>, 'is_control': True}, {'rank': 6, 'token_str': <SpecialTokens.end_tools: '[/AVAILABLE_TOOLS]'>, 'is_control': True}, {'rank': 7, 'token_str': <SpecialTokens.begin_tool_results: '[TOOL_RESULTS]'>, 'is_control': True}, {'rank': 8, 'token_str': <SpecialTokens.end_tool_results: '[/TOOL_RESULTS]'>, 'is_control': True}, {'rank': 9, 'token_str': <SpecialTokens.tool_calls: '[TOOL_CALLS]'>, 'is_control': True}, {'rank': 10, 'token_str': <SpecialTokens.img: '[IMG]'>, 'is_control': True}, {'rank': 11, 'token_str': <SpecialTokens.pad: '<pad>'>, 'is_control': True}, {'rank': 12, 'token_str': <SpecialTokens.img_break: '[IMG_BREAK]'>, 'is_control': True}, {'rank': 13, 'token_str': <SpecialTokens.img_end: '[IMG_END]'>, 'is_control': True}, {'rank': 14, 'token_str': <SpecialTokens.prefix: '[PREFIX]'>, 'is_control': True}, {'rank': 15, 'token_str': <SpecialTokens.middle: '[MIDDLE]'>, 'is_control': True}, {'rank': 16, 'token_str': <SpecialTokens.suffix: '[SUFFIX]'>, 'is_control': True}, {'rank': 17, 'token_str': <SpecialTokens.begin_system: '[SYSTEM_PROMPT]'>, 'is_control': True}, {'rank': 18, 'token_str': <SpecialTokens.end_system: '[/SYSTEM_PROMPT]'>, 'is_control': True}, {'rank': 19, 'token_str': <SpecialTokens.begin_tool_content: '[TOOL_CONTENT]'>, 'is_control': True}). This behavior will be deprecated going forward. Please update your tokenizer file and include all special tokens you need.
warnings.warn(
API Key Setup
Set up your Mistral API key for authentication.
Enter your MISTRAL_API_KEY: ··········
Embedding and Chunking Configuration
Define parameters for code embedding and chunking.
In our experiments, we find that chunking with small chunk size (3000 characters, ~512 tokens) and overlap (1000 characters), leads to much better retrieval for RAG.
Download and Prepare Repository Structures
Download and extract repository structures for the SWE-bench Lite dataset.
Downloading... From (original): https://drive.google.com/uc?id=1wG1CcfVHi-70FoAd5wPI59WdI4g1LkpS From (redirected): https://drive.google.com/uc?id=1wG1CcfVHi-70FoAd5wPI59WdI4g1LkpS&confirm=t&uuid=ff0d2b10-c817-4e53-abd5-44d7755dd926 To: /content/min_swebench_repo_structure.zip 100%|██████████| 200M/200M [00:01<00:00, 139MB/s]
Utility Functions for Data Processing
Define helper functions for flattening repository structures, chunking code, formatting documents, and extracting file paths from patches.
Embedding, Indexing, and Retrieval Functions
Functions for generating embeddings, building FAISS indexes, retrieving relevant files, and evaluating recall.
Load SWE-bench Lite Dataset
Load the SWE-bench Lite dataset and extract ground truth file changes for evaluation.
Loading SWE-bench Lite dataset...
Downloading readme: 0%| | 0.00/3.67k [00:00<?, ?B/s]
Downloading data files: 0%| | 0/2 [00:00<?, ?it/s]
Downloading data: 0%| | 0.00/120k [00:00<?, ?B/s]
Downloading data: 0%| | 0.00/1.12M [00:00<?, ?B/s]
Extracting data files: 0%| | 0/2 [00:00<?, ?it/s]
Generating dev split: 0%| | 0/23 [00:00<?, ? examples/s]
Generating test split: 0%| | 0/300 [00:00<?, ? examples/s]
Loaded 300 instances from SWE-bench Lite
Main Evaluation Loop
For each instance in the dataset, index the repository, retrieve relevant files for the problem statement, and compute recall@5.
Processing instance 1 of 300 Indexing repository for astropy__astropy-12907... Found 872 Python files Chunking 872 files... Created 6660 chunks from 872 files (size increase: 7.6x) Preparing texts for embedding... Getting embeddings...
Processing embedding batches: 100%|██████████| 263/263 [05:22<00:00, 1.23s/it]
Creating FAISS index... Saved index for astropy__astropy-12907 with 6660 chunks
Processing embedding batches: 100%|██████████| 1/1 [00:00<00:00, 3.96it/s]
Retrieved files: ['astropy/modeling/separable.py', 'astropy/modeling/tests/test_separable.py', 'astropy/modeling/tests/test_models.py', 'astropy/modeling/core.py', 'astropy/modeling/tests/test_compound.py'] Ground truth files: ['astropy/modeling/separable.py'] astropy__astropy-12907: Recall@5 = 1.000
Results and Summary
Calculate and save the average recall@5 and detailed results for all evaluated instances.
Evaluation complete! Average Recall@5: 1.000 Results saved to /content/swebench_results.json