CodeMalt

CodeMalt is a high-performance, code-specialized static embedding model created through Model2Vec distillation of sentence-transformers/all-mpnet-base-v2. This model achieves 73.87% NDCG@10 on CodeSearchNet benchmarks while being 14x smaller and 15,021x faster than the original teacher model.

🏆 Performance Highlights

• NDCG@10: 0.7387 (Best among all distilled models)
• Mean Reciprocal Rank (MRR): 0.7010
• Recall@5: 0.8017
• Model Size: 7.6M parameters (vs 109M original)
• Inference Speed: 15,021x faster than teacher model
• Memory Usage: <1GB RAM (vs 8+ GB VRAM for original)

📊 CodeSearchNet Performance by Language

🔧 Model Details

• Teacher Model: sentence-transformers/all-mpnet-base-v2
• Distillation Method: Model2Vec + Tokenlearn training on CodeSearchNet
• Architecture: Static embeddings (no neural network inference required)
• Embedding Dimensions: 256
• Training Data: CodeSearchNet code-comment pairs across 6 programming languages
• Optimization: PCA dimensionality reduction + SIF weighting + Zipf regularization
• Vocabulary Size: 29,528
• Parameters: 7.6M
• Size: 14.4MB

🎯 Distiller: Code-Specialized Embedding Toolkit

Distiller is an independent toolkit built upon Model2Vec and Tokenlearn for creating code-specialized static embeddings. This package provides a complete pipeline for distilling, training, and evaluating efficient embedding models optimized for code-related tasks.

Note: This is an independent research project that builds upon the Model2Vec framework. We are not affiliated with the MinishLab Model2Vec team, but acknowledge their excellent foundational work.

[!Important] Check out the comprehensive REPORT.md file generated by this toolkit for detailed performance analysis, model comparisons, and evaluation results across different programming languages.

[!Warning] Research Finding: See NOTES.md for critical analysis showing that C4 fine-tuning significantly degraded performance (-16.8% NDCG@10) compared to simple Model2Vec distillation. Recommendation: Use basic distillation without additional training for optimal code embedding performance.

The distiller package provides a complete pipeline for:

1. Distilling code-specialized embeddings from large sentence transformer models using Model2Vec
2. Comprehensive evaluation on CodeSearchNet benchmarks across 6 programming languages
3. Performance benchmarking (speed, memory, model size analysis)
4. Advanced training with tokenlearn for enhanced code understanding
5. Analysis and reporting with visualizations and comparison charts
6. Cloud-scale processing with Beam support for distributed execution

Key Benefits

• 🚀 Performance: Up to 500x faster inference with 50x smaller models
• 📊 Code-Optimized: Specialized for code search, classification, and similarity tasks
• 🔬 Comprehensive: Full evaluation pipeline with CodeSearchNet metrics
• ☁️ Scalable: Local and cloud execution with Beam support
• 📈 Analytical: Rich reporting with performance charts and comparisons

🚀 Quick Start

Installation

# Install with all dependencies
pip install model2vec[train] torch transformers datasets sentence-transformers
pip install typer pydantic plotly matplotlib seaborn

# Install the distiller package (assuming local development)
pip install -e .

Basic Usage

# Simple distillation of a teacher model
distiller distill

# Distillation with advanced CodeSearchNet training  
distiller distill --train

# Evaluate distilled models on CodeSearchNet
distiller evaluate

# Generate comprehensive analysis report
distiller analyze

Python API

from distiller import distill, evaluate, analyze

# Distill a specific model
results = distill.run_local_distillation(
    teacher_models=["microsoft/codebert-base"],
    enable_training=True,  # Include CodeSearchNet fine-tuning
    pca_dims=256
)

# Evaluate on CodeSearchNet
evaluation_results = evaluate.run_evaluation(
    models=["."],
    max_queries=1000,
    languages=["python", "javascript", "java", "go", "php", "ruby"]
)

# Generate analysis report
analyze.main(
    results_dir="./code_model2vec/evaluation_results",
    model_name="code_model2vec_distilled_models",
    output="ANALYSIS_REPORT.md"
)

📋 Features

🔬 Distillation Engine

• Multiple Teacher Models: Support for 15+ pre-configured teacher models including:

  • Code-specialized: microsoft/codebert-base, BAAI/bge-code-v1, Salesforce/SFR-Embedding-Code-2B_R
  • General-purpose: sentence-transformers/all-mpnet-base-v2, BAAI/bge-m3
  • Instruction-tuned: Alibaba-NLP/gte-Qwen2-1.5B-instruct

• Advanced Training Pipeline: Optional tokenlearn-based training following the POTION approach:

  1. Model2Vec distillation (basic static embeddings)
  2. Feature extraction using sentence transformers
  3. Tokenlearn training on CodeSearchNet data
  4. Post-training re-regularization (PCA + SIF weighting)

• Robust Model Handling: Automatic compatibility checks and specialized handling for problematic models

📊 Evaluation Framework

• CodeSearchNet Evaluation: Standard code search benchmarks across 6 programming languages
• Retrieval Metrics: NDCG@k, MRR, Recall@k, Mean/Median Rank
• Performance Benchmarking:
  • Model size analysis (disk usage, parameters, memory footprint)
  • Inference speed testing (various batch sizes and text lengths)
  • CPU vs GPU performance comparison
  • Memory scaling analysis

