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.

[ajianaz]

Hey @Mahnoor-Zaffar! Good news — #499 (the dependency for incremental per-symbol embedding) is already merged, so this issue is unblocked and ready to go. Appreciate the interest in picking it up! 🙌

A few things to keep in mind as you work on this:

1. This is a full feature, not a bugfix
~410 lines across 6 files with two new dependencies (ort, tokenizers). The scope is well-defined in the issue body, but reach out if anything's unclear.

2. ONNX must be opt-in — this is non-negotiable
Cora runs on everything from beefy workstations to 2GB VPS containers. ONNX inference is heavy (~154MB model in RAM, ~50-100ms per symbol on CPU). To keep low-spec machines unaffected:

• Feature-gate at compile time: onnx-embed feature, NOT in default — same pattern as the existing pretrained-embed flag
• Config opt-in at runtime: brain.embedding: onnx in .cora.yaml, default stays auto (pretrained → hashing fallback chain)
• Lazy model download: user must explicitly run cora brain --setup — never auto-download 154MB
• Graceful fallback: if model file is missing → warn + fall back to pretrained/hashing, never crash

3. Non-brain operations must have zero overhead
cora review, cora scan, cora index (non-brain) should not load or initialize the ONNX runtime at all. The session should only be created lazily when brain mode is actually invoked.

If that all sounds good, go for it! Feel free to open a draft PR early so we can iterate on the architecture before the full implementation. 🚀

[ajianaz]

Quick note: this is a free open-source project — there's no compensation or bounty for working on this. If you're still interested after reading that, just let us know! If not, no worries at all — we'll get to it ourselves. 😊

[ajianaz]

⚠️ Incorrectly closed by PR #513

This issue was auto-closed because PR #513 (feat(license): relicense MIT to Apache-2.0 + add CLA v1.1) included Closes #500 in its PR body.

PR #513 contains: Relicense MIT → Apache-2.0 + CLA documents (legal/licensing)
Issue #500 contains: ONNX runtime integration with Jina v2 base code embeddings (feature/engineering)

These two are completely unrelated. Likely a typo in the issue number in the PR body.

This issue should remain open with status:deferred until the ONNX integration is actually implemented.

---

Reopened on 2026-08-10 — clarified incorrect auto-close from unrelated PR.

[Mahnoor-Zaffar]

Sure, I will be happy to take it

[ajianaz]

Spike result: jina-v2-base-code q8 ONNX on CPU (measured, 2026-08-28)

Ran a standalone benchmark (onnxruntime 1.22.1, CPUExecutionProvider, q8 model from HF, real Rust-style code text, Apple Silicon 8-core) before any integration — the 50-100 ms/symbol assumption needed validation.

config	ms/seq
4 threads, seq 128	102
4 threads, seq 256	273
4 threads, seq 512	378
8 threads, batch 32, seq 256	160 (best batched)
CoreML EP	278-489 (worse: q8 model maps only 938/1278 nodes)

Session load: ~350 ms (fine). Pooling → 768-dim L2-normalized: verified.

Verdict vs the decision gate (≤40 ms/seq @4-core → integrate)

FAILS at 2.5-7x over. Full-embed math at ~130-270 ms/seq: 5k symbols ≈ 11-23 min one-time; 50k ≈ 1.5-4 h. Query-time single embed (~100-270 ms) is fine interactively, and embed_fingerprint keeps incremental runs cheap — but a tier that costs hours on monorepos cannot be the default-on quality tier on CPU.

Implications

1. jina-v2-base-code stays viable ONLY as an opt-in "max quality" tier with background-indexing UX (progress + --watch-style async), or on FP16 + CoreML/ANE (640 MB download, untested).
2. The ~150M-param transformer class has a fundamental ~100 ms/seq wall on CPU ORT — models like nomic-embed-text (137M) land in the same class. Going meaningfully faster means ~22M-class models (MiniLM) with a real quality drop, or a different runtime (CoreML FP16 / DirectML / int4).
3. EmbeddingGemma-300m Q4: 2x params + int4 — likely lands in the same wall class; needs its own spike before betting (tracked in research: small CPU-only code embeddings + adoptable optimizations from peer tools (embed tier roadmap) #544).
4. Suggestion: re-scope this issue to "tier-2 = opt-in async background tier" OR evaluate smaller distills (e.g., 30-80M code embeds) as the default-on tier-2. vecq store (feat(brain): vecq as opt-in vector backend for v0.15.0 (replace usearch long-term) #542) composes with any outcome.

Spike scripts: /tmp/jina-spike (throwaway); happy to re-run on Linux/Windows CI runners for a cross-platform table.