onnx

package module
v0.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Aug 9, 2026 License: MIT Imports: 10 Imported by: 0

README

provider/local/onnx

Learned sparse encoding in your own process, through ONNX Runtime. No server, no network, no Python at runtime.

This directory is a separate Go module. go get github.com/regularkevvv/agentic does not pull it, nothing in the root module's build reaches it, and it is deliberately absent from go.work — working on it means cd provider/local/onnx. That is what keeps CGO out of the default contributor loop and out of every application that imports Agentic.

Install this module independently with:

go get github.com/regularkevvv/agentic/provider/local/onnx@v0.1.0

Two things must be on disk before anything compiles

Both are discovered at link time otherwise, as library 'tokenizers' not found or as a runtime failure to load libonnxruntime. Neither is a Go dependency and neither is fetched by go build.

1. libtokenizers.a, per platform

github.com/daulet/tokenizers is a CGO wrapper around the Rust crate transformers uses. The Go side is on proxy.golang.org; the compiled static library is a release asset you fetch yourself.

cd provider/local/onnx
mkdir -p lib
# darwin/arm64 — the only platform this has been verified on
curl -L https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.darwin-arm64.tar.gz | tar xz -C lib
# linux/amd64
curl -L https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.linux-amd64.tar.gz | tar xz -C lib

export CGO_LDFLAGS="-L$PWD/lib"

-L is all that is needed. The binding's own #cgo LDFLAGS already supply -ltokenizers -ldl -lm -lstdc++, and repeating them only produces ld: warning: ignoring duplicate libraries.

The release publishes darwin-{arm64,x86_64}, linux-{amd64,arm64,aarch64, ppc64le,s390x}, and musl variants of each. Pick the one matching where the binary will run, not where it is built: this links statically.

2. ONNX Runtime, a separate download
# darwin/arm64
curl -L -O https://github.com/microsoft/onnxruntime/releases/download/v1.28.0/onnxruntime-osx-arm64-1.28.0.tgz
tar xzf onnxruntime-osx-arm64-1.28.0.tgz
export AGENTIC_ONNX_LIBRARY=$PWD/onnxruntime-osx-arm64-1.28.0/lib/libonnxruntime.dylib

# linux/amd64
curl -L -O https://github.com/microsoft/onnxruntime/releases/download/v1.28.0/onnxruntime-linux-x64-1.28.0.tgz
tar xzf onnxruntime-linux-x64-1.28.0.tgz
export AGENTIC_ONNX_LIBRARY=$PWD/onnxruntime-linux-x64-1.28.0/lib/libonnxruntime.so

This one is loaded at runtime rather than linked, so it can also be pointed at with onnx.WithLibraryPath. The environment variable exists because ONNX Runtime has no standard install location on any platform here; the option always wins over it.

What all of that costs

Measured 2026-08-01 on darwin/arm64.

Artifact Size Where it goes
ONNX Runtime, macOS arm64 32 MB download, 39 MB .dylib loaded at runtime
ONNX Runtime, Linux x64 9 MB download loaded at runtime
libtokenizers.a 14 MB download, 39 MB on disk ~18 MB linked into your binary
Exported Granite graph ~117 MiB (124 MB on disk) read at construction

The 18 MB is the difference between this package's test binary (29 MB) and a comparable pure-Go binary that imports Agentic (12 MB). It is static: the resulting binary needs no libtokenizers at runtime.

Getting the model

Granite ships no ONNX export, so producing one is a documented one-time step rather than a download. provider/local/onnx/export_onnx.py does it and writes the PyTorch reference this package's tests assert against:

uv run provider/local/onnx/export_onnx.py --out /some/writable/directory

One command, and it leaves nothing behind: the script declares its dependencies in a PEP 723 header, so uv builds a throwaway environment pinned by export_onnx.py.lock and discards it afterwards. Python is a build tool here, in the same sense a compiler is — this package links against no Python and executes none. Without uv, pip install torch==2.13.0 transformers onnx onnxscript does the same thing unpinned.

