embed

package
v0.3.0 Latest Latest
Warning

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

Go to latest
Published: Sep 4, 2026 License: Apache-2.0 Imports: 18 Imported by: 0

Documentation

Overview

Package embed generates sentence embeddings with a local ONNX model (arctic-embed-xs by default; see models.go for the registry). Inference runs fully in-process against a single ONNX Runtime session — no network, no API key. The runtime library and model files are downloaded once (see download.go) and cached under the OS cache dir (see paths.go).

Index

Constants

View Source
const DefaultModel = "arctic-embed-xs"

DefaultModel is what a caller who expresses no preference gets.

View Source
const OrtVersion = "1.29.0"

OrtVersion is the ONNX Runtime release whose C API onnxruntime_go v1.33.0 is built against. The two are coupled: the binding compiles against one version of the headers and dlopens whatever this constant downloaded, so bumping the Go module without bumping this constant produces a binary that loads a library it was not compiled for. Nothing in CI catches that — the tests skip inference when no model is installed — so the versions move together, in one commit, or not at all.

Variables

This section is empty.

Functions

func Check

func Check() error

Check reports whether the ONNX runtime library and model files are all present on disk, without initializing the runtime. Returns nil when embedding is ready, or an error naming the first missing piece — suitable for `semantic status`. Get performs the same checks lazily.

func CosineSim

func CosineSim(a, b Vec) float64

CosineSim returns cosine similarity in [-1, 1]. Vectors from Get are already L2-normalized, so for those this reduces to the dot product, but the full formula is kept so callers can pass un-normalized inputs.

func DownloadAll

func DownloadAll(logf func(string, ...any)) error

DownloadAll downloads the ONNX Runtime library and model files. logf receives progress messages.

func DownloadModel

func DownloadModel(logf func(string, ...any)) error

DownloadModel downloads and caches model.onnx and tokenizer.json.

func DownloadOrt

func DownloadOrt(logf func(string, ...any)) error

DownloadOrt downloads and caches the ONNX Runtime shared library.

func EnsureModel added in v0.1.3

func EnsureModel(logf func(string, ...any)) error

EnsureModel downloads whatever Check reports missing, and does nothing when everything is already cached. It exists so a command that is about to embed can heal itself instead of failing with an instruction to run `init` and be run a second time.

A checkpoint change is the case it is really for: ModelCacheDir moves with the checkpoint, so an upgraded binary finds nothing at the new path, on a machine whose owner has run `init` once already and reasonably believes the model is installed. The reindex that same upgrade triggers is automatic, and a manual step in the middle of an otherwise invisible migration is the part a user would experience as breakage.

Get does not call this. A library that fetches a hundred-odd MB because something imported it is a surprise no caller asked for; the choice to spend the bandwidth belongs to the command, which is also the thing with somewhere to report progress.

$SEMANTIC_NO_DOWNLOAD turns it back into a check, for a sandbox or an air-gapped machine where an unasked-for hundred-MB fetch is worse than the error it avoids. The test suite sets it so no script can spend the bandwidth by accident.

func Installed added in v0.1.3

func Installed(m *Model) bool

Installed reports whether m's files are already on disk, without disturbing the current selection. `semantic models` uses it to say which checkpoints a switch would cost a download, and which are already paid for.

func ModelCacheDir

func ModelCacheDir() string

ModelCacheDir returns the directory where model files are cached. $SEMANTIC_MODEL_DIR overrides.

The selected checkpoint names the path, for the reason OrtCacheDir keys on OrtVersion: model.onnx at a shared path would make a checkpoint change find the previous weights already there and skip the download. That failure is worse than the runtime's, because it is silent — the old weights load, pooling still runs, and the index fills with vectors from a model the caller did not choose. Keying by checkpoint also lets several models sit side by side, which is what makes switching back and forth cost nothing after the first fetch of each.

func ModelNames added in v0.1.3

func ModelNames() []string

ModelNames returns the known names in a form suitable for an error message.

func OrtCacheDir

func OrtCacheDir() string

OrtCacheDir returns the directory where the ONNX Runtime library is cached.

The version is part of the path. The binding is compiled against one release of the C API and will not load a library from another, so a single unversioned path would make an upgrade find the old file already present, skip the download, and fail at the first embed with "Error setting ORT API base". Keying by version makes the upgrade fetch what it needs and leaves the superseded library sitting harmlessly beside it.

func OrtDownloadURL

func OrtDownloadURL() string

OrtDownloadURL returns the GitHub release URL for the ORT shared library archive for the current platform.

func OrtLibFilename

func OrtLibFilename() string

OrtLibFilename returns the platform-appropriate shared library filename.

func RepresentationID

func RepresentationID() string

RepresentationID names the vector space the current model produces. It is the package-level spelling of Model.RepresentationID, kept because the index asks the package what it is embedding with, not which Model object is selected.

func Select added in v0.1.3

func Select(name string) error

Select switches the checkpoint by name and tears down any live ONNX session so the next embed builds one against the new weights. An empty name resolves $SEMANTIC_MODEL, then DefaultModel, which lets the CLI pass its flag value through unconditionally.

Callers that switch mid-process — the benchmark is the only one — get a clean session per model. Callers that never call it get DefaultModel.

func SetProgress added in v0.1.3

func SetProgress(fn Progress)

SetProgress installs the hook downloads report through. Pass nil to silence it. Not safe against a download already running; the CLI sets it once at startup.

Types

type Model added in v0.1.3

type Model struct {
	// Name is the checkpoint's identity, as the user types it and as the
	// index records it. It also names the cache directory, so two models
	// coexist on disk rather than overwriting each other's weights.
	Name string

	// Dim is the width of the output vector.
	Dim int

	// MaxSeqLen caps tokens per embed. Raising it changes the vector for every
	// chunk long enough to have been truncated at the old cap, so it is pinned
	// per checkpoint rather than shared.
	MaxSeqLen int

	// Pooling is how the token states collapse to one vector.
	Pooling Pooling

	// QueryPrefix is prepended by GetQuery and by nothing else. Asymmetric
	// models are trained to see a marker on the query side so they can tell a
	// short question from the long passage answering it; a symmetric model
	// leaves this empty and GetQuery becomes Get.
	QueryPrefix string

	// DocPrefix is prepended by Get, so it marks everything that goes into the
	// index. Most asymmetric checkpoints mark only the query and leave this
	// empty; the E5 family marks both sides and ranks badly if either marker
	// is missing. Unlike QueryPrefix this rewrites the stored vectors, so
	// RepresentationID names it.
	DocPrefix string

	// ModelURL and TokenizerURL are where `semantic init` fetches the files.
	ModelURL     string
	TokenizerURL string

	// ApproxMB is the model download's rough size, for the progress line. It
	// is what the user is about to spend, so it is worth saying before it is
	// spent.
	ApproxMB int
}

Model is one embedding checkpoint and everything that makes its vectors what they are. Two indexes built with different Models hold cosine-incomparable vectors even at the same dimension, which is why every field here feeds RepresentationID.

func Current added in v0.1.3

func Current() *Model

Current returns the checkpoint in force, resolving $SEMANTIC_MODEL on first use and falling back to DefaultModel.

An unrecognized $SEMANTIC_MODEL is ignored here rather than reported, because Current has nowhere to report it and silently embedding with a model nobody asked for is the worse of the two outcomes only if it is also the quiet one. The CLI calls Select up front, which does validate and does fail loudly, so the ignore path is reachable only by a library caller who set the variable without going through it.

func Lookup added in v0.1.3

func Lookup(name string) (*Model, bool)

Lookup finds a checkpoint by name, case-insensitively — the registry key is lowercased because `all-MiniLM-L6-v2` is not a name anyone types the same way twice.

func Models added in v0.1.3

func Models() []*Model

Models returns every known checkpoint, ordered by name, for `semantic models`.

func (*Model) RepresentationID added in v0.1.3

func (m *Model) RepresentationID() string

RepresentationID names the vector space this model produces. The index stores it and rebuilds itself when it stops matching, so anything that changes what Get returns for the same input must be reflected here.

Every component is load-bearing:

  • the checkpoint, because different weights mean different vectors;
  • the pooling and normalization, because mean-vs-CLS pooling or dropping the L2 norm rewrites the space without changing its dimension;
  • the dimension, which is the one mismatch that would fail loudly anyway;
  • the sequence cap, because raising it changes the vector for every chunk long enough to have been truncated at the old cap — silently, and only for the long chunks, which is the worst kind of drift to debug.

It is derived from the fields it names rather than written out, so it cannot drift from them. QueryPrefix is deliberately absent: it touches the query alone and leaves every stored vector unchanged, so an index built before a prefix existed stays valid.

DocPrefix is the opposite case and does appear, because it is embedded into every stored vector. It appears only when set, so the checkpoints that predate the field keep the IDs they already stamped into existing indexes.

type Pooling added in v0.1.3

type Pooling string

Pooling names how a model reduces its per-token hidden states to one vector. It is a property of the checkpoint, not a preference: a model is trained against one of these, and using the other yields a coherent-looking vector that ranks badly, with no error to notice.

const (
	// PoolCLS takes the first token. Models trained with a [CLS] objective
	// (the BGE family) put the sentence representation there.
	PoolCLS Pooling = "cls"
	// PoolMean averages the unmasked tokens. What the sentence-transformers
	// MiniLM checkpoints are trained with.
	PoolMean Pooling = "mean"
)

type Progress added in v0.1.3

type Progress func(name string, done, total int64)

Progress reports how far a single file's download has got. total is the size the server advertised, or 0 when it advertised none — a caller showing a percentage has to handle that, because Hugging Face redirects to a CDN that does not always send Content-Length.

It is a package-level hook rather than a parameter threaded through five functions because it is presentation, and every one of those functions otherwise has nothing to say about how bytes are displayed. nil means the caller wants none, which is the default and what every test gets.

type Vec

type Vec []float32

Vec is a float32 embedding vector.

func Get

func Get(text string) (Vec, error)

Get returns a normalized embedding vector for text using the local ONNX model. Returns an error if the model is not installed — run the binary's `init` command to download it (~160 MB).

v1 has no warm daemon: every process pays the ONNX session cold-start (~800ms) on its first Get. The session then stays warm for the life of the process, so batch indexing amortizes the cost.

func GetQuery added in v0.1.3

func GetQuery(text string) (Vec, error)

GetQuery embeds text as a search query. Use it for what the user typed; use Get for the content being searched.

An asymmetric checkpoint is trained to see a marker in front of a query and nothing in front of the passages it ranks, because a short question and the long passage answering it are not the same kind of text. Model.QueryPrefix carries whichever marker the selected checkpoint wants, and is empty for a symmetric one, which makes this Get.

The prefix touches the query alone, so stored vectors are unaffected and RepresentationID does not name it — an index built before this existed stays valid, and no reindex is needed to benefit.

Symmetric comparisons — `dupes`, where both sides are passages from the corpus — deliberately do not use this. There is no query in that pairing, and prefixing one arbitrary side would tilt it. They call Get, which is also what gives both sides the document marker a checkpoint like E5 expects.

Jump to

Keyboard shortcuts

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