Documentation
¶
Overview ¶
Package registry holds the model configuration catalog. Every supported embedding model is registered here with its identity (name, dimension, max input tokens), file layout, ONNX graph signature, pooling strategy, and download source.
Backend-agnostic: both the ONNX adapter and future backends (Ollama, CoreML) read model metadata from this package. Adding a new model is a single registry entry — no code changes in backend packages.
Index ¶
Constants ¶
const DefaultModelName = "bge-m3"
DefaultModelName is used when no model is explicitly specified.
Variables ¶
This section is empty.
Functions ¶
func KnownDims ¶
KnownDims returns a model's standard MRL truncation dimensions (ascending, native last), or nil when the model has no MRL ladder (only its native dim is valid). Unknown models return nil.
func PositionIDsExtraInput ¶
PositionIDsExtraInput returns position_ids [0..seqLen) broadcast over batch. Used by decoder-only ONNX exports (Qwen2, Mistral).
func QueryInstruct ¶
QueryInstruct returns a model's query-side task instruction, or "" when the model embeds queries and passages symmetrically (no query prefix). Unknown models return "".
func ZeroExtraInput ¶
ZeroExtraInput returns a [batch*seqLen] slice of zeros. Used for BERT-family token_type_ids (single-segment input).
Types ¶
type ExtraInputFn ¶
ExtraInputFn builds a synthetic input tensor that the tokenizer doesn't produce but the ONNX graph requires. BERT models need token_type_ids (zeros); decoder-only models need position_ids.
Returned slice is row-major [batch * seqLen], int64.
type ModelConfig ¶
type ModelConfig struct {
Name string // stable identifier, e.g. "bge-large-en-v1.5"
Dim int // output vector dimension
MaxInput int // max input tokens (sequences longer get truncated)
Normalize string // "l2" or ""
// QueryInstruct is the task description used to build an asymmetric model's
// query-side prompt ("Instruct: {QueryInstruct}\nQuery: {q}", Qwen3). Empty
// for symmetric models (bge-*), which embed queries and passages the same
// way — those get no query prefix.
QueryInstruct string
// KnownDims is the standard ladder of MRL truncation dimensions for an
// MRL-trained model (ascending, native dim last). --embed-dim must be one of
// these, keeping indexes on a consistent, comparable set of dimensions. Nil
// for non-MRL models, where only the native Dim is valid.
KnownDims []int
// File layout (relative to model directory)
OnnxFile string // e.g. "onnx/model.onnx"
TokenizerFile string // e.g. "tokenizer.json"
// ONNX graph signature
InputOrder []string // input tensor names in order
Outputs []string // output tensor names
ExtraInputs map[string]ExtraInputFn // synthetic inputs beyond tokenizer output
// Pooling strategy for sentence vector extraction
Pooling PoolingMode
// Download source
HFRepo string // HuggingFace repository, e.g. "BAAI/bge-large-en-v1.5"
// Memory estimate in MB for the build pipeline's pre-flight check.
// Includes weights + runtime + working set + compile headroom.
// 0 disables the check for this model.
EstimatedRAMMB uint64
}
ModelConfig describes one embedding model: identity, file layout, ONNX graph signature, pooling strategy, and download source.
func Lookup ¶
func Lookup(name string) (ModelConfig, error)
Lookup returns the config for the named model, or an error listing all known models.
func (ModelConfig) DefaultModelDir ¶
func (c ModelConfig) DefaultModelDir() (string, error)
DefaultModelDir returns ~/.cache/ckv/models/<name>.
func (ModelConfig) FetchFiles ¶
func (c ModelConfig) FetchFiles() []string
FetchFiles returns the relative paths the downloader needs to fetch.
type PoolingMode ¶
type PoolingMode int
PoolingMode selects how the inference backend collapses the per-token hidden states [batch, seqLen, hidden] into a single sentence vector [batch, hidden].
const ( // PoolingCLS takes the [CLS] token (position 0) hidden state. // Used by BERT-family models (bge-large-en-v1.5). PoolingCLS PoolingMode = iota // PoolingMean averages all attended token hidden states. // Used by sentence-BERT variants and bge-m3. PoolingMean // PoolingLastToken takes the last attended token. Used by // decoder-only models (Qwen2-based, e5-mistral). PoolingLastToken )
func (PoolingMode) String ¶
func (p PoolingMode) String() string
Source Files