The 117 MiB graph cannot live in the repository, which is why the live tests are gated on a local path. See docs/multi-representation-inference.md for the export step and the runtime setup in one place.

Using it

import (
    agentic "github.com/regularkevvv/agentic"
    "github.com/regularkevvv/agentic/provider/local/onnx"
)

encoder, err := onnx.New(
    "/models/granite-sparse.onnx",
    "/models/tokenizer.json",
    agentic.VectorSpace{
        Model:    "ibm-granite/granite-embedding-30m-sparse",
        Revision: "ad82b1fd09541c998c8d45045d601c51fdb8a9b7",
    },
)
if err != nil {
    return err
}
defer encoder.Close()

resp, err := agentic.EncodeDocuments(ctx, encoder, docs, agentic.RepresentationSparse)

Close is not optional. The session and the tokenizer hold memory outside Go's heap that no garbage collector reclaims.

The vector space is yours to declare, and the parts it can prove are filled in: Dimensions comes from the graph's own output width, Provider defaults to onnx, Metric to dot product, and the ID is the canonical hash of the rest. Record Revision — an index keyed on an unrevisioned space cannot detect a silent re-export.

What it does and does not do

Sparse only. Dense and multi-vector requests return agentic.ErrUnsupportedRepresentation rather than an answer from the wrong reduction. The pooling is SPLADE's — log1p(relu(logits)), masked, max over positions — the reduction the model's own sentence-transformers pooling config declares, and the reason a document about an "automobile" carries weight on "car".

Nothing is downloaded. Both constructor arguments are filesystem paths. There is no HTTP client in this package.

Over-long inputs are rejected, not truncated. The limit is the model's positional bound, 512 tokens by default; past it the graph indexes off the end of its position table rather than degrading. Silently dropping the tail of a document would produce a vector for text nobody asked about.

Batch width is the cost, not batch size. Every row in a forward pass is padded to the widest one. Inputs are therefore ordered by token length and grouped with neighbors, so no row is padded past twice its own length and a long document never drags short ones up to its width. A second ceiling caps the logits tensor a pass may allocate — at Granite's vocabulary a single 512-token input is already 103 MB of logits.

Calls are serialized. ONNX Runtime already spreads one graph across every core it is configured for, so concurrent Encode calls would contend rather than scale.

Running the tests

The unit tests need the tokenizer library on the link line and nothing else:

CGO_LDFLAGS="-L$PWD/lib" GOWORK=off go test ./...

The live tests are the gate for this package. They skip cleanly without their three variables and assert against testdata/granite_reference.json:

Variable Meaning
AGENTIC_ONNX_MODEL the exported .onnx graph
AGENTIC_ONNX_TOKENIZER the matching tokenizer.json
AGENTIC_ONNX_LIBRARY libonnxruntime.dylib or libonnxruntime.so
export AGENTIC_ONNX_MODEL=/models/granite-sparse.onnx
export AGENTIC_ONNX_TOKENIZER=/models/tokenizer.json
export AGENTIC_ONNX_LIBRARY=/opt/onnxruntime/lib/libonnxruntime.dylib
CGO_LDFLAGS="-L$PWD/lib" GOWORK=off go test -v -run Live ./...

They check four separable claims, so a failure says which one broke:

  • the Go tokenizer reproduces transformers' input ids from the raw text;
  • every coordinate index matches PyTorch's and every weight is within 1e-4 — measured 4.56e-06 on 2026-08-01, which is why the gate is 1e-4 and not looser;
  • a row padded to twice its own length carries the same coordinate indices as the same row alone, and weights differing by at most 1.46e-06 — the wider forward pass's arithmetic rather than the padding;
  • the padding id does not reach the result: that same padded row is identical to the bit under two padding ids, which is what says the mask is doing the excluding.

Platform support and licenses

