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%

Both benchmarks were run with the extension's full search (this model plus keyword and name matching). With 100 and 28 tasks, a difference of a few points is within noise.

Training

The model was trained in two rounds (4,500 steps in total) on the code, docstrings and documentation of 15 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.

Limitations

• It is an embedding model only; it does not generate text.
• Queries should be in English. Each text is limited to 512 tokens, so embed functions or classes rather than whole files.
• It may reflect the conventions of the projects it was trained on.

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.