📈 Analysis & Reporting

• Comprehensive Reports: Automated generation of analysis reports with:

  • Performance comparison tables
  • Language-specific radar charts
  • Efficiency analysis (performance vs model size)
  • Peer model comparisons

• Rich Visualizations: Plotly and Matplotlib charts including:

  • Multi-model performance heatmaps
  • Batch size scaling curves
  • Memory usage patterns
  • Model efficiency scatter plots

☁️ Cloud Integration

• Beam Support: Distributed execution on Beam cloud infrastructure
• Volume Management: Persistent storage with checkpoint support
• Resource Optimization: GPU-optimized configurations (A100-40G default)
• Automatic Syncing: Seamless model and result synchronization

🛠️ CLI Reference

distiller distill

Distill teacher models into efficient static embeddings.

distiller distill [OPTIONS]

Options:
  --use-beam              Use Beam cloud for distillation
  --train                 Enable advanced training (CodeSearchNet fine-tuning)  
  --teacher-models TEXT   Specific teacher models to distill (can be repeated)
  --pca-dims INTEGER      PCA dimensions (default: 256)
  --clear-cache          Clear HuggingFace cache for problematic models

Examples:

# Basic distillation of all default models
distiller distill

# Train specific models with advanced CodeSearchNet fine-tuning
distiller distill --train --teacher-models microsoft/codebert-base --teacher-models BAAI/bge-code-v1

# Use Beam cloud with custom PCA dimensions
distiller distill --use-beam --train --pca-dims 512

distiller evaluate

Evaluate models on CodeSearchNet benchmarks with performance analysis.

distiller evaluate [OPTIONS]

Options:
  --use-beam              Use Beam cloud for evaluation
  --skip-third-party      Skip third-party models evaluation
  --skip-benchmark        Skip performance benchmarking  
  --max-queries INTEGER   Maximum queries per language (default: 100)

Examples:

# Comprehensive evaluation with benchmarking
distiller evaluate --max-queries 1000

# Quick evaluation without performance benchmarks
distiller evaluate --skip-benchmark --max-queries 100

# Cloud-based evaluation
distiller evaluate --use-beam --max-queries 500

distiller analyze

Generate comprehensive analysis reports with visualizations.

distiller analyze [OPTIONS]

Options:
  --results-dir PATH      Results directory (default: code_model2vec/evaluation_results)
  --model-name TEXT       Model name for analysis (default: gte_qwen2_m2v_code (Ours))
  --output PATH           Output report file (default: REPORT.md)
  --export-csv PATH       Export results to CSV file

Examples:

# Generate standard analysis report
distiller analyze

# Custom analysis with CSV export
distiller analyze --model-name "my_distilled_model" --output custom_report.md --export-csv results.csv

# Analyze specific results directory
distiller analyze --results-dir ./custom_results --output analysis.md

📁 Directory Structure

The distiller uses a standardized directory structure:

code_model2vec/
├── base/                    # Basic distilled models (Step 1)
│   └── code_model2vec_{teacher_name}/
├── final/                   # Final models (copied from base or after training)
│   └── code_model2vec_{teacher_name}[_fine_tuned]/
├── evaluation_results/      # CodeSearchNet evaluation results
│   └── comprehensive_eval_{model}.json
├── benchmark_results/       # Performance benchmark results  
├── analysis_results/        # Analysis reports and charts
│   └── charts/
├── checkpoints/            # Training checkpoints
└── cache/                  # Temporary cache files

⚙️ Configuration

Teacher Models

Default supported teacher models (configured in config.py):

TEACHER_MODELS = [
    "Alibaba-NLP/gte-Qwen2-1.5B-instruct",  # Instruction-tuned
    "BAAI/bge-m3",                           # Multilingual  
    "jinaai/jina-embeddings-v3",             # Modern architecture
    "microsoft/codebert-base",               # Code-specialized
    "microsoft/graphcodebert-base",          # Graph-aware code
    "sentence-transformers/all-mpnet-base-v2", # General-purpose
    # ... and more
]

Distillation Parameters

# Model2Vec distillation settings
optimal_pca_dims: int = 256
sif_coefficient: float = 1e-3  
apply_zipf: bool = True

# Tokenlearn training settings (when --train is enabled)
tokenlearn_dataset: str = "sentence-transformers/codesearchnet"
tokenlearn_text_key: str = "code"  # Use code field for training

Evaluation Settings

# CodeSearchNet evaluation
evaluation_languages = ["python", "java", "javascript", "php", "ruby", "go"]
max_queries_per_language: int = 1000
evaluation_metrics = ["ndcg@1", "ndcg@5", "ndcg@10", "mrr", "recall@1", "recall@5", "recall@10"]

📄 License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.

🙏 Acknowledgments

This independent research project builds upon several excellent open-source foundations:

• Model2Vec by MinishLab - Core static embedding distillation framework
• Tokenlearn by MinishLab - Advanced token-level training methodology
• CodeSearchNet by GitHub - Code search benchmark dataset and evaluation framework
• Sentence Transformers by UKP Lab - Teacher model ecosystem and training framework
• Beam - Distributed cloud computing infrastructure
• Transformers by Hugging Face - Model loading and tokenization utilities

Note: While this toolkit leverages Model2Vec and Tokenlearn, it is an independent research contribution and is not officially associated with or endorsed by the MinishLab team.