darwin/arm64 and linux/amd64 have both been verified end to end against the real exported graph, on 2026-08-01. Every coordinate index matched the PyTorch golden on both, and the largest weight difference was 4.56e-06 on darwin/arm64 and 4.38e-06 on linux/amd64 — the two platforms disagree with PyTorch by about as much as they disagree with each other, and both sit far inside the 1e-4 the live test gates on.

No other platform has been run. The bindings publish artifacts for linux/arm64, musl, and darwin/x86_64 among others; those are untested here.

Both bindings are MIT: github.com/daulet/tokenizers v1.27.0 and github.com/yalue/onnxruntime_go v1.31.0. Both are pinned exactly. ONNX Runtime itself is MIT and is not vendored here.

Documentation

Overview

Package onnx implements agentic.RepresentationEncoder in this process, through ONNX Runtime, with no server and no network. This directory is a separate Go module and deliberately not part of github.com/regularkevvv/agentic, because it requires CGO, a native ONNX Runtime shared library, and a statically linked tokenizer — none of which can be a condition of using the library. That is why `go get github.com/regularkevvv/agentic` never pulls this package, and why a directory under provider/ is absent from the root module's build. Working on it means `cd provider/local/onnx`, with the setup in README.md done first.

What it produces

Learned sparse vectors, and nothing else. Dense and multi-vector requests return agentic.UnsupportedRepresentationError rather than an answer derived from the wrong reduction. The target is the SPLADE family — a masked-language -model head pooled into vocabulary weights — which is what makes a document about an "automobile" carry weight on "car", a word it never contained.

Nothing is downloaded

New takes filesystem paths and reads them. There is no model cache, no registry lookup, and no HTTP client in this package; a model you have not already exported is an error rather than a fetch. Producing the graph is a documented one-time step — see provider/local/onnx/export_onnx.py — and keeping it a step is the point: a 117 MiB artifact that appears by surprise during a test run is not a dependency anyone agreed to.

Batch width, not batch size, is the cost

Every row in one forward pass is padded to the widest row in it, and padding buys compute nobody asked for. Measured on 2026-08-01, three short inputs padded to a common width of 18 took 20 ms as one call against 13 ms as three. The encoder therefore orders inputs by token length and groups neighbors, so a long document never drags short ones up to its width and no row is ever padded past twice its own length. The rule, and the bound it buys, are in batch.go.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

This section is empty.

Types

type Encoder

type Encoder struct {
	// contains filtered or unexported fields
}

Encoder runs an exported masked-language-model graph and pools its logits into a learned sparse vector.

It owns memory outside Go's heap and must be closed; see Encoder.Close.

func New

func New(modelPath, tokenizerPath string, space agentic.VectorSpace, opts ...Option) (encoder *Encoder, err error)

New loads an exported SPLADE-family graph and its tokenizer and encodes with them in this process.

modelPath is an ONNX file whose graph takes int64 input_ids and attention_mask and returns float32 logits of shape [batch, sequence, vocabulary], with the batch and sequence axes dynamic. provider/local/onnx/export_onnx.py produces exactly that. Pooling is deliberately outside the graph, so what the model contributes and what this package contributes stay separable.

tokenizerPath is a Hugging Face tokenizer.json — the portable fast-tokenizer format, not a directory and not a sentencepiece model.

space is the identity that will be persisted beside every vector. Provider defaults to "onnx", Kind must be sparse, Metric defaults to dot product because a learned sparse weight is part of the score rather than a direction, Dimensions is filled from the graph's own output width when left zero, and ID is derived from the rest when left empty. Model is required: a file name is not an identity, since two exports of different weights can share one.

Nothing is downloaded, then or later. Both paths are read from the local filesystem and this package contains no HTTP client at all.

The graph is loaded twice here — once to read the vocabulary from its output shape, once for the session that will run it — which is what makes the vocabulary observed rather than assumed. Measured on darwin/arm64 against the 117 MiB Granite export, that costs 0.12 s of a 0.30 s construction. An encoder is meant to outlive the request that needed it.

