Documentation
¶
Overview ¶
Package native is a pure Go transformer inference runtime for hush.
It loads a distilroberta model from the .hbin format produced by training/scripts/export_hbin.py and runs a classifier forward pass without any CGO dependency. This exists so hush can ship as a single static binary with no libonnxruntime requirement at runtime.
Index ¶
- Constants
- type Arena
- type Bundle
- type Detector
- type Labels
- type Layer
- type MaybeWeight
- type Meta
- type Model
- type QuantWeight
- type RawTensor
- type Scorer
- type Span
- type Tensor
- func Add(a, b *Tensor) *Tensor
- func AddBias(x *Tensor, bias []float32) *Tensor
- func AddInPlace(a, b *Tensor) *Tensor
- func ApplyAdditiveMask(x *Tensor, mask []float32) *Tensor
- func BatchMatMul(a, b *Tensor) *Tensor
- func FromSlice(shape []int, data []float32) *Tensor
- func GELU(x *Tensor) *Tensor
- func Gather(table *Tensor, indices []int32, B, T int) *Tensor
- func LayerNorm(x *Tensor, gamma, beta []float32, eps float32) *Tensor
- func MatMul(a, b *Tensor) *Tensor
- func MatMulInt8(a *Tensor, w *QuantWeight) *Tensor
- func NewTensor(shape ...int) *Tensor
- func Reshape(t *Tensor, shape ...int) *Tensor
- func ScaleInPlace(t *Tensor, s float32) *Tensor
- func Softmax(x *Tensor) *Tensor
- func Transpose(t *Tensor, axes []int) *Tensor
Constants ¶
const ( DTypeF32 uint8 = 1 DTypeI8 uint8 = 2 DTypeI32 uint8 = 3 )
const ModelVersion = "hush-model-v1"
ModelVersion is the identifier of the embedded v1 sequence-classification model. Keep in sync with the asset filename and CHANGELOG.
const ModelVersionV2 = "hush-model-v2"
ModelVersionV2 identifies the embedded v2 token-classification (NER) model.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type Arena ¶ added in v0.1.4
type Arena struct {
// contains filtered or unexported fields
}
Arena is a per-forward scratch allocator for *Tensor buffers. The forward pass creates hundreds of short-lived tensors (one per op); allocating each via make/NewTensor turned into ~500 allocs and ~4 MB per call, all of which then had to be GC'd.
The arena hands out []float32 buffers from a pool of reusable slices, growing the pool only when a request exceeds all existing buffers. Reset() marks every buffer free so the next forward reuses them in place. Allocation is O(B) per Get where B is the number of live buffers; B stays tiny (tens).
This is explicitly not thread-safe. Each forward pass owns its arena; parallel inference creates multiple arenas.
type Bundle ¶
Bundle is the result of reading an hbin file: topology + all tensors keyed by their ONNX initializer names.
type Detector ¶ added in v0.1.8
type Detector struct {
// contains filtered or unexported fields
}
Detector wraps a token-classification (NER) Model and tokenizer and produces character-offset spans over arbitrary input text. It owns the sliding-window strategy for inputs that exceed the model's effective context.
Concurrent use: as of v0.1.11, Detect serializes calls through an internal mutex so library callers can wire one Detector into a worker pool without crashing on tensor-buffer races. Throughput is therefore bounded by one forward pass at a time; callers that need parallel throughput should construct N Detectors and round-robin across them.
func LoadDetector ¶ added in v0.1.8
LoadDetector reads a v2 token-classification Model from modelPath (.hbin) and a tokenizer from tokenizerPath (HF tokenizer.json).
func LoadDetectorReader ¶ added in v0.1.8
LoadDetectorReader is the io.Reader variant for embedded bytes.
func NewBundledDetector ¶ added in v0.1.8
NewBundledDetector constructs a Detector from the int8 v2 NER model and tokenizer embedded via go:embed. No filesystem access is required.
func NewDetector ¶ added in v0.1.8
NewDetector returns a Detector using an already-loaded token-classification Model and tokenizer. Returns an error if the model isn't a token classification model or its label metadata is missing.
func (*Detector) Close ¶ added in v0.1.8
Close is a no-op (no CGO resources). Present for API parity with Scorer.
func (*Detector) Detect ¶ added in v0.1.8
Detect tokenizes text in overlapping char windows, runs Forward on each, BIO-decodes the per-token logits, shifts spans back to absolute char offsets, and dedupes overlaps across windows. Returns spans sorted by Start ascending.
func (*Detector) ModelVersion ¶ added in v0.1.8
ModelVersion returns the embedded model identifier.
type Labels ¶ added in v0.1.8
type Labels struct {
Id2Label map[string]string `json:"id2label"`
Label2Id map[string]int `json:"label2id"`
NumLabels int `json:"num_labels"`
SeqLen int `json:"seq_len"`
}
Labels carries token classification label metadata when Task == "token_classification". It is omitted for sequence classification models.
type Layer ¶
type Layer struct {
QueryW, KeyW, ValueW *MaybeWeight
QueryB, KeyB, ValueB []float32
AttnOutW *MaybeWeight
AttnOutB []float32
Attn1LN_W, Attn1LN_B []float32 // attention output LayerNorm
InterW *MaybeWeight // intermediate.dense [H, FFN]
InterB []float32
OutputW *MaybeWeight // output.dense [FFN, H]
OutputB []float32
Out2LN_W, Out2LN_B []float32 // output LayerNorm
}
Layer holds all weights for one transformer block.
type MaybeWeight ¶
type MaybeWeight struct {
F32 *Tensor // non-nil for fp32 weights (possibly eager-dequantized)
I8 *QuantWeight // non-nil for int8 weights kept as int8 in RAM
}
MaybeWeight is a tagged union holding either a dense fp32 weight or a per-output-channel int8 quantized weight. Exactly one of F32 or I8 is non-nil. Helper MatMul dispatches to the right kernel.
func (*MaybeWeight) IsInt8 ¶
func (w *MaybeWeight) IsInt8() bool
IsInt8 reports whether this is an int8 quantized weight (still int8 in RAM).
func (*MaybeWeight) MatMul ¶
func (w *MaybeWeight) MatMul(a *Tensor) *Tensor
MatMul dispatches to fp32 MatMul or MatMulInt8 depending on storage. For fp32 weights, the caller must have already oriented the weight to [K, N] (matching MatMul semantics). For int8 weights, the stored layout is always [In, Out] == [K, N].
type Meta ¶
type Meta struct {
Model string `json:"model"`
Hidden int `json:"hidden"`
Layers int `json:"layers"`
Heads int `json:"heads"`
FFN int `json:"ffn"`
Vocab int `json:"vocab"`
MaxPosition int `json:"max_position"`
TokenTypeCount int `json:"token_type_count"`
PaddingIdx int `json:"padding_idx"`
SeqLen int `json:"seq_len"`
OutputClasses int `json:"output_classes"`
Task string `json:"task,omitempty"`
Labels *Labels `json:"labels,omitempty"`
}
Meta captures the topology of the embedded model. The hbin exporter infers these from the ONNX graph and writes them as JSON in the header.
func EmbeddedV2Meta ¶ added in v0.1.8
EmbeddedV2Meta reads the embedded v2 hbin and returns just its Meta. The hbin format requires sequential reads, so this loads tensors too; the bundle is dropped immediately and Meta is a small struct, so the cost is transient. Use this from the CLI to auto-detect Task without committing to a full Detector load.
func (*Meta) IsTokenClassification ¶ added in v0.1.8
IsTokenClassification reports whether this meta describes a token classification (NER) model. Empty Task defaults to sequence classification for backward compatibility with v1 hbin files.
type Model ¶
type Model struct {
Meta Meta
WordEmb *Tensor // [V, H]
PosEmb *Tensor // [MaxPos, H]
TypeEmb *Tensor // [1, H]
EmbLN_W []float32
EmbLN_B []float32
Layers []Layer
// Classifier dense + out_proj. Stored so that MatMul(a, W) produces
// the expected output: for fp32 this is the already-transposed [in, out]
// layout; for int8 the exporter pre-transposes so the stored weight is
// also [in, out].
ClsDenseW *MaybeWeight // [H, H]
ClsDenseB []float32
ClsOutW *MaybeWeight // [H, num_classes]
ClsOutB []float32
// Token classification head (v2). Loaded only when Meta.Task == "token_classification".
ClassifierW *MaybeWeight // [H, num_labels]
ClassifierB []float32 // [num_labels]
}
Model is a loaded distilroberta classifier.
func LoadModel ¶
LoadModel constructs a Model from an hbin Bundle. Every weight must be present with an expected name and shape; missing or mis-shaped weights return an error rather than silently zero-filling.
func (*Model) Forward ¶
Forward runs a classifier forward pass for a single example. inputIDs and attentionMask are [seqLen].
Return contract depends on Meta.Task:
- v1 sequence classification (default / empty Task): returns logits of length OutputClasses ([num_classes]).
- v2 token classification (Task == "token_classification"): returns row-major logits of length T*num_labels, where T is the effective (post pad-trim) sequence length. Caller must consult Meta.Task and Meta.Labels to interpret the slice.
The runtime trims trailing padding tokens before running. Transformers with masked attention are length invariant, so this does not change the numerics — only skips wasted computation on pad positions that would be zeroed out anyway. For typical hush inputs (~60 tokens out of 384) this is a 20-40x speedup.
func (*Model) ForwardBatch ¶ added in v0.1.6
ForwardBatch runs the classifier over B examples in a single pass. inputIDs and attentionMask are both length B*seqLen, each example occupying a contiguous row. Returns a slice of length B.
Per-example return contract mirrors Forward:
- v1 sequence classification: each entry holds OutputClasses logits.
- v2 token classification: each entry holds T*num_labels row-major logits, where T is the (shared, per-batch) effective sequence length.
Matches looping Forward numerically to within fp32 reassociation drift. Dynamic T trim is per-batch (max effective length).
type QuantWeight ¶
QuantWeight holds a per-output-channel symmetrically quantized int8 weight matrix in [In, Out] layout along with fp32 scales of length [Out]. Dequantized value at position (k, j) is:
W[k, j] = int8Data[k*Out + j] * Scale[j]
Zero point is 0 (symmetric quantization). Scales are always positive.
func NewQuantWeight ¶
func NewQuantWeight(in, out int, data []int8, scale []float32) (*QuantWeight, error)
NewQuantWeight constructs a QuantWeight from raw buffers, validating shape consistency.
func (*QuantWeight) DequantizeToF32 ¶
func (qw *QuantWeight) DequantizeToF32() *Tensor
DequantizeToF32 eagerly materializes an int8 QuantWeight into a fp32 *Tensor of shape [In, Out]. Callers can wrap the result in a MaybeWeight to run the fast fp32 path with no per-call dequant allocation.
Memory impact: 4x vs keeping int8, but for sub-100MB matmul weights this is still small compared to the whole process. Use this at load time if you care about throughput more than RAM.
type RawTensor ¶
type RawTensor struct {
Name string
DType uint8
Shape []int
F32 []float32 // populated when DType == DTypeF32
I8 []int8 // populated when DType == DTypeI8
I32 []int32 // populated when DType == DTypeI32
}
RawTensor is a weight tensor read straight from the hbin file.
type Scorer ¶
type Scorer struct {
// contains filtered or unexported fields
}
Scorer wraps a loaded Model plus its tokenizer and presents the scanner.Scorer contract used by the rest of hush. It is the single integration point for the pure Go runtime.
Construct one via NewScorer when you already have a loaded *Model, or via LoadScorer which reads model + tokenizer from paths. Safe for concurrent use only when the underlying Model is — the current implementation is not (it mutates tensor buffers during Forward). Guard with a sync.Mutex or one Scorer per goroutine if you need parallelism.
func LoadScorer ¶
LoadScorer reads a Model from modelPath (.hbin) and a tokenizer from tokenizerPath (HuggingFace tokenizer.json) and returns a ready Scorer.
func LoadScorerReader ¶
LoadScorerReader is like LoadScorer but takes io.Reader sources so callers can feed embedded bytes or any other stream.
func NewBundledScorer ¶
NewBundledScorer constructs a Scorer from the int8 v1 model and tokenizer embedded in the binary via go:embed. No filesystem access is required.
func NewScorer ¶
NewScorer returns a Scorer using an already-loaded model and tokenizer. maxLen of <= 0 falls back to model.Meta.SeqLen.
func (*Scorer) BatchScore ¶ added in v0.1.6
func (s *Scorer) BatchScore(triples []scanner.SpanTriple) ([]float64, error)
BatchScore scores multiple candidates in a single transformer forward pass. It tokenizes each triple, pads each example to the batch's max encoded length (capped by maxLen), stacks the result into [B, T] int32 arrays, and runs Model.ForwardBatch once.
Returns probabilities in the same order as the input. An empty input yields an empty result. Numerically matches calling Score on each triple to within fp32 reassociation drift; see scanner.Scan for the fallback path used when a scorer doesn't implement this.
Satisfies scanner.BatchScorer.
func (*Scorer) Close ¶
Close is a no-op today (no CGO resources to release), present so the Scorer satisfies the same contract as pkg/classifier.Classifier.
func (*Scorer) ModelVersion ¶
ModelVersion returns the embedded model version from the bundle's meta. The native runtime doesn't hard-code a version the way pkg/classifier does, so we expose the one the user loaded.
type Span ¶ added in v0.1.8
type Span struct {
Start int // char offset, inclusive
End int // char offset, exclusive
Type string // entity type from BIO tag, e.g. "secret", "pii", "noise"
Score float32 // mean softmax-max prob across span tokens
}
Span represents a decoded entity span over character offsets.
func DecodeBIO ¶ added in v0.1.8
func DecodeBIO( logits []float32, K int, id2label map[int]string, offsets [][2]int, attentionMask []int32, ) []Span
DecodeBIO greedy-argmaxes per-token logits, drops O, merges contiguous B-X / I-X runs into a single Span. Lenient: I-X with no preceding B-X (or following a different type) is treated as B-X.
logits is row-major [T, K]. offsets[i] == [2]int{0,0} marks special tokens (skip). attentionMask[i] == 0 marks padding (skip).
type Tensor ¶
type Tensor struct {
Shape []int
Data []float32
// Packed, if non-nil, holds a pre-packed copy of Data in the layout
// consumed by matmulPacked: for a [K, N] tensor this is
// [ceil(N/packNR), K, packNR] (with a smaller tail panel)
// Callers that use this tensor as the B matrix of a matmul should
// prefer MatMulPacked when Packed != nil. Data is still kept for
// correctness fallbacks and for tensors that are transposed/consumed
// differently.
Packed []float32
}
Tensor is a row-major float32 n-dimensional array. Methods prefer operating on existing storage (out-of-place) to simplify reasoning; hot paths can reuse an output tensor via SetData or fused helpers.
func Add ¶
Add element-wise with broadcasting only along leading dims. Shapes must match in the trailing axes, and b's shape must be a prefix of a's. For the common case both shapes equal.
func AddBias ¶
AddBias adds a 1D bias of shape [M] across the last axis of x. x: [..., M]; bias: [M]. Mutates x in place and returns it.
func AddInPlace ¶
AddInPlace computes a += b (same shape). Returns a for chaining.
func ApplyAdditiveMask ¶
ApplyAdditiveMask adds a large negative value to entries of x where mask is 0, so that subsequent softmax zeroes them out. mask must be broadcast-compatible with x over trailing axes.
func BatchMatMul ¶
BatchMatMul computes [B,M,K] x [B,K,N] -> [B,M,N].
func GELU ¶
GELU applies the exact GELU activation using math.Erf, matching the ONNX export path that uses Erf (not the tanh approximation).
GELU(x) = x * 0.5 * (1 + erf(x / sqrt(2)))
func Gather ¶
Gather is the embedding lookup: given a [V, H] table and an index tensor of shape [B, T], returns a [B, T, H] tensor.
func LayerNorm ¶
LayerNorm applies x = (x - mean) / sqrt(var + eps) * gamma + beta over the last axis. Standard fp32 impl matching ONNX LayerNormalization.
func MatMulInt8 ¶
func MatMulInt8(a *Tensor, w *QuantWeight) *Tensor
MatMulInt8 computes A [M, K] fp32 x W (int8 [K, N] + fp32 scale [N]) -> out [M, N] fp32. Correct-first implementation: dequantize each K-row of W into an fp32 buffer once and feed into the standard blocked matmul. This keeps the hot inner loop identical to the fp32 path while paying the dequant cost only once per call (K*N multiplies).
Activations stay fp32; no activation-side quantization.
func Reshape ¶
Reshape returns a tensor sharing t.Data with a new shape (no copy). Panics if the totals do not match.
func ScaleInPlace ¶
ScaleInPlace multiplies every element by s. Returns t.
func Transpose ¶
Transpose rearranges axes. Supports arbitrary n-D tensors. General path is slow; callers wanting attention-shape transposes should use Transpose4D below.
func (*Tensor) PackForMatMul ¶ added in v0.1.4
PackForMatMul pre-packs a 2D [K, N] tensor into the layout consumed by matmulPacked and stores it on t.Packed. Safe to call multiple times (re-packs). Panics if t is not 2D.
Source Files