--- language: - code license: mit base_model: microsoft/graphcodebert-base tags: - code-search - semantic-search - graphcodebert - erlang - cpp library_name: transformers pipeline_tag: feature-extraction --- # GraphCode-CErl — Semantic Code Search for Erlang & C++ Fine-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. ## Model Description This 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. - **Base model:** `microsoft/graphcodebert-base` - **Architecture:** GraphCodeBERT encoder with mean pooling + L2 normalization (no LM head) - **Languages trained on:** Erlang, C++ - **Task:** Semantic code search / function retrieval ### Architecture detail The model wraps the GraphCodeBERT encoder in a lightweight `CodeSearchModel`: ```python # Mean pooling over all token positions (not CLS) def mean_pooling(last_hidden_state, attention_mask): mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float() return torch.sum(last_hidden_state * mask, 1) / torch.clamp(mask.sum(1), min=1e-9) ``` Embeddings are L2-normalized, so retrieval is a plain dot product (equivalent to cosine similarity). --- ## Training ### Data Training triplets were constructed from two sources: | Language | Source | Records | |----------|--------|---------| | C++ | [`codeparrot/xlcost-text-to-code`](https://huggingface.co/datasets/codeparrot/xlcost-text-to-code) (C++-program-level) | 8,650 | | Erlang | Private dataset (not released) | — | Each record is a `(code, good_docstring, bad1_docstring, bad2_docstring)` tuple. Negatives were mined as follows: - **60% hard negatives** — BM25-retrieved docstrings that are lexically similar to the positive but semantically wrong (top-20 BM25 candidates, sampled randomly) - **30% cross-language negatives** — docstrings sampled from the opposite language to discourage language-specific shortcuts - **10% random negatives** — uniform random docstrings as easy negatives ### Loss Temperature-scaled cross-entropy over augmented scores. For each batch the score matrix is extended with both negatives: ``` augmented_scores = [good_scores | bad1_scores | bad2_scores] loss = CrossEntropyLoss(augmented_scores / τ, diagonal_labels) ``` where `τ = 0.05`. ### Hyperparameters | Parameter | Value | |-----------|-------| | Base model | `microsoft/graphcodebert-base` | | Batch size | 32 | | Epochs | 10 | | Learning rate | 2e-5 | | LR schedule | Linear warmup (10%) → linear decay to 0 | | Optimizer | AdamW | | Gradient clipping | 1.0 | | Code max length | 256 tokens | | NL max length | 128 tokens | | Temperature (τ) | 0.05 | | Early stopping patience | 3 (not triggered) | | Seed | 42 | ### Training curve | Epoch | Loss | |-------|------| | 1 | 1.4135 | | 2 | 0.4685 | | 3 | 0.3438 | | 4 | 0.2738 | | 5 | 0.2308 | | 6 | 0.1997 | | 7 | 0.1671 | | 8 | 0.1507 | | 9 | 0.1425 | | **10** | **0.1348** ← best | Training ran for all 10 epochs without triggering early stopping (patience = 3). Best model saved at epoch 10. --- ## Usage This 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. ### Quick start ```bash git clone https://github.com/MatthewsO3/GraphCode-CErl-base cd "GraphCode-CErl-base/Code Search/Evaluation" python setup.py # creates .venv, installs deps, builds erlang.so source .venv/bin/activate # Index a repository (auto-discovers Erlang + C++ + Python) python code_search.py index \ --repo /path/to/your/repo \ --model MatthewsO3/GraphCode-CErl-codesearch \ --output corpus.jsonl \ --index corpus_index.pt # Search interactively python code_search.py search \ --model MatthewsO3/GraphCode-CErl-codesearch \ --jsonl corpus.jsonl \ --index corpus_index.pt \ --top 5 ``` Language-specific flags are also available and can be combined freely: ```bash # Erlang only python code_search.py index --erlang /path/to/erl_repo ... # C++ only python code_search.py index --cpp /path/to/cpp_repo ... # Explicit mix python code_search.py index --erlang /path/erl --cpp /path/cpp --python /path/py ... ``` ### Using the model directly ```python from transformers import AutoTokenizer, AutoModel import torch tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base") model = AutoModel.from_pretrained("MatthewsO3/GraphCode-CErl-codesearch") model.eval() def encode(texts): enc = tokenizer(texts, return_tensors="pt", truncation=True, padding=True, max_length=256) with torch.no_grad(): out = model(**enc) # Mean pooling mask = enc["attention_mask"].unsqueeze(-1).float() emb = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9) return emb / emb.norm(dim=1, keepdim=True) query = encode(["handle TCP connection timeout"]) code = encode(["handle_timeout(Socket, State) -> gen_tcp:close(Socket), {stop, timeout, State}."]) score = (query @ code.T).item() print(f"Similarity: {score:.4f}") ``` > **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. --- ## Supported Languages | Language | Extractor | Extensions | |----------|-----------|------------| | Erlang | tree-sitter (WhatsApp grammar) + custom `ErlangParser` + regex fallback | `.erl`, `.hrl` | | C++ | tree-sitter + regex fallback | `.cpp`, `.cc`, `.cxx`, `.c`, `.h`, `.hpp` | | Python | tree-sitter + regex fallback | `.py` | > **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. --- ## Limitations - Not trained on Python — cross-language transfer to Python is best-effort - The Erlang training set is private and not released - Functions without docstrings or comments are embedded on code tokens alone, which may reduce retrieval accuracy for ambiguous natural language queries - Running on CPU is fully supported but slow for large corpora at index-build time; a GPU is recommended --- ## Repository Training code, indexing tool, and setup scripts are available at: [github.com/MatthewsO3/GraphCode-CErl-base](https://github.com/MatthewsO3/GraphCode-CErl-base) --- ## Citation If you use this model, please cite the original GraphCodeBERT paper: ```bibtex @inproceedings{guo2021graphcodebert, title = {GraphCodeBERT: Pre-training Code Representations with Data Flow}, author = {Guo, Daya and Ren, Shuo and Lu, Shuai and Feng, Zhangyin and Tang, Duyu and Liu, Shujie and Zhou, Long and Duan, Nan and Svyatkovskiy, Alexey and Fu, Shengyu and others}, booktitle = {International Conference on Learning Representations}, year = {2021} } ```