Upload folder using huggingface_hub

README.md

@@ -1,3 +1,59 @@
 ---
+language: en
+tags:
+- code
+- semantic-search
+- jepa
+- code-search
 license: mit
+datasets:
+- claudios/code_search_net
+metrics:
+- mrr
 ---
+
+# Repo-JEPA: Semantic Code Navigator (SOTA 0.90 MRR)
+
+A **Joint Embedding Predictive Architecture** (JEPA) for semantic code search, trained on 411,000 real Python functions using an NVIDIA H100.
+
+## 🏆 Performance
+
+Tested on 1,000 unseen real-world Python functions from CodeSearchNet.
+
+| Metric | Result | Target |
+|--------|--------|--------|
+| **MRR** | **0.9052** | 0.60 |
+| **Hits@1** | **86.2%** | - |
+| **Hits@5** | **95.9%** | - |
+| **Hits@10** | **97.3%** | - |
+| **Median Rank** | **1.0** | - |
+
+## 🧩 Usage (AutoModel)
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+# 1. Load Model
+model = AutoModel.from_pretrained("uddeshya-k/RepoJepa", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
+
+
+# 2. Encode Code
+code = "def handle_login(user): return auth.verify(user)"
+code_embed = model.encode_code(**tokenizer(code, return_tensors="pt"))
+
+# 3. Encode Query
+query = "how to authenticate users?"
+query_embed = model.encode_query(**tokenizer(query, return_tensors="pt"))
+
+# 4. Search
+similarity = (code_embed @ query_embed.T).item()
+print(f"Similarity: {similarity:.4f}")
+```
+
+## 🏗️ Technical Details
+
+- **Backbone**: CodeBERT (RoBERTa-style)
+- **Loss**: VICReg (Variance-Invariance-Covariance Regularization)
+- **Hardware**: NVIDIA H100 PCIe (80GB VRAM)
+- **Optimizer**: AdamW + OneCycleLR

config.json

@@ -0,0 +1,18 @@
+{
+  "model_type": "repo-jepa",
+  "architectures": ["RepoJEPAModel"],
+  "hidden_dim": 768,
+  "num_encoder_layers": 12,
+  "num_attention_heads": 12,
+  "intermediate_dim": 3072,
+  "hidden_dropout_prob": 0.1,
+  "attention_dropout_prob": 0.1,
+  "vocab_size": 50265,
+  "max_seq_len": 512,
+  "pad_token_id": 1,
+  "base_model": "microsoft/codebert-base",
+  "auto_map": {
+    "AutoConfig": "modeling_repo_jepa.RepoJEPAConfig",
+    "AutoModel": "modeling_repo_jepa.RepoJEPAModel"
+  }
+}

model.safetensors

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

modeling_repo_jepa.py

@@ -0,0 +1,138 @@
+"""
+Hugging Face Export for Repo-JEPA
+
+This file enables loading Repo-JEPA with AutoModel.from_pretrained()
+using trust_remote_code=True.
+"""
+
+import copy
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig, PreTrainedModel, RobertaModel
+
+
+class RepoJEPAConfig(PretrainedConfig):
+    """Configuration for Repo-JEPA model."""
+    
+    model_type = "repo-jepa"
+    
+    def __init__(
+        self,
+        hidden_dim: int = 768,
+        num_encoder_layers: int = 12,
+        num_attention_heads: int = 12,
+        intermediate_dim: int = 3072,
+        hidden_dropout_prob: float = 0.1,
+        attention_dropout_prob: float = 0.1,
+        vocab_size: int = 50265,
+        max_seq_len: int = 512,
+        pad_token_id: int = 1,
+        base_model: str = "microsoft/codebert-base",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.num_encoder_layers = num_encoder_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_dim = intermediate_dim
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_dropout_prob = attention_dropout_prob
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.pad_token_id = pad_token_id
+        self.base_model = base_model
+
+
+class ProjectionHead(nn.Module):
+    """MLP projection head."""
+    
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(input_dim, output_dim),
+            nn.BatchNorm1d(output_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(output_dim, output_dim),
+            nn.BatchNorm1d(output_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(output_dim, output_dim),
+        )
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+class RepoJEPAModel(PreTrainedModel):
+    """
+    Repo-JEPA: Joint Embedding Predictive Architecture for Code Search.
+    
+    Use for semantic code search (encode_code) and retrieval queries (encode_query).
+    """
+    
+    config_class = RepoJEPAConfig
+    
+    def __init__(self, config: RepoJEPAConfig):
+        super().__init__(config)
+        
+        # In the HF model, we store both encoders
+        self.context_encoder = RobertaModel.from_pretrained(
+            config.base_model,
+            add_pooling_layer=False,
+        )
+        self.target_encoder = RobertaModel.from_pretrained(
+            config.base_model,
+            add_pooling_layer=False,
+        )
+        
+        # Projection heads
+        hidden_size = self.context_encoder.config.hidden_size
+        self.context_projector = ProjectionHead(hidden_size, config.hidden_dim)
+        self.target_projector = ProjectionHead(hidden_size, config.hidden_dim)
+        
+        self.post_init()
+    
+    def encode_code(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode code snippet into embedding space."""
+        outputs = self.context_encoder(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = self._mean_pool(outputs.last_hidden_state, attention_mask)
+        return self.context_projector(pooled)
+    
+    def encode_query(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode search query (docstring) into embedding space."""
+        outputs = self.target_encoder(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = self._mean_pool(outputs.last_hidden_state, attention_mask)
+        return self.target_projector(pooled)
+
+    def _mean_pool(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
+            sum_hidden = torch.sum(hidden_states * mask, dim=1)
+            sum_mask = torch.clamp(mask.sum(dim=1), min=1e-9)
+            return sum_hidden / sum_mask
+        return hidden_states.mean(dim=1)
+
+    def forward(self, **kwargs):
+        # HF requires forward(), we default to code encoding or raise error
+        if "input_ids" in kwargs:
+            return self.encode_code(kwargs["input_ids"], kwargs.get("attention_mask"))
+        raise NotImplementedError("Use .encode_code() or .encode_query() specifically.")
+
+
+# Register with Auto classes
+try:
+    from transformers import AutoConfig, AutoModel
+    AutoConfig.register("repo-jepa", RepoJEPAConfig)
+    AutoModel.register(RepoJEPAConfig, RepoJEPAModel)
+except:
+    pass
+