feat(brain): ONNX runtime integration with Jina v2 base code embeddings

[ajianaz]

Summary

Add ONNX Runtime-based contextual code embeddings using jina-embeddings-v2-base-code (154MB quantized ONNX, Apache-2.0). This provides full transformer inference for Brain Mode semantic search, dramatically improving query relevance over static lookup methods.

Background

Corona-code currently has two embedding backends (neither is a full neural network):

1. Hashing trick 256d — token hash → pseudo-random vector. Zero dependency, very fast, but low quality (~15-20 estimated CoIR NDCG@10).
2. Nomic distilled 768d — static token→vector lookup table (40,856 tokens × 768d int8 = 31MB embedded in binary). Medium quality (~30-35 estimated CoIR NDCG@10). Not contextual — same token always gets same vector regardless of surrounding code.

Both lack contextual understanding — a token like render gets the same embedding whether it appears in render_template, render_html, or render_pipeline.

Proposed: Tier 3 — ONNX Full Transformer

Model: jinaai/jina-embeddings-v2-base-code — quantized ONNX variant.

Property	Value
ONNX size (int8 quantized)	154 MB
Embedding dimension	768
Max sequence length	8,192 tokens
License	Apache-2.0 ✅ (commercial-safe)
Code-specific training	✅ (30+ code languages)
Tokenizer	BPE (tokenizer.json — HuggingFace standard)
Architecture	BERT 12-layer, mean pooling, L2 normalization
CoIR NDCG@10 (estimated)	~57

Why Jina v2 base code?

Evaluated 9 models. Jina v2 base code is the only model that passes all criteria:

Model	ONNX?	License	Size	Verdict
jina-v2-base-code	✅ 3 variants	✅ Apache-2.0	154MB (q8)	CHOSEN
jina-code-0.5b	❌	⚠️ CC-BY-NC-4.0	942MB	License blocked
codet5p-110m	❌	BSD-3	439MB	No ONNX, seq=512
SFR-Code-400M_R	✅	⚠️ CC-BY-NC-4.0	424MB (q8)	License blocked
nomic-embed-code	❌	Apache-2.0	27 GB	Non-starter
bge-code-v1	❌	MIT	5.9 GB	Non-starter

Proposed Changes

1. New Dependencies (Cargo.toml)

[dependencies]
ort = { version = "2.0", features = ["half"] }       # ONNX Runtime (CPU)
tokenizers = { version = "0.21", features = ["onnx"] } # HuggingFace BPE

Note: ort links against ONNX Runtime shared library (~15MB). Consider static linking or bundling.

2. New Module: src/embed/onnx.rs

pub struct OnnxEmbedder {
    session: ort::Session,        // ONNX inference session
    tokenizer: tokenizers::Tokenizer,
    dim: usize,                   // 768
}

impl OnnxEmbedder {
    pub fn load(model_path: &Path) -> Result<Self>;
    pub fn embed(&self, code: &str) -> Result<Vec<f32>>;  // BPE → forward pass → mean pool → L2 norm
}

3. Lazy Model Download

$ cora brain --setup
🧠 Brain Mode Setup

Select embedding backend:
  > auto (recommended)
    pretrained (nomic 768d, embedded)
    onnx (Jina v2 base code, 154MB download)
    hashing (256d, zero dependency)

Downloading jina-embeddings-v2-base-code (quantized)...
  → ~/.local/share/cora/models/jina-code-q8.onnx (154 MB)
  → ~/.local/share/cora/models/jina-tokenizer.json (1.2 MB)
✅ Done! Brain Mode will use ONNX embeddings.

Download URL: https://huggingface.co/jinaai/jina-embeddings-v2-base-code/resolve/main/onnx/model_quantized.onnx

4. Config

# .cora.yaml
brain:
  embedding: onnx
  onnx_model_path: ~/.local/share/cora/models/jina-code-q8.onnx

5. Incremental Embedding (builds on #499)

With 154MB model + 50-100ms inference per symbol, full re-embed is expensive:

Project	Full re-embed	Incremental (1 file)
Uteke (1106 symbols)	~110 seconds	~5-10 seconds
Corin (537 symbols)	~54 seconds	~3-5 seconds

Per-symbol fingerprint tracking from #499 is mandatory for ONNX backend.

6. Brain Search Query Embedding

Brain search must embed the query string using the same ONNX model:

fn vector_search(conn, project_id, query, limit) {
    let query_vec = match active_backend() {
        Backend::Onnx(model) => model.embed(query),     // ~50-100ms
        Backend::Pretrained => embed_code_pretrained(query), // ~0.01ms
        Backend::Hashing => tokens::embed_code(query),       // ~0.01ms
    };
    // usearch KNN search...
}

Query embedding latency: ~50-100ms per cora brain invocation. Acceptable for interactive CLI.

Performance Targets

Metric	Target
Model download	< 2 minutes on broadband
Full index (1106 symbols)	< 120 seconds
Incremental (1 file, ~10 symbols)	< 5 seconds
Brain search query	< 300ms (including query embed)
Memory usage	< 500MB during inference

Risk Assessment

Risk	Mitigation
ort crate binary size (+15-20MB)	Feature-gate behind onnx-embed
ONNX Runtime platform differences	CI test on Linux/macOS/Windows
Model download failures	Retry + resume + checksum verify
Inference too slow on old CPUs	Benchmark on minimum hardware, document requirements
Tokenizer mismatch	Use official tokenizer.json from HF repo

Dependencies

• Requires feat(brain): runtime embedding config + incremental per-symbol embedding #499 — incremental per-symbol embedding must exist first
• Requires v0.13.0 release — config schema + runtime dispatch must be in place

Files to Create/Change

File	Change	Est. Lines
src/embed/onnx.rs	NEW — ONNX inference module	~200
src/embed/mod.rs	Add OnnxBackend to dispatch	~30
src/commands/brain_setup.rs	NEW — interactive setup wizard	~150
src/config/schema.rs	Add ONNX fields to BrainConfig	~10
Cargo.toml	Add ort, tokenizers, onnx-embed feature	~5
src/index/brain.rs	Query embedding via active backend	~15

Total: ~410 lines new/changed code.

Testing

• cargo test --features tree-sitter,onnx-embed
• Model download + load works on Linux/macOS/Windows
• Embedding output matches Python reference (cosine sim > 0.99)
• Brain search with ONNX returns vector signal with high relevance
• Incremental embed: only changed symbols re-embedded
• Memory stays under 500MB during full index
• Config: embedding: onnx without model file → helpful error message

[Mahnoor-Zaffar]

Hey! Thanks for reporting this.

The issue proposes the integration of ONNX Runtime-based contextual code embeddings into the Corona-code repository, specifically using the jina-embeddings-v2-base-code model. The expected behavior is to enhance query relevance in semantic searches by providing contextual understanding of code tokens. Currently, the observed behavior is limited to two embedding methods that lack contextual awareness, resulting in lower quality embeddings.

The most likely root cause of the issue is the absence of a full neural network embedding backend that can provide contextual embeddings. The current implementations (hashing trick and static lookup) do not account for the context in which tokens appear. The affected code primarily involves the embedding logic, which would be implemented in the proposed new module src/embed/onnx.rs and the modifications to Cargo.toml to include dependencies for ONNX Runtime and tokenizers.

If no one is currently working on this, I'd be happy to take a look and put together a fix.