func (*Encoder) Capabilities

func (e *Encoder) Capabilities() agentic.RepresentationCapabilities

Capabilities implements agentic.RepresentationEncoder.

Truncation is not advertised. The exported graph accepts sequences up to the model's positional limit and nothing beyond it, and an over-long input is rejected with its token count rather than clipped — silently dropping the end of a document produces a vector for text the caller never asked about.

MaximumBatchSize is zero because this encoder splits a request into forward passes itself; the caller's batch is a request shape, not a hardware one.

An empty sparse vector is not allowed. A SPLADE head that predicts no vocabulary entry at all for a non-empty input is a broken graph rather than a short document, and storing the empty vector would put an unmatchable row in an index instead of failing.

func (*Encoder) Close

func (e *Encoder) Close() error

Close releases the session and the tokenizer, and is safe to call twice.

Both hold memory outside Go's heap that no garbage collector reclaims, so an encoder that is never closed leaks a loaded model for the life of the process. Encode after Close returns an error rather than entering a destroyed session, which would be a crash.

The process-global ONNX Runtime environment is deliberately left up: another encoder may still be using it. A program that wants it down calls onnxruntime_go.DestroyEnvironment after closing every encoder.

func (*Encoder) Encode

Encode implements agentic.RepresentationEncoder.

Cancellation is observed between forward passes rather than inside one. ONNX Runtime's Run has no cancellation, so a pass that has already started runs to completion; ctx bounds how many more begin.

func (*Encoder) Name

func (e *Encoder) Name() string

Name implements agentic.RepresentationEncoder. It reports the model identifier from the vector space, which is the only model name this package has: an ONNX file declares no identity of its own.

type Option

type Option func(*config)

Option configures an Encoder.

func WithLibraryPath

func WithLibraryPath(path string) Option

WithLibraryPath points ONNX Runtime at its shared library — libonnxruntime.dylib, libonnxruntime.so, or onnxruntime.dll.

Without it the AGENTIC_ONNX_LIBRARY environment variable is used, and without that the binding falls back to the platform's default name, which resolves only if the library is already on the loader's search path.

The runtime environment is global to the process: whichever encoder is constructed first settles the path, and a later different one is ignored.

func WithLimits

func WithLimits(limits agentic.RepresentationLimits) Option

WithLimits replaces the request-size ceilings, which default to agentic.DefaultRepresentationLimits.

The response-side ceilings still apply. MaxSparseNonZero is the one that matters here: a SPLADE head is sparse by training rather than by construction, so a long document can carry thousands of nonzero coordinates and nothing in the graph bounds that.

func WithMaxBatchBytes

func WithMaxBatchBytes(bytes int) Option

WithMaxBatchBytes caps the logits tensor a single forward pass allocates, defaulting to 256 MiB.

See [defaultBatchBytes] for what that buys. Lower it on a memory-constrained machine; raising it does not make encoding faster, because the encoder groups by token length and a group is closed by padding waste before it is closed by this ceiling for any realistic mix of inputs.

func WithMaxTokens

func WithMaxTokens(tokens int) Option

WithMaxTokens sets the longest tokenized input the encoder will accept, defaulting to 512.

Inputs above it are rejected with their token count rather than truncated. Raise it only for a model whose positional embeddings actually extend that far; the limit exists because exceeding it is a runtime fault inside the graph rather than a degraded result.

func WithPadTokenID

func WithPadTokenID(id int64) Option

WithPadTokenID sets the id written into padded positions, defaulting to 0.

The value does not affect the result: padded positions are masked out of attention and out of pooling, which the live tests assert by encoding one padded row under two padding ids and comparing every coordinate. It still has to be a token the model has, since it indexes the embedding table, and zero is the only id every vocabulary contains. Set the model's real padding id if you would rather the tensors read the way its own tooling writes them.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL