hasitha-gallella/live-graph-coder-v1

Live Graph Coder v1 - Beta release

Live Graph Coder  •  Try VS Code extension for free ↗

---

This Beta release contains only the mixed-precision 6-layer code embedding model: int4 weights with int8 for Q/K/V and embeddings. This attention-based code-search encoder turns code and natural-language queries into embeddings, so a question can find the code it is about. At 66 MB, it is light enough to run on a CPU.

Used in the VS Code extension Live Graph Coder (Beta V1): install it for free from the VS Code Marketplace.

Architecture

Layers	6 (post-norm transformer, no biases)
Hidden size	768
Attention	12 heads × 64, rotary positions (base 1000)
Feed-forward	SwiGLU, 3,072
Vocabulary	30,522 (lower-case WordPiece, vocab.txt)
Parameters	80.1M
Output	768 numbers, unit length
File	model_int4_qkv8.onnx, 66 MB: embeddings int8, Q/K/V 8-bit, other matrices 4-bit (blocks of 32)

Usage

Queries start with Represent this query for searching relevant code: ; code is embedded as is.

import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from transformers import BertTokenizerFast

repo = "hasitha-gallella/live-graph-coder-v1"
tokenizer = BertTokenizerFast(hf_hub_download(repo, "vocab.txt"), do_lower_case=True)
session = ort.InferenceSession(hf_hub_download(repo, "model_int4_qkv8.onnx"), providers=["CPUExecutionProvider"])

def embed(texts):
    batch = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="np")
    feeds = {k: batch[k].astype(np.int64) for k in ("input_ids", "attention_mask")}
    return session.run(None, feeds)[0]  # (n, 768), unit length

query = embed(["Represent this query for searching relevant code: read settings from a JSON file"])
code = embed(["def load_config(path):\n    with open(path) as f:\n        return json.load(f)",
              "def add(a, b):\n    return a + b"])
print(query @ code.T)  # about [[0.57, -0.03]]: the JSON loader matches

The same file runs in Node.js and in browsers with onnxruntime-web.

Results

Two benchmarks whose code was not used in training:

Benchmark	Measure	This model (66 MB)	12-layer CodeRankEmbed teacher (547 MB)
RepoQA (Python functions, found from a description)	Right function first / in the top 10	70% / 93%	77% / 94%
SWE-Atlas questions	Named code in the top 5	29%	36%

Training

The model was trained in two rounds (4,500 steps in total) on the code, docstrings and documentation of open-source repositories from SWE-bench Lite and SWE-bench Multimodal, released under the MIT, BSD, Apache-2.0 or Matplotlib licence: Astropy, Matplotlib, seaborn, Flask, Requests, xarray, pytest, scikit-learn, Sphinx, Chart.js, react-pdf, OpenLayers, Lighthouse, highlight.js and Prism. The training objective has three parts: matching the teacher's vectors, keeping its similarity ranking within each batch, and matching each docstring with its code.

References

1. The base design and the WordPiece vocabulary (vocab.txt). J. Devlin et al. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. 2018. arXiv:1810.04805.
2. The encoder design (rotary positions, SwiGLU feed-forward, no biases). Z. Nussbaum et al. Nomic Embed: Training a Reproducible Long Context Text Embedder. 2024. arXiv:2402.01613.
3. This model is distilled from CodeRankEmbed's layers 0, 2, 4, 6, 8 and 10 (trained on the CoRNStack code-retrieval data) and uses the CodeRankEmbed query prefix. T. Suresh et al. CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking. 2024. arXiv:2412.01007.