Upload CodeCompass-Embed model (NDCG@10=0.9228 on CodeSearchNet-Python)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,202 @@
+---
+license: apache-2.0
+language:
+- en
+- code
+library_name: transformers
+tags:
+- code
+- embeddings
+- retrieval
+- code-search
+- semantic-search
+- feature-extraction
+- sentence-transformers
+datasets:
+- CoIR-Retrieval/CodeSearchNet-python
+- bigcode/the-stack
+pipeline_tag: feature-extraction
+base_model: Qwen/Qwen2.5-Coder-0.5B
+model-index:
+- name: CodeCompass-Embed
+  results:
+  - task:
+      type: retrieval
+      name: Code Retrieval
+    dataset:
+      type: CoIR-Retrieval/CodeSearchNet-python
+      name: CodeSearchNet Python
+    metrics:
+    - type: ndcg@10
+      value: 0.9228
+      name: NDCG@10
+    - type: mrr@10
+      value: 0.9106
+      name: MRR@10
+---
+
+# CodeCompass-Embed
+
+**CodeCompass-Embed** is a code embedding model fine-tuned from [Qwen2.5-Coder-0.5B](https://huggingface.co/Qwen/Qwen2.5-Coder-0.5B) for semantic code search and retrieval tasks.
+
+## Model Highlights
+
+- 🏆 **SOTA on CodeSearchNet-Python**: NDCG@10 = 0.9228, MRR@10 = 0.9106
+- ⚡ **Efficient**: 494M parameters, runs on consumer GPUs
+- 🔄 **Bidirectional Attention**: Converted from causal to bidirectional for embedding tasks
+- 📏 **Flexible Context**: Trained at 512 tokens, supports up to 32K via RoPE extrapolation
+- 🎯 **Mean Pooling**: Robust to variable-length inputs
+
+## Model Details
+
+| Property | Value |
+|----------|-------|
+| Base Model | Qwen2.5-Coder-0.5B |
+| Parameters | 494M |
+| Embedding Dimension | 896 |
+| Max Sequence Length | 512 (training) / 32K (inference) |
+| Pooling | Mean |
+| Normalization | L2 |
+| Attention | Bidirectional (all 24 layers) |
+
+## Benchmark Results (CoIR)
+
+We evaluate on the [CoIR Benchmark](https://github.com/CoIR-team/coir) (ACL 2025), the gold standard for code retrieval evaluation.
+
+### Per-Task Results
+
+| Task | NDCG@10 | MRR@10 | Recall@10 |
+|------|---------|--------|-----------|
+| **codesearchnet-python** | **0.9228** | **0.9106** | 0.9600 |
+| stackoverflow-qa | 0.6480 | 0.6156 | 0.7500 |
+| synthetic-text2sql | 0.5673 | 0.4853 | 0.8220 |
+| codefeedback-st | 0.4080 | 0.3698 | 0.5300 |
+| codetrans-dl | 0.3305 | 0.2161 | 0.7167 |
+| apps | 0.1277 | 0.1097 | 0.1860 |
+| **Average** | **0.5007** | **0.4512** | - |
+
+### Comparison with SOTA Models
+
+| Model | Params | Avg NDCG@10 | CodeSearchNet-Python |
+|-------|--------|-------------|---------------------|
+| SFR-Embedding-Code-400M | 400M | 0.6786 | - |
+| CodeRankEmbed | 137M | 0.6303 | - |
+| Jina-Code-v2 | 161M | 0.5789 | - |
+| BGE-M3 | 568M | 0.5547 | - |
+| **CodeCompass-Embed (ours)** | **494M** | **0.5007** | **0.9228** |
+| CodeT5+-110M | 110M | 0.4817 | - |
+
+> **Note**: CodeCompass achieves state-of-the-art on CodeSearchNet-Python (NL→Code retrieval), which is the primary use case for code search applications.
+
+## Usage
+
+### With Transformers
+
+```python
+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]}}...")
+```
+
+### With Sentence Transformers (Coming Soon)
+
+```python
+from sentence_transformers import SentenceTransformer
+
+model = SentenceTransformer("faisalmumtaz/codecompass-embed")
+embeddings = model.encode(["def hello(): print('world')"])
+```
+
+## Instruction Templates
+
+For optimal performance, use these instruction prefixes for queries:
+
+| Task | Instruction Template |
+|------|---------------------|
+| NL → Code | `Instruct: Find the most relevant code snippet given the following query:\nQuery: {{query}}` |
+| Code → Code | `Instruct: Find an equivalent code snippet given the following code snippet:\nQuery: {{query}}` |
+| Tech Q&A | `Instruct: Find the most relevant answer given the following question:\nQuery: {{query}}` |
+| Text → SQL | `Instruct: Given a natural language question and schema, find the corresponding SQL query:\nQuery: {{query}}` |
+
+**Note**: Document/corpus texts do NOT need instruction prefixes.
+
+## Training Details
+
+- **Base Model**: Qwen2.5-Coder-0.5B
+- **Training Data**: 8.8M samples from CoRNStack, StackOverflow, CodeSearchNet
+- **Architecture Modification**: Converted all 24 attention layers from causal to bidirectional
+- **Pooling**: Mean pooling (robust for variable-length extrapolation)
+- **Loss**: InfoNCE with temperature τ=0.05
+- **Hard Negatives**: 7 per sample (embedding-mined)
+- **Effective Batch Size**: 1024 (via GradCache)
+- **Training Steps**: 950 (early stopping at best MRR)
+- **Hardware**: NVIDIA H100 (95GB)
+
+## Limitations
+
+- Optimized for **NL → Code** retrieval; weaker on code translation tasks
+- Trained primarily on Python/JavaScript/Java/Go/PHP/Ruby
+- May not generalize well to low-resource programming languages
+
+## Citation
+
+```bibtex
+@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

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,61 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 24,
+  "model_type": "qwen2",
+  "num_attention_heads": 14,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.40.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936,
+  "torch_dtype": "bfloat16",
+  "model_name": "Qwen/Qwen2.5-Coder-0.5B",
+  "embedding_dim": 896,
+  "max_seq_len": 512,
+  "use_lora": false,
+  "pooling": "mean",
+  "normalize": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0fe7769104ed825d2feba8e063b6f4ea499668985ed73a63c067c5ccc2cd1db
+size 988096088

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8ca367b714d9d9f3790a7626974cf2ece38c267640196be1af09af8f3367ae7
+size 11422162

tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}