hasitha-gallella/live-graph-coder-v1

Live Graph Coder v1 - Beta release

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. 66 MB, light weight to run on a CPU.

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

Architecture

text or code
   │  WordPiece tokenizer (lower-case, 30,522 tokens), up to 512 tokens
   ▼
token embedding 30,528 × 768 (int8) + token-type embedding → LayerNorm
   │
   ▼
   ┌── transformer layer, repeated 6 times ───────────────────
   │  x ─► Q, K, V = x · W_qkv         768 → 3 × 768   (8-bit)
   │       12 heads × 64, rotary positions on Q and K
   │       attention = softmax(Q·Kᵀ / √64) · V
   │    ─► output projection           768 → 768       (4-bit)
   │    ─► h = LayerNorm(x + attention output)
   │  h ─► SwiGLU: W₁h ⊙ SiLU(W₂h)     768 → 3,072     (4-bit)
   │    ─► down projection             3,072 → 768     (4-bit)
   │    ─► LayerNorm(h + feed-forward output)
   └─────────────────────────────────────────────────────────────────
   │
   ▼
[CLS] vector → L2 normalisation → embedding (768)

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 it 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 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 use the extension's full search (this model, plus keyword and name matching). With 100 and 28 tasks, a few points either way is noise.

Training

Trained in two rounds (4,500 steps) on the code, docstrings and documentation of open-source repositories (from SWE-bench Multimodal and SWE-bench Lite) under MIT, BSD, Apache-2.0 or the Matplotlib licence: astropy, matplotlib, seaborn, flask, requests, xarray, pytest, scikit-learn, sphinx, Chart.js, react-pdf, openlayers, lighthouse, highlight.js, prism. The loss: match the teacher's vectors, keep its similarity ranking within each batch, and match each docstring with its code.

Limitations

• An embedding model only: no text generation.
• English questions; at most 512 tokens per text, 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) [3]. 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 the CodeRankEmbed query prefix. T. Suresh et al. CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking. 2024. arXiv:2412.01007.