MatthewsO3/GraphCode-CErl-codesearch

1---2language:3  - code4license: mit5base_model: microsoft/graphcodebert-base6tags:7  - code-search8  - semantic-search9  - graphcodebert10  - erlang11  - cpp12library_name: transformers13pipeline_tag: feature-extraction14---15 16# GraphCode-CErl — Semantic Code Search for Erlang & C++17 18Fine-tuned [GraphCodeBERT](https://huggingface.co/microsoft/graphcodebert-base) for semantic code search over **Erlang** and **C++** codebases. Given a natural language query, the model retrieves the most semantically relevant functions from an indexed repository.19 20## Model Description21 22This is a bi-encoder trained with contrastive learning. It encodes both natural language queries and code snippets into a shared embedding space, enabling efficient cosine-similarity-based retrieval at search time.23 24- **Base model:** `microsoft/graphcodebert-base`25- **Architecture:** GraphCodeBERT encoder with mean pooling + L2 normalization (no LM head)26- **Languages trained on:** Erlang, C++27- **Task:** Semantic code search / function retrieval28 29### Architecture detail30 31The model wraps the GraphCodeBERT encoder in a lightweight `CodeSearchModel`:32 33```python34# Mean pooling over all token positions (not CLS)35def mean_pooling(last_hidden_state, attention_mask):36    mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()37    return torch.sum(last_hidden_state * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)38```39 40Embeddings are L2-normalized, so retrieval is a plain dot product (equivalent to cosine similarity).41 42---43 44## Training45 46### Data47 48Training triplets were constructed from two sources:49 50| Language | Source | Records |51|----------|--------|---------|52| C++ | [`codeparrot/xlcost-text-to-code`](https://huggingface.co/datasets/codeparrot/xlcost-text-to-code) (C++-program-level) | 8,650 |53| Erlang | Private dataset (not released) | — |54 55Each record is a `(code, good_docstring, bad1_docstring, bad2_docstring)` tuple. Negatives were mined as follows:56- **60% hard negatives** — BM25-retrieved docstrings that are lexically similar to the positive but semantically wrong (top-20 BM25 candidates, sampled randomly)57- **30% cross-language negatives** — docstrings sampled from the opposite language to discourage language-specific shortcuts58- **10% random negatives** — uniform random docstrings as easy negatives59 60### Loss61 62Temperature-scaled cross-entropy over augmented scores. For each batch the score matrix is extended with both negatives:63 64```65augmented_scores = [good_scores | bad1_scores | bad2_scores]66loss = CrossEntropyLoss(augmented_scores / τ, diagonal_labels)67```68 69where `τ = 0.05`.70 71### Hyperparameters72 73| Parameter | Value |74|-----------|-------|75| Base model | `microsoft/graphcodebert-base` |76| Batch size | 32 |77| Epochs | 10 |78| Learning rate | 2e-5 |79| LR schedule | Linear warmup (10%) → linear decay to 0 |80| Optimizer | AdamW |81| Gradient clipping | 1.0 |82| Code max length | 256 tokens |83| NL max length | 128 tokens |84| Temperature (τ) | 0.05 |85| Early stopping patience | 3 (not triggered) |86| Seed | 42 |87 88### Training curve89 90| Epoch | Loss |91|-------|------|92| 1 | 1.4135 |93| 2 | 0.4685 |94| 3 | 0.3438 |95| 4 | 0.2738 |96| 5 | 0.2308 |97| 6 | 0.1997 |98| 7 | 0.1671 |99| 8 | 0.1507 |100| 9 | 0.1425 |101| **10** | **0.1348** ← best |102 103Training ran for all 10 epochs without triggering early stopping (patience = 3). Best model saved at epoch 10.104 105---106 107## Usage108 109This model is intended to be used with [`code_search.py`](https://github.com/MatthewsO3/GraphCode-CErl-base/tree/main/Code%20Search), a unified indexing and search tool included in the repository.110 111### Quick start112 113```bash114git clone https://github.com/MatthewsO3/GraphCode-CErl-base115cd "GraphCode-CErl-base/Code Search/Evaluation"116python setup.py          # creates .venv, installs deps, builds erlang.so117source .venv/bin/activate118 119# Index a repository (auto-discovers Erlang + C++ + Python)120python code_search.py index \121    --repo /path/to/your/repo \122    --model MatthewsO3/GraphCode-CErl-codesearch \123    --output corpus.jsonl \124    --index corpus_index.pt125 126# Search interactively127python code_search.py search \128    --model MatthewsO3/GraphCode-CErl-codesearch \129    --jsonl corpus.jsonl \130    --index corpus_index.pt \131    --top 5132```133 134Language-specific flags are also available and can be combined freely:135 136```bash137# Erlang only138python code_search.py index --erlang /path/to/erl_repo ...139 140# C++ only141python code_search.py index --cpp /path/to/cpp_repo ...142 143# Explicit mix144python code_search.py index --erlang /path/erl --cpp /path/cpp --python /path/py ...145```146 147### Using the model directly148 149```python150from transformers import AutoTokenizer, AutoModel151import torch152 153tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base")154model = AutoModel.from_pretrained("MatthewsO3/GraphCode-CErl-codesearch")155model.eval()156 157def encode(texts):158    enc = tokenizer(texts, return_tensors="pt", truncation=True,159                    padding=True, max_length=256)160    with torch.no_grad():161        out = model(**enc)162    # Mean pooling163    mask = enc["attention_mask"].unsqueeze(-1).float()164    emb = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)165    return emb / emb.norm(dim=1, keepdim=True)166 167query = encode(["handle TCP connection timeout"])168code  = encode(["handle_timeout(Socket, State) -> gen_tcp:close(Socket), {stop, timeout, State}."])169 170score = (query @ code.T).item()171print(f"Similarity: {score:.4f}")172```173 174> **Note:** The tokenizer is loaded from `microsoft/graphcodebert-base` since it is identical to the fine-tuned model's tokenizer and avoids a redundant download.175 176---177 178## Supported Languages179 180| Language | Extractor | Extensions |181|----------|-----------|------------|182| Erlang | tree-sitter (WhatsApp grammar) + custom `ErlangParser` + regex fallback | `.erl`, `.hrl` |183| C++ | tree-sitter + regex fallback | `.cpp`, `.cc`, `.cxx`, `.c`, `.h`, `.hpp` |184| Python | tree-sitter + regex fallback | `.py` |185 186> **Note:** Python indexing is supported by `code_search.py` but the model was not trained on Python data. Results for Python queries may be less accurate.187 188---189 190## Limitations191 192- Not trained on Python — cross-language transfer to Python is best-effort193- The Erlang training set is private and not released194- Functions without docstrings or comments are embedded on code tokens alone, which may reduce retrieval accuracy for ambiguous natural language queries195- Running on CPU is fully supported but slow for large corpora at index-build time; a GPU is recommended196 197---198 199## Repository200 201Training code, indexing tool, and setup scripts are available at:202[github.com/MatthewsO3/GraphCode-CErl-base](https://github.com/MatthewsO3/GraphCode-CErl-base)203 204---205 206## Citation207 208If you use this model, please cite the original GraphCodeBERT paper:209 210```bibtex211@inproceedings{guo2021graphcodebert,212  title     = {GraphCodeBERT: Pre-training Code Representations with Data Flow},213  author    = {Guo, Daya and Ren, Shuo and Lu, Shuai and Feng, Zhangyin and Tang, Duyu214               and Liu, Shujie and Zhou, Long and Duan, Nan and Svyatkovskiy, Alexey215               and Fu, Shengyu and others},216  booktitle = {International Conference on Learning Representations},217  year      = {2021}218}219```