Documentation
¶
Overview ¶
Package encoder loads and (in subsequent commits) runs the nomic-ai/CodeRankEmbed neural reranker as a pure-Go forward pass.
This file implements Milestone 1 of outputs/ken-rerank-plan.md: config + weight loader with strict shape validation against the dumped checkpoint schema. Forward pass arrives in M2.
The package is a sibling of internal/embed (Model2Vec, the first-stage retriever); the two share tokenize.go + safetensors.go but use different inference algorithms and load different artifacts, so they stay separate per plan §3.
Example ¶
Rerank candidates with the CodeRankEmbed transformer: encode the query (with the mandatory query prefix) and each candidate, then score by cosine of the returned hidden states. Higher-fidelity than the Model2Vec first stage — the typical two-stage retrieve-then-rerank setup.
Needs a checkpoint on disk, so this example is illustrative (compiled, not run). See the repo README for fetching CodeRankEmbed.
package main
import (
"fmt"
"github.com/townsendmerino/aikit/encoder"
)
func main() {
m, err := encoder.Load("path/to/coderankembed-snapshot")
if err != nil {
return
}
query, _ := m.Encode("how do I parse json", true) // isQuery = true
doc, _ := m.Encode("func parseJSON(b []byte) (*Config, error) { ... }", false)
_, _ = query, doc
fmt.Println(m.HiddenDim())
}
Output:
Index ¶
- Constants
- func EncodeDoc(tok *embed.Tokenizer, text string, maxLen int) ([]int32, error)
- func EncodeQuery(tok *embed.Tokenizer, query string, maxLen int) ([]int32, error)
- func RegisterBackend(name string, factory func() (Backend, error))
- type BERT
- type Backend
- type Config
- type CrossEncoder
- type Encoder
- type LayerWeights
- type LayerWeightsQ8
- type Model
- func (m *Model) Encode(text string, isQuery bool) ([]float32, error)
- func (m *Model) EncodeBatch(texts []string, isQueries []bool, concurrency int) ([][]float32, error)
- func (m *Model) EncodeTokens(text string, isQuery bool) ([]float32, int, error)
- func (m *Model) EncodeTokensWithIDs(text string, isQuery bool) ([]float32, []int32, error)
- func (m *Model) HiddenDim() int
- func (m *Model) SetMaxSeqLength(n int)
- type ModelQ8
- type SPLADE
- type Weights
- type WeightsQ8
Examples ¶
Constants ¶
const DefaultMaxSeqLength = 512
DefaultMaxSeqLength caps query+candidate token length per plan §5. CodeRankEmbed's tokenizer_config.max_length is 512; the model itself supports up to n_positions=8192 via RoPE, but rerank candidates are chunk-sized and 512 keeps latency bounded. Truncation is right-side (tokenizer_config truncation_side=right), preserving the [CLS] prefix.
const QueryPrefix = "Represent this query for searching relevant code: "
QueryPrefix is the mandatory CodeRankEmbed query-side instruction from the model card. Document (code) inputs MUST NOT carry this prefix; only queries do. Misusing it on docs degrades cosine sharply.
Variables ¶
This section is empty.
Functions ¶
func EncodeDoc ¶
EncodeDoc tokenizes a code/document candidate for CodeRankEmbed: no prefix, just [CLS]/text/[SEP] with right truncation to maxLen. maxLen ≤ 0 means "use DefaultMaxSeqLength".
func EncodeQuery ¶
EncodeQuery tokenizes a query for CodeRankEmbed: prepend the mandatory query prefix, then wrap with [CLS]/[SEP], truncating to maxLen tokens from the right. maxLen ≤ 0 means "use DefaultMaxSeqLength".
func RegisterBackend ¶ added in v0.4.0
RegisterBackend registers a named Backend factory. The goinfer/gpu module calls this from init() (under `-tags gpu`) to make "webgpu" available without aikit importing the cgo WebGPU implementation. Safe for concurrent use; a later registration of the same name replaces the earlier one.
Types ¶
type BERT ¶ added in v1.4.0
type BERT struct {
// contains filtered or unexported fields
}
BERT is a loaded MiniLM-class encoder. Immutable after load; the forward is read-only-safe for concurrent use (no shared mutable state).
func LoadBERT ¶ added in v1.4.0
LoadBERT loads a sentence-transformers BERT model (config.json + model.safetensors with BERT tensor names) from dir. It validates the two architecture assumptions this forward implements: GELU activation and learned absolute positions.
func (*BERT) Embed ¶ added in v1.4.0
Embed returns the mean-pooled, L2-normalized sentence embedding for token ids.
type Backend ¶ added in v0.4.0
type Backend interface {
// Name identifies the backend ("cpu", "webgpu").
Name() string
// MatmulBT computes dst[M,N] = a[M,K] · b[N,K]ᵀ — the PyTorch [out,in]
// weight layout the safetensors checkpoints already store (so no
// transpose copy), matching the encoder's matmulBT convention.
MatmulBT(a, b, dst []float32, M, K, N int)
// Close releases backend resources (GPU buffers, etc.). No-op on CPU.
Close() error
}
Backend abstracts the one hot primitive the encoder forward pass is bound by — the big weight matmuls (Wqkv, OutProj, the MLP projections). Norms, RoPE, softmax and elementwise ops are cheap and stay on the CPU even when a GPU matmul backend is in use, which avoids a host↔device round-trip per layer.
The default "cpu" backend is always registered (pure-Go SIMD matmul, no cgo). A WebGPU backend lives in the opt-in github.com/townsendmerino/goinfer/gpu module: importing it under `-tags gpu` calls RegisterBackend("webgpu", …) on init, so the encoder gains GPU acceleration WITHOUT ever importing github.com/cogentcore/webgpu (cgo) into aikit's dependency graph.
func NewBackend ¶ added in v0.4.0
NewBackend returns the named backend. "" and "cpu" always resolve to the pure-Go CPU backend. Other names resolve through the registry; "webgpu" falls back to CPU with an explanatory error (rather than hard-failing) when goinfer/gpu has not been imported, so a `--backend webgpu` flag still runs on a build without the GPU module.
type Config ¶
type Config struct {
VocabSize int `json:"vocab_size"`
HiddenDim int `json:"n_embd"`
NumLayers int `json:"n_layer"`
NumHeads int `json:"n_head"`
IntermediateDim int `json:"n_inner"`
MaxPositions int `json:"n_positions"`
MaxTrainedPositions int `json:"max_trained_positions"`
TypeVocabSize int `json:"type_vocab_size"`
RoPEBase float64 `json:"rotary_emb_base"`
RoPEFraction float64 `json:"rotary_emb_fraction"`
RoPEInterleaved bool `json:"rotary_emb_interleaved"`
LayerNormEpsilon float64 `json:"layer_norm_epsilon"`
ActivationFunction string `json:"activation_function"`
Prenorm bool `json:"prenorm"`
UseRMSNorm bool `json:"use_rms_norm"`
QKVProjBias bool `json:"qkv_proj_bias"`
MLPFc1Bias bool `json:"mlp_fc1_bias"`
MLPFc2Bias bool `json:"mlp_fc2_bias"`
ScaleAttnWeights bool `json:"scale_attn_weights"`
Causal bool `json:"causal"`
ParallelBlock bool `json:"parallel_block"`
// contains filtered or unexported fields
}
Config captures the architecture constants from CodeRankEmbed's config.json that the forward pass depends on. Loaded from the checkpoint rather than hardcoded so a drop-in compatible checkpoint can override; ValidateAssumptions then fails loudly on any dimension/feature this loader/forward-pass doesn't implement.
func (*Config) ValidateAssumptions ¶
ValidateAssumptions errors if the config contradicts any baked-in assumption of the forward pass: post-norm only, standard LayerNorm, SwiGLU MLP, no biases on QKV/MLP, bidirectional, sequential block, full-rotation NeoX-style RoPE. Fail loudly at load time rather than silently produce junk activations. The plan §1 calls every one of these out — this is the runtime guard that pins them.
type CrossEncoder ¶ added in v1.5.0
type CrossEncoder struct {
// contains filtered or unexported fields
}
CrossEncoder is a loaded BERT cross-encoder reranker.
func LoadCrossEncoder ¶ added in v1.5.0
func LoadCrossEncoder(dir string) (*CrossEncoder, error)
LoadCrossEncoder loads a BertForSequenceClassification cross-encoder (config.json + model.safetensors) from dir: the BERT trunk (via LoadBERT) plus the pooler and the classification head. The number of labels is read from the classifier shape (1 for a ms-marco-style relevance reranker).
func (*CrossEncoder) Score ¶ added in v1.5.0
func (ce *CrossEncoder) Score(query, doc string) (float32, error)
Score returns the relevance logit for a (query, document) pair — higher is more relevant. Rank a candidate list by descending Score to rerank. (For a model with more than one label, this is label 0; use ScoreAll for the rest.)
type Encoder ¶
type Encoder interface {
Encode(text string, isQuery bool) ([]float32, error)
EncodeBatch(texts []string, isQueries []bool, concurrency int) ([][]float32, error)
HiddenDim() int
}
Encoder is the surface used by NeuralReranker. Both *Model (f32) and *ModelQ8 (int8) implement it; the reranker doesn't care which it holds. Lets the CLI / MCP layer pick a precision at startup without the rest of the code knowing.
type LayerWeights ¶
type LayerWeights struct {
Wqkv []float32 // [3*HiddenDim, HiddenDim] fused Q/K/V input projection, no bias
OutProj []float32 // [HiddenDim, HiddenDim] attention output projection, NO bias (verified against checkpoint)
Norm1W []float32 // [HiddenDim] post-attention LayerNorm weight
Norm1B []float32 // [HiddenDim] post-attention LayerNorm bias
Fc11 []float32 // [IntermediateDim, HiddenDim] SwiGLU gate (not fused with Fc12 in the checkpoint)
Fc12 []float32 // [IntermediateDim, HiddenDim] SwiGLU value (not fused with Fc11 in the checkpoint)
Fc2 []float32 // [HiddenDim, IntermediateDim] output projection, no bias
Norm2W []float32 // [HiddenDim] post-MLP LayerNorm weight
Norm2B []float32 // [HiddenDim] post-MLP LayerNorm bias
}
LayerWeights bundles one transformer block's tensors. Matrices are stored in PyTorch's [out, in] row-major layout — matmul is then A · Bᵀ without a transpose copy, matching internal/embed's convention.
Lifetime: every []float32 here aliases the underlying SafetensorsFile bytes (zero-copy unsafe slice). Do not mutate; do not let the parent Weights drop while these are in use.
type LayerWeightsQ8 ¶
type LayerWeightsQ8 struct {
Wqkv linalg.WeightMat // [3*HiddenDim, HiddenDim]
OutProj linalg.WeightMat // [HiddenDim, HiddenDim]
Fc11 linalg.WeightMat // [IntermediateDim, HiddenDim]
Fc12 linalg.WeightMat // [IntermediateDim, HiddenDim]
Fc2 linalg.WeightMat // [HiddenDim, IntermediateDim]
// LN weights stay f32 — small (768 each) so no memory saving, and
// LN is the parity-sensitive op the f32 forward already uses f64
// accumulators on; quantizing γ/β here introduces compounding noise
// across the 12 layers for no benefit.
Norm1W, Norm1B []float32
Norm2W, Norm2B []float32
}
LayerWeightsQ8 is the M8 int8-quantized per-layer bundle. Only the big linear projections are quantized — the LayerNorm parameters (Norm1W/B, Norm2W/B) are tiny ([D]=768 each) and stay f32 (no memory saving, plus LN is parity-sensitive: float-noise from quantizing γ/β here would compound across the 12 layers).
Each big matrix is a weight-only Q8 linalg.WeightMat (per-row int8 + per-row f32 scales; reconstruction f32[i,j] ≈ int8[i*K+j] * scale[i]). The forward still drives the encoder's own baked-scale matmulBTQ8Into over the int8 codes/scales pulled via WeightMat.Int8() — storage is shared, the kernel stays the encoder's (it is numerically distinct from linalg.MatmulBTQ8: large-M dequant-once-into-scratch).
type Model ¶
type Model struct {
// contains filtered or unexported fields
}
Model is the loaded CodeRankEmbed reranker: weights + tokenizer. Goroutine-safe for concurrent Encode calls (all internal state is immutable after Load; per-call buffers are stack/heap-local).
func Load ¶
Load reads a CodeRankEmbed snapshot from dir (config.json, model.safetensors, tokenizer.json — the standard HF layout). Cf. embed.LoadFromFS for the analogous Model2Vec loader.
func LoadFromFS ¶
LoadFromFS reads from fsys rooted at dir. Same file shape as Load.
func (*Model) Encode ¶
Encode tokenizes `text` (prepending the mandatory query prefix iff isQuery is true), wraps with [CLS]/[SEP], runs the transformer forward pass, and returns the raw (UN-normalized) CLS hidden state.
The caller is responsible for L2-normalizing if the consumer is cosine; ranking is invariant to it but the parity goldens compare raw vectors so this is the natural Model output.
Goroutine-safe. The Weights and Tokenizer are immutable; every per-call buffer (token ids, hidden states, attention scores, …) is allocated fresh inside the call.
func (*Model) EncodeBatch ¶
EncodeBatch runs N forward passes in parallel across concurrency workers, returning one CLS vector per (text, isQuery) input. Two layers of parallelism (M3 + M7):
- WORKERS: input is statically split into `concurrency` chunks (default NumCPU). Workers run independently — no shared state other than input/output slices.
- BATCHED FORWARD: each worker calls forwardBatch on its chunk, processing all its candidates as one padded batch through the 12 layers' big matmuls (Wqkv, OutProj, fc11, fc12, fc2). The batched matmuls amortize per-call overhead and keep hidden states in cache across layers — the M7 win.
Static partitioning (vs M3's job-channel) is fine here because each worker now does a coalesced batch, not per-input pop-from-queue work. Load imbalance can show on adversarial inputs (worker A's chunk is all 500-token; worker B's all 5-token), but rerank batches are typically uniform-length per M3's measurements.
On error, returns the first error and a nil result slice (no partial results). `concurrency` ≤ 0 means runtime.NumCPU().
func (*Model) EncodeTokens ¶ added in v0.2.0
EncodeTokens tokenizes `text` (query or doc per the prefix rule) and returns the per-token hidden-state matrix from the forward pass. Returns (vectors []float32 of length L*D, L int, error). Row i is vectors[i*D : (i+1)*D]; D is m.HiddenDim().
Token-id boundaries: with isQuery=true the token sequence is "[CLS] <prefix tokens…> <query tokens…> [SEP]". Callers that want to exclude the instruction prefix can find its boundary via EncodeQueryPrefixIDs (TODO if needed) or by re-encoding just the prefix and counting; for v0 the MaxSim probe in ken handles the exclusion by looking up well-known [CLS]/[SEP] ids and the prefix token sequence directly.
Used by ken's MaxSim rerank probe; mirror of Encode but returning every position instead of CLS-pooling.
func (*Model) EncodeTokensWithIDs ¶ added in v0.2.0
EncodeTokensWithIDs is EncodeTokens plus the underlying token-id sequence — useful for the MaxSim probe to mask out the query-prefix range and special tokens without re-tokenizing.
func (*Model) SetMaxSeqLength ¶
SetMaxSeqLength overrides the per-call truncation cap. Useful for unit tests (small L = much faster) and for callers who want to trade recall on long inputs for latency. 0 or negative resets to default.
type ModelQ8 ¶
type ModelQ8 struct {
// contains filtered or unexported fields
}
ModelQ8 is the int8-quantized sibling of Model. Same API surface (the Encoder interface is the common contract) but the per-layer big linear projections store int8 + per-row scales instead of float32 weights, cutting weight bytes ~4× (137M params × 4B = 547MB → ~140 MB resident). Forward pass routes through matmulBTQ8 for those layers.
Accuracy cost: end-to-end cosine ≥ 0.97 vs the f32 Model (pinned by TestModelQ8_cosineMatchesF32). For NDCG: M0's measured CoIR lift was +0.165 at β=1; the int8 model is expected to reproduce that to within bench noise (~±0.01) because per-matmul ~0.8% relative error attenuates by the time it's squashed through 12 LayerNorms.
func LoadQ8 ¶
LoadQ8 reads + quantizes the rerank model at dir. Same disk layout as Load (config.json + tokenizer.json + model.safetensors).
func (*ModelQ8) EncodeBatch ¶
func (m *ModelQ8) EncodeBatch(texts []string, isQueries []bool, concurrency int) ([][]float32, error)
EncodeBatch implements Encoder. Same static-partition + batched- forward-per-worker shape as Model.EncodeBatch; routes through the q8 batched forward instead.
func (*ModelQ8) SetMaxSeqLength ¶
SetMaxSeqLength mirrors Model.SetMaxSeqLength.
type SPLADE ¶ added in v1.4.0
type SPLADE struct {
// contains filtered or unexported fields
}
SPLADE is a loaded SPLADE model: a BERT encoder plus the BertForMaskedLM head.
func LoadSPLADE ¶ added in v1.4.0
LoadSPLADE loads a BertForMaskedLM SPLADE model (config.json + model.safetensors) from dir: the BERT encoder (via LoadBERT) plus the masked-LM head.
type Weights ¶
type Weights struct {
Cfg Config
WordEmb []float32 // [VocabSize, HiddenDim]
TokenTypeEmb []float32 // [TypeVocabSize, HiddenDim] — only row 0 used (single segment)
EmbLN_W []float32 // [HiddenDim]
EmbLN_B []float32 // [HiddenDim]
Layers []LayerWeights
// contains filtered or unexported fields
}
Weights is the immutable per-checkpoint bundle returned by Load*. Multiple concurrent forward passes share one Weights instance.
Plan §1 correction: the checkpoint stores the SwiGLU gate and value as TWO separate tensors (mlp.fc11 [3072,768] and mlp.fc12 [3072,768]), not a fused mlp.fc1 [6144,768]. There is also no out_proj bias and no final encoder LayerNorm beyond the last block's norm2.
func LoadWeights ¶
LoadWeights reads config.json + model.safetensors from a real on-disk directory. As of M8 the .safetensors blob is mmapped (not heap-copied) so the 547 MB CodeRankEmbed checkpoint stays in the OS page cache instead of dominating Go heap RSS. config.json (small) still goes through fs.ReadFile.
Use LoadWeightsFromFS for fs.FS-backed (MapFS, embed.FS) paths — that route stays heap-backed because fs.FS doesn't expose a file descriptor.
func LoadWeightsFromFS ¶
LoadWeightsFromFS reads config.json + model.safetensors from fsys/dir, validates every tensor's shape against Cfg, and returns the bundle. Returns the first error encountered (tensor name included) so the failure mode is one clear "tensor X has shape Y, want Z" rather than silent activation drift.
fs.FS-backed (heap copy via fs.ReadFile). For mmap-backed loads from a real directory, use LoadWeights (M8 path).
type WeightsQ8 ¶
type WeightsQ8 struct {
Cfg Config
WordEmb []float32 // [VocabSize, HiddenDim] — stays f32 (embedding lookup, not matmul)
TokenTypeEmb []float32 // [TypeVocabSize, HiddenDim] — stays f32
EmbLN_W []float32 // [HiddenDim] — LN bias, parity-sensitive
EmbLN_B []float32
Layers []LayerWeightsQ8
// contains filtered or unexported fields
}
WeightsQ8 mirrors Weights but with the 5 big linear projections per layer quantized to int8 + per-row scales. Embeddings (WordEmb) and emb_ln stay f32 — embedding rows feed an integer lookup, not a matmul, so quantizing them has no perf benefit; emb_ln is the same "tiny + parity-sensitive" exception as the per-layer LNs.
Memory footprint: ~140 MB total vs Weights' ~547 MB (4× reduction for the linear-projection bulk; embeddings still cost ~92 MB).
func LoadWeightsQ8 ¶
LoadWeightsQ8 reads the f32 checkpoint via the mmap path (M8) and quantizes the 5 big per-layer linear matrices to int8 at load time. Calibration is per-row symmetric (max(|row|) / 127).
After this returns, the original f32 weight bytes are released via Close() on the safetensors handle — we hold int8 copies on the Go heap instead. The model footprint drops from ~547 MB (heap) + 547 MB (mmap shadow) to ~140 MB int8 + the small f32 tail.
Source Files