Documentation
¶
Overview ¶
Package bgeonnx is the production Embedder backend running ONNX models locally via CGO. The package is organized into three concerns:
- bgeonnx.go (this file): the Embedder facade — identity, lifecycle, Embed() orchestration. Model-agnostic.
- model_config.go: per-model registry of ONNX input signature, pooling strategy, file layout. Adding a new model is one entry.
- tokenizer.go / tokenizer_impl.go: text → (input_ids, attention_mask).
- session.go / session_impl.go: tokens → pooled+normalized vectors.
Default build (no `bgeonnx` tag) keeps the stub variants so users without libonnxruntime + libtokenizers can still build CKV with the mock embedder. `-tags bgeonnx` wires the real CGO implementations.
Model configuration bridge: re-exports registry types so existing bgeonnx internal code (session, pooling, tokenizer) can reference them without a package-wide import change. New code outside bgeonnx should import internal/embed/registry directly.
Index ¶
- Constants
- Variables
- func EstimatedRAMMB(opts Options) uint64
- func FetchModel(name, destDir string, progress func(string)) (string, error)
- type Adapter
- func (a *Adapter) Close() error
- func (a *Adapter) Dimension() int
- func (a *Adapter) Embed(ctx context.Context, batch []string) ([][]float32, error)
- func (a *Adapter) EstimatedRAMMB() uint64
- func (a *Adapter) Identity() types.EmbeddingIdentity
- func (a *Adapter) MaxInputTokens() int
- func (a *Adapter) ModelDir() string
- func (a *Adapter) Name() string
- func (a *Adapter) Provider() string
- type ExtraInputFn
- type ModelConfig
- type Options
- type PoolingMode
- type Session
- type TokenizedBatch
- type Tokenizer
Constants ¶
const ( PoolingCLS = registry.PoolingCLS PoolingMean = registry.PoolingMean PoolingLastToken = registry.PoolingLastToken )
const DefaultModelName = registry.DefaultModelName
Variables ¶
var ( ZeroExtraInput = registry.ZeroExtraInput PositionIDsExtraInput = registry.PositionIDsExtraInput )
var ErrNotImplemented = errors.New("bgeonnx: ONNX runtime not wired — rebuild with -tags bgeonnx (see docs/d1-installation-guide.md)")
ErrNotImplemented is returned by stub Tokenizer / Session when callers compile without `-tags bgeonnx`. Makes "you didn't enable the real backend" a clear runtime signal rather than a mysterious zero-vector output.
Functions ¶
func EstimatedRAMMB ¶
EstimatedRAMMB resolves model options and returns the registered memory estimate. Used by the build pipeline's pre-flight check.
Types ¶
type Adapter ¶
type Adapter struct {
// contains filtered or unexported fields
}
Adapter is the Embedder facade. One per process: Session construction is expensive (~1.5s cold start) so callers hold this long-lived and share it across queries.
func (*Adapter) Close ¶
Close releases the underlying Session and Tokenizer. The Tokenizer interface deliberately omits Close — only some implementations need it (the HF binding holds a Rust object) — so we test for io.Closer dynamically. Idempotent.
func (*Adapter) Embed ¶
Embed orchestrates tokenizer → session. Pooling and normalization happen inside Session per the model's ModelConfig.Pooling.
func (*Adapter) EstimatedRAMMB ¶
EstimatedRAMMB exposes ModelConfig.EstimatedRAMMB to callers via duck typing. The build pipeline's memory guard reads this through an anonymous interface, so we deliberately don't widen the public Embedder interface — embedders without an estimate (the mock) just don't implement this method and the guard treats them as unknown.
func (*Adapter) Identity ¶
func (a *Adapter) Identity() types.EmbeddingIdentity
Identity reports the embedding space this adapter produces, sourced entirely from the model registry (model_config.go) so it stays correct as models are added or swapped — no per-model value is hardcoded here.
func (*Adapter) MaxInputTokens ¶
func (*Adapter) ModelDir ¶
ModelDir returns the directory the adapter loaded model.onnx + tokenizer.json from. Empty for a nil adapter. Read by the health endpoint via duck typing — embedders that don't have a model dir (the mock) don't implement this method and the field stays empty in the health response.
func (*Adapter) Provider ¶
Provider reports the underlying session's execution backend ("cpu", "coreml", "coreml-fallback-to-cpu", "stub"). The build pipeline uses this to label footprint events so a slow run can be diagnosed against the backend that produced it. Returns empty when the adapter or its session is nil.
type ExtraInputFn ¶
type ExtraInputFn = registry.ExtraInputFn
type ModelConfig ¶
type ModelConfig = registry.ModelConfig
Type aliases for backward compatibility within the bgeonnx package.
func LookupModel ¶
func LookupModel(name string) (ModelConfig, error)
func RegisteredModels ¶
func RegisteredModels() []ModelConfig
type Options ¶
type Options struct {
ModelDir string
ModelName string // optional; inferred from ModelDir basename if blank
Tokenizer Tokenizer
Session Session
}
Options control adapter construction.
ModelDir + ModelName resolution order:
- ModelName explicit → registry lookup; ModelDir defaults to ~/.cache/ckv/models/<ModelName>.
- ModelDir explicit, ModelName empty → infer ModelName from ModelDir's basename, then registry lookup.
- Both empty → DefaultModelName + ~/.cache/ckv/models/<default>.
Tokenizer / Session override the default factories. Tests inject stubs here to bypass CGO; production callers leave them nil.
type PoolingMode ¶
type PoolingMode = registry.PoolingMode
type Session ¶
type Session interface {
// Run executes one batch through the graph. Output is
// len(batch) × ModelDim float32 vectors, already L2-normalized.
Run(ctx context.Context, tokens TokenizedBatch) ([][]float32, error)
// Provider names the execution backend the session ended up using.
// Values are short tokens suitable for log fields:
// "cpu" — default ORT CPU EP, no acceleration attached.
// "coreml" — CoreML EP V2 attached and accepted on Apple
// hardware.
// "coreml-fallback-to-cpu" — CoreML attach attempted but failed
// at session creation; ORT silently uses CPU.
// "stub" — non-bgeonnx build; no real session.
// Surfaced via Adapter.Provider() so the build footprint records
// which backend ran an index.
Provider() string
// Close releases the underlying ONNX session + environment.
// Idempotent.
Close() error
}
Session is the ONNX Runtime concern: take tokenized tensors, run the model graph, return pooled+normalized 1024-d vectors.
Production impl: `yalue/onnxruntime_go` behind the `bgeonnx` build tag so existing CI without libonnxruntime stays green.
Pooling + normalization happen here, not at the caller. bge-* models use mean pooling over the token dimension (masked by attention), then L2 normalize. Doing it inside Session keeps the interface uniform across embedders.
type TokenizedBatch ¶
type TokenizedBatch struct {
InputIDs [][]int64
AttentionMask [][]int64
TokenTypeIDs [][]int64 // reserved; populated only if a model registers it via Tokenize
}
TokenizedBatch is the on-the-wire shape for one Embed() call. All slices are sized [batchSize][seqLen] where seqLen is uniform across the batch (padded with 0 attention). TokenTypeIDs is left nil — extra inputs the ONNX graph needs (token_type_ids for BERT, position_ids for Qwen2, etc.) are synthesized in Session.Run from ModelConfig.ExtraInputs.
type Tokenizer ¶
type Tokenizer interface {
// Tokenize returns input_ids + attention_mask for batch. maxLen
// is the hard upper bound (model-specific, e.g. 512 for
// bge-large-en-v1.5) — tokens beyond it MUST be truncated, never
// silently retained.
Tokenize(ctx context.Context, batch []string, maxLen int) (TokenizedBatch, error)
}
Tokenizer turns a batch of strings into the token tensors ONNX Runtime expects. The shape `[][]int64` is row-major: outer index is batch position, inner index is token position. Truncation / padding-to-max is the tokenizer's responsibility, not the session's.
Production impl: `daulet/tokenizers` wrapping the HuggingFace Rust tokenizers crate; reads the model's tokenizer.json directly.
Source Files