nlp

package
v0.1.0 Latest Latest
Warning

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

 Go to latest Published: Jul 13, 2026 License: MPL-2.0 Imports: 19 Imported by: 0

Documentation

Overview

Package nlp builds transformer language models — the components that turn a stream of token ids into next-token predictions, and the decoding strategies that turn those predictions back into text.

For AI practitioners

nlp sits at layer L4 on top of the tensor/backend/autograd/nn stack. It provides the pieces a GPT- or Llama-family model is made of and the machinery to run and train them:

  • Tokenization: a byte-level BPE Tokenizer (GPT-2 compatible, encode bit-exact vs tiktoken, decode∘encode byte-exact under fuzzing, §T37), a GPT-2/HF byte-level BPETokenizer, a Unigram tokenizer (SentencePiece, 1-best Viterbi over a scored vocabulary, §T108), and BERT's WordPiece — loadable from a GGUF model's embedded vocabulary (BPEFromGGUF / UnigramFromGGUF) or from HuggingFace tokenizer JSON, so a .gguf or HF checkpoint tokenizes end-to-end.
  • Attention: fused multi-head attention (MHA) with a hand-derived SDPA VJP, grouped/multi-query variants (GQA/MQA), causal masking, ALiBi and sliding-window biases, and a KV-cache for O(1)-per-step incremental decode.
  • Models: the GPT decoder (pre-LN, causal MHA, GELU FFN) and the Llama/Llama-2 decoder (pre-norm RMSNorm, RoPE queries/keys, grouped-query attention, SwiGLU, no biases), assembled from the same verified primitives and loadable straight from safetensors weights or GGUF files — including quantized GGUF, which QuantLlama decodes directly from the ggml Q-blocks without dequantizing to f32.
  • Sampling: a Sampler implementing the standard HuggingFace pipeline (temperature → top-k → top-p nucleus → min-p → epsilon/eta → locally-typical → multinomial), Mirostat adaptive sampling, and repetition / frequency / presence logit penalties — all behind the TokenSampler interface that every sequential generation loop accepts.
  • Search & contrast: beam search and Diverse Beam Search, contrastive search (degeneration penalty), contrastive decoding (expert−amateur logits), DoLa layer-contrast decoding (DoLaDecode, via early-exit layer logits), and Classifier-Free Guidance with negative prompts (CFGDecode).
  • Constrained & marked output: regex/FSM-guided constrained decoding and red/green-list LLM watermarking with statistical detection.
  • Accelerated decoding: lossless speculative decoding with a draft model, draft-model-free prompt-lookup (n-gram), Medusa multi-head drafting (trainable MedusaHeads over ForwardHidden + MedusaGenerate with typical acceptance; MedusaGenerateTree verifies a topK candidate tree in ONE masked forward; the GPU loop lives in llamagpu), and Jacobi parallel decoding (JacobiDecode / GPT.JacobiGenerate).
  • Long-context inference: attention-sink streaming (StreamGenerate), bounded KV-cache eviction policies, SnapKV prompt compression, PyramidKV per-layer cache budgets, an 8-bit quantized KV-cache, and Self-Extend length extrapolation (Llama.SelfExtendForward / SelfExtendGenerate — grouped attention, no fine-tuning; generation stays coherent at 4× training length).
  • Training objectives & data: BERT masked-LM, T5 span corruption, UL2 mixture-of-denoisers, Fill-in-the-Middle transformation, and sequence packing.

Every algorithm is validated on the §V16 ladder: tier-1 bit- or tolerance-exact parity against an official reference (torch / tiktoken / gguf), tier-2 the defining paper, cited in SPEC.md §R. Forward and backward paths are checked against real PyTorch at f64 rtol ~1e-12 (inference) and ~1e-9 (gradients).

For everyone else

A language model reads text as a list of numbers ("tokens"), and for each position it outputs a score for every possible next token — think of a giant weighted dice with one face per word-piece, re-weighted after every token. This package contains both halves of that process:

  • the model that produces the scores (attention, the mechanism that lets each word look back at the earlier words, plus the surrounding layers), and
  • the "decoder" that chooses the next token from the scores. Choosing always the single highest-scoring token is greedy decoding; adding a little controlled randomness (temperature, top-k, top-p, min-p) makes the output more varied and natural. Beam search explores several candidate continuations at once; speculative decoding uses a small fast model to guess ahead and a large model to check the guesses, going faster without changing the result.

The runnable Example functions below show each of these at three levels: a one-line trivial call, a realistic use case, and a piece embedded in a larger decode pipeline.

Package nlp is layer L4: transformer/LLM building blocks (§T21, §T23).

Example (PenaltyThenDecode)

A realistic decode step chains a logit transform with a decoder: penalize tokens already produced, then choose. Token 0 leads the raw logits, but it was already generated, so a repetition penalty of 1.5 divides its logit down (3.0 → 2.0) and the greedy choice shifts to token 1 — anti-repetition steering in two lines. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	logits := []float64{3.0, 2.9, 0.5}
	generated := []int{0} // token 0 already emitted

	nlp.ApplyPenalties(logits, generated, 1.5 /*repetition*/, 0, 0)
	next := nlp.Greedy().Sample(logits)
	fmt.Println(next)
}

Output:
1

Index

Examples

Constants

View Source 
const (
	MedusaEpsilon = 0.3  // hard probability floor ε (posterior_threshold)
	MedusaDelta   = 0.09 // entropy-adaptive scale δ (posterior_alpha)
)

Default typical-acceptance thresholds: a candidate is accepted when the original model's probability exceeds min(epsilon, delta·exp(−H)). These numeric defaults are from the reference implementation (FasterDecoding/Medusa medusa/model/utils.py: posterior_threshold=0.3, posterior_alpha=0.09); the paper itself gives no fixed values (it sweeps ε∈[0.01,0.25] and notes the Hewitt α=√ε relation, §3.3.2).

View Source 
const MLMIgnoreLabel = -100

MLMIgnoreLabel marks a position that is NOT predicted (excluded from the loss) — the same −100 convention as HuggingFace / PyTorch's ignore_index. Real token ids are non-negative, so it never collides with an actual label.

View Source 
const WatermarkZThreshold = 4.0

WatermarkZThreshold is the default z-score above which Detect's result is treated as watermarked (Kirchenbauer et al. 2023 §4: z>4 ⇒ FPR ≈ 3e−5).

Variables

This section is empty.

Functions

func ApplyLoRAGPT 

func ApplyLoRAGPT(g *GPT, r int, alpha float64, seed uint64) ([]*tensor.Tensor, error)

ApplyLoRAGPT attaches rank-r LoRA adapters (Hu et al. 2021) to every attention projection (q, k, v, o) of every block (§T504) and returns the adapters' trainable tensors — the PEFT workflow for the built-in GPT: freeze the model (do NOT pass its Params to the optimizer), train ONLY the returned adapter params, and the base weights stay bit-identical. B is zero-initialized, so attaching adapters does not change the model's output until training. The adapters live on the blocks' MHA.LoRA maps; detach by clearing those maps.

func ApplyPenalties 

func ApplyPenalties(logits []float64, generated []int, repetition, frequency, presence float64)

ApplyPenalties adjusts logits in place using the history of already-generated tokens:

  • repetition penalty (CTRL, Keskar et al. 2019): for each seen token, divide its logit by θ if the logit is positive, else multiply by θ. θ>1 discourages repeats; θ=1 (or ≤0) is a no-op. The sign-aware form (HuggingFace) is used so a negative logit is pushed further negative rather than toward zero.
  • frequency penalty (OpenAI): subtract frequency·count_i (scales with how many times the token already appeared).
  • presence penalty (OpenAI): subtract presence once for any token that appeared.

Tokens that have not been generated are left unchanged. Out-of-range history entries are ignored.

Example

ApplyPenalties discourages repetition before sampling. Token 0 was already generated, so a repetition penalty of 1.5 divides its (positive) logit by 1.5, lowering the chance of picking it again; the other tokens are untouched. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	logits := []float64{3.0, 1.0, 0.5}
	nlp.ApplyPenalties(logits, []int{0}, 1.5 /*repetition*/, 0 /*frequency*/, 0 /*presence*/)
	fmt.Printf("%.4f %.1f %.1f\n", logits[0], logits[1], logits[2])
}

Output:
2.0000 1.0 0.5

func CFGDecode 

func CFGDecode(model *GPT, prompt, negPrompt []int, maxNew int, gamma float64, s TokenSampler) ([]int, error)

CFGDecode generates up to maxNew tokens with Classifier-Free Guidance (§T440, the generation loop around GuidedLogits): the SAME model decodes two streams with their own KV caches — the conditional one seeded with prompt, the unconditional one with negPrompt (the CFG paper's "context-free" branch, or a negative prompt to steer away from) — and at every step the two next-token logit rows are combined by GuidedLogits(·,·,gamma) before sampler s draws. Every generated token is appended to BOTH streams. γ=1 is exactly the model's own decoding of prompt; γ>1 sharpens prompt adherence; with a negative prompt, larger γ pushes further away from it. negPrompt must be non-empty (the unconditional stream needs at least one token to condition its first logits on — typically a BOS token or a minimal neutral prefix). Generation stops at the context limit of the longer stream.

func ContrastiveDecode 

func ContrastiveDecode(expert, amateur *GPT, prompt []int, maxNew int, alpha, beta float64, s TokenSampler) ([]int, error)

ContrastiveDecode generates up to maxNew tokens after prompt with contrastive decoding (§R82): at each step it forms ContrastiveLogits from the expert and amateur next-token distributions and draws from them with sampler s (Greedy for argmax CD). Both models decode with their own KV-cache. With beta=0 the amateur is ignored, so the result is exactly the expert's own decoding.

func ContrastiveLogits 

func ContrastiveLogits(expert, amateur []float64, alpha, beta float64) []float64

ContrastiveLogits combines expert and amateur next-token logits into contrastive scores (§R82). The PLAUSIBILITY CONSTRAINT keeps only tokens the expert rates at ≥ alpha of its own max probability (V_head); all others score −∞. For a kept token v the score is (1+beta)·log p_expert(v) − beta·log p_amateur(v), so beta=0 reduces to the expert's own log-prob (the amateur is ignored) and larger beta contrasts more strongly. alpha≤0 disables the constraint; the paper's default is alpha=0.1. The returned scores are ready to pass to a Sampler (argmax = greedy CD).

Example

The plausibility constraint excludes tokens the expert finds unlikely even when the amateur loves them — here the expert is confident on token 0, so the amateur's favourite (token 2, low expert probability) is masked out of the contrast. 

package main

import (
	"fmt"
	"math"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	expert := []float64{5, 1, -3} // token 2: very low expert probability
	amateur := []float64{0, 0, 5} // amateur strongly prefers token 2
	s := nlp.ContrastiveLogits(expert, amateur, 0.1, 1.0)
	fmt.Println(math.IsInf(s[2], -1)) // excluded by the plausibility constraint
}

Output:
true

func ContrastiveScore 

func ContrastiveScore(probs, maxSim []float64, alpha float64) []float64

ContrastiveScore implements the token ranking of contrastive search (Su, Lan, Wang, Yogatama, Kong & Collier 2022, "A Contrastive Framework for Neural Text Generation" / SimCTG, NeurIPS, arXiv:2202.06417). Ordinary likelihood decoding degenerates into repetitive text because the model keeps re-selecting tokens whose representations are nearly identical to the recent context. Contrastive search counters this by ranking the top-k candidates with a balance of model confidence and a DEGENERATION PENALTY: 

score(v) = (1−α)·prob(v)  −  α·max_j cos(h_v, h_{x_j})

The first term is the model's probability for candidate v; the second is the maximum cosine similarity of v's representation h_v to the representations of the already generated tokens (a high value means v would repeat the context). The candidate with the highest score is chosen. α ∈ [0,1] balances the two: α=0 is greedy decoding by probability, larger α more strongly discourages repetition (the paper uses α=0.6, top-k=4-8).

probs holds the candidates' model probabilities and maxSim their per-candidate maximum cosine similarity to the context (compute it with MaxContextCosine). The returned scores are aligned with the candidates; the argmax is the selected token.

Example 
// two candidates: token 0 is likely but repeats the context; token 1 is a bit less
// likely but novel. With α=0.6 contrastive search prefers the novel token 1.
probs := []float64{0.55, 0.45}
maxSim := []float64{0.95, 0.10}
scores := nlp.ContrastiveScore(probs, maxSim, 0.6)
fmt.Println(argmax(scores))

Output:
1

func DoLaDecode 

func DoLaDecode(model *GPT, prompt []int, maxNew int, layers []int, alpha float64, s TokenSampler) ([]int, error)

DoLaDecode generates up to maxNew tokens with DoLa layer-contrast decoding (§T442, the generation loop around DoLaLogits): each step runs ForwardEarlyExit over the full running sequence, contrasts the mature next-token logits against the JSD-selected premature layer among layers, and draws with sampler s from the contrastive scores (restricted to the adaptive-plausibility set, α as in DoLaLogits). No KV cache — every step is a full forward (O(n) forwards; DoLa taps hidden states DecodeStep does not expose), so it suits analysis-scale runs. α=1 keeps only the mature argmax, making greedy DoLa exactly greedy decoding.

func DoLaLogits 

func DoLaLogits(mature []float64, premature [][]float64, alpha float64) ([]float64, int)

DoLaLogits implements DoLa decoding (Chuang, Xie, Luo, Kim, Glass & He 2023, "DoLa: Decoding by Contrasting Layers Improves Factuality of Large Language Models", arXiv:2309.03883, ICLR'24). Factual knowledge in a transformer tends to be resolved in the LATER layers, so DoLa amplifies it by contrasting the final ("mature") layer's next-token distribution against that of an earlier ("premature") layer — obtained by projecting an intermediate hidden state through the SAME output head (early exit). Unlike model-level contrastive decoding (ContrastiveLogits, which needs a second amateur model), DoLa contrasts LAYERS of the one model.

Given the mature-layer logits and the logits of the candidate premature layers, DoLa (1) selects the premature layer M whose distribution is MOST divergent from the mature one, M = argmax_{j} JSD(q_N ‖ q_j) (the most contrastive layer), and (2) returns the contrastive score F(x) = log q_N(x) − log q_M(x) restricted to the adaptive-plausibility set V = { x : q_N(x) ≥ α·max_w q_N(w) } (α default 0.1), scoring −∞ elsewhere so implausible tokens are never selected. The returned scores and the selected premature-layer index are ready to pass to a Sampler (argmax = greedy DoLa).

Example 
package main

import (
	"fmt"
	"math"

	"github.com/jxsl13/goai/nlp"
)

func argmax(x []float64) int {
	best, bv := 0, math.Inf(-1)
	for i, v := range x {
		if v > bv {
			best, bv = i, v
		}
	}
	return best
}

func main() {
	// mature layer favors token 1; the (only) premature layer favors token 0, so
	// contrasting sharpens token 1 as the factual choice.
	mature := []float64{1, 2, 0}
	premature := [][]float64{{3, 0, 0}}
	out, sel := nlp.DoLaLogits(mature, premature, 0.1)
	fmt.Printf("selected=%d argmax=%d\n", sel, argmax(out))
}

Output:
selected=0 argmax=1

func DocumentCausalMask 

func DocumentCausalMask(docIDs []int) *tensor.Tensor

DocumentCausalMask builds the [L,L] additive attention mask for a packed block given the per-position document ids: entry [i,j] is 0 when token j is in the same document as token i AND j ≤ i (intra-document causal), and −∞ otherwise. This is block-diagonal — each document is an independent causal block — so no token attends across a document boundary. A block of a single document reproduces the ordinary causal mask.

func DocumentPositions 

func DocumentPositions(docIDs []int) []int

DocumentPositions returns the position ids for a packed block: they restart at 0 at each document boundary, so every document sees positions as if decoded alone (e.g. document ids [0,0,0,1,1] → [0,1,2,0,1]). Documents are assumed contiguous within the block, as produced by PackSequences.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// Two documents (lengths 3 and 2) packed together: the mask keeps them separate and the
	// position ids restart per document.
	_, docIDs := nlp.PackSequences([][]int{{10, 11, 12}, {20, 21}}, 5)
	fmt.Println(docIDs[0], nlp.DocumentPositions(docIDs[0]))
}

Output:
[0 0 0 1 1] [0 1 2 0 1]

func FIMReconstruct 

func FIMReconstruct(fim []int, s FIMSentinels, spm bool) (doc []int, ok bool)

FIMReconstruct recovers the original document (prefix+middle+suffix) from a FIM-transformed sequence, returning ok=false if it is not a well-formed FIM sequence for the given mode. It is the exact inverse of FIMTransform, so FIMReconstruct(FIMTransform(doc)) == doc.

func FIMTransform 

func FIMTransform(tokens []int, s FIMSentinels, spm bool, rng *rand.Rand) []int

FIMTransform reorders tokens for FIM training. It splits the document at two random points into prefix/middle/suffix (document = prefix+middle+suffix) and emits either the PSM layout 

<PRE> prefix <SUF> suffix <MID> middle

or, when spm is true, the SPM layout 

<PRE> <SUF> suffix <MID> prefix middle

so in both cases the middle is produced last. rng draws the two split points; the sentinels must not appear among the document tokens. The original tokens are not modified.

Example 
package main

import (
	"fmt"
	"math/rand/v2"
	"slices"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// Reorder a document into PSM fill-in-the-middle layout, then recover the original.
	s := nlp.FIMSentinels{Pre: 1000, Suf: 1001, Mid: 1002}
	doc := []int{1, 2, 3, 4}
	rng := rand.New(rand.NewPCG(5, 5))
	fim := nlp.FIMTransform(doc, s, false, rng)
	back, _ := nlp.FIMReconstruct(fim, s, false)
	fmt.Println(slices.Equal(back, doc))
}

Output:
true

func GatherRows 

func GatherRows(t *tensor.Tensor, idx []int) *tensor.Tensor

GatherRows selects rows idx (in order) from t[n,d] → [len(idx),d]. Used to apply an eviction policy's kept-index list to a cached K or V tensor (inference-only, no gradient path).

func GuidedLogits 

func GuidedLogits(conditional, unconditional []float64, gamma float64) []float64

GuidedLogits applies Classifier-Free Guidance to autoregressive decoding (Sanchez, Fan, Spangher, Levi, Ammanabrolu & Biderman 2023, "Stay on topic with Classifier-Free Guidance", arXiv:2306.17806). At each step the model is run twice — once WITH the prompt/context (conditional next-token logits) and once WITHOUT it (or with a negative prompt: unconditional logits) — and the two next-token distributions are linearly extrapolated in log-probability space to sharpen the influence of the prompt: 

log P_cfg(w) = log P(w) + γ·( log P(w|c) − log P(w) )
            = γ·log P(w|c) − (γ−1)·log P(w)

γ = 1 recovers the ordinary conditional distribution; γ > 1 (typically ≈1.5) pushes the distribution toward tokens the prompt makes MORE likely than the unconditional model, improving prompt adherence. Because the unconditional branch can be a NEGATIVE prompt, the same rule steers generation AWAY from it (negative prompting).

conditional and unconditional are the two raw next-token logit vectors (same length). The returned scores are in log-probability space (a valid logit vector for a Sampler, which re-normalizes with softmax — argmax = greedy CFG).

Example 
// prompt strongly favors token 1; guidance γ=1.5 sharpens that preference, so the
// greedy (argmax) choice is token 1.
cond := []float64{0, 3, 1}
uncond := []float64{1, 0, 1} // no/negative prompt
guided := nlp.GuidedLogits(cond, uncond, 1.5)
fmt.Println(argmax(guided))

Output:
1

func H2OKeep 

func H2OKeep(scores []float64, recent, budget int) []int

H2OKeep returns the token indices an H2O cache of budget size retains given the per-token accumulated-attention scores: the union of the most recent `recent` tokens and the top heavy-hitters (highest scores) among the rest, up to budget (Zhang et al. 2023, Alg. 1). When the token count exceeds budget the lowest- scoring non-recent tokens are evicted. Returns ascending indices (temporal order preserved). recent is capped at budget; if n ≤ budget all are kept.

func H2OScores 

func H2OScores(attn [][]float64) []float64

H2OScores returns the H2O accumulated-attention score of each key: the column sum of the post-softmax attention matrix, score(j) = Σ_i A[i,j] over all query rows i (Zhang et al. 2023). attn is [queries][keys]; the result is length-keys. A key that received a lot of attention across queries is a "heavy hitter".

func JacobiDecode 

func JacobiDecode(step JacobiStep, prompt, init []int, maxIters int) (gen []int, iters int)

JacobiDecode runs Jacobi parallel decoding for len(init) continuation tokens after prompt, starting from the guess init and iterating (at most maxIters times) until the generated block stops changing. It returns the generated tokens — equal to sequential greedy decoding of the same model — and the number of parallel iterations used. The initial guess only affects the iteration count, never the converged result.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

// mockLM is a deterministic causal model: the token predicted after position j is the sum
// of the prefix seq[:j+1] modulo vocab — a genuine dependency on the whole prefix.
func mockLM(vocab int) nlp.JacobiStep {
	return func(seq []int) []int {
		out := make([]int, len(seq))
		sum := 0
		for j, t := range seq {
			sum += t
			out[j] = ((sum % vocab) + vocab) % vocab
		}
		return out
	}
}

// greedyRef decodes genLen tokens sequentially (the ground truth Jacobi must match).
func greedyRef(step nlp.JacobiStep, prompt []int, genLen int) []int {
	y := make([]int, 0, genLen)
	for i := 0; i < genLen; i++ {
		full := append(append([]int(nil), prompt...), y...)
		preds := step(full)
		y = append(y, preds[len(full)-1])
	}
	return y
}

func main() {
	// A toy model: the next token is the running prefix-sum mod 5. Jacobi decoding recovers
	// the same sequence as sequential greedy, in a few parallel iterations.
	step := mockLM(5)
	gen, iters := nlp.JacobiDecode(step, []int{1, 2}, make([]int, 4), 100)
	fmt.Println(gen, "in", iters, "iters ==", greedyRef(step, []int{1, 2}, 4))
}

Output:
[3 1 2 4] in 5 iters == [3 1 2 4]

func JensenShannon 

func JensenShannon(p, q []float64) float64

JensenShannon returns the Jensen-Shannon divergence (in nats) between two probability distributions p and q of equal length: JSD = ½·KL(p‖m) + ½·KL(q‖m) with m = (p+q)/2. It is symmetric and bounded in [0, ln 2]. Terms with p_i = 0 (or q_i = 0) contribute 0, per the 0·log0 = 0 convention.

func LlamaGGUFQuantMix 

func LlamaGGUFQuantMix(ts map[string]*tensor.Tensor) map[string]gguf.QuantType

LlamaGGUFQuantMix returns the per-tensor quantization map for the llama.cpp LLAMA_FTYPE_MOSTLY_Q4_K_M mix (§R101), given the GGUF tensor map from LlamaToGGUF. Pass the result to gguf.WriteQuantized to write a real ~4.5-bit Q4_K_M model: 

meta, ts := nlp.LlamaToGGUF(m)
gguf.WriteQuantized(w, &gguf.File{Version: 3, Metadata: meta, Tensors: ts}, nlp.LlamaGGUFQuantMix(ts))

The recipe (research-lite CONFIRMED unanimous vs llama-quant.cpp llama_tensor_get_type): base type Q4_K for every 2-D projection; output.weight and the use_more_bits layers of attn_v.weight / ffn_down.weight are upgraded to the higher-precision Q6_K (§R99); the 1-D RMSNorm gains and token_embd stay their natural type (norms F32, token_embd Q4_K).

A k-quant super-block is 256 elements along a row, so a tensor whose GGUF row length (shape[1]) is not a multiple of 256 CANNOT be k-quantized; the helper leaves it out of the map (written F32) rather than producing an invalid file — real Llama geometries (Dim, Hidden multiples of 256) are always covered, tiny/odd ones degrade gracefully.

Example

LlamaGGUFQuantMix picks the llama.cpp Q4_K_M per-tensor type recipe for a Llama's GGUF tensors, ready to hand to gguf.WriteQuantized for a ~4.5-bit model file. 

package main

import (
	"fmt"
	"strings"

	"github.com/jxsl13/goai/format/gguf"
	"github.com/jxsl13/goai/nlp"
)

func main() {
	m, _ := nlp.NewLlama(nlp.LlamaConfig{
		Vocab: 8, Ctx: 8, Dim: 256, Heads: 4, KVHeads: 4, Layers: 1, Hidden: 256, Eps: 1e-5, RopeBase: 10000,
	}, 1)
	_, ts := nlp.LlamaToGGUF(m)
	qm := nlp.LlamaGGUFQuantMix(ts)
	// output.weight and the (single, use_more_bits) layer's attn_v/ffn_down are Q6_K; the
	// other projections are Q4_K; the three 1-D norms stay F32 (absent from the map).
	q6 := 0
	for _, qt := range qm {
		if qt == gguf.Q6_K {
			q6++
		}
	}
	fmt.Println("quantized:", len(qm), "Q6_K:", q6, "output.weight:", qm["output.weight"] == gguf.Q6_K,
		"norms F32:", !strings.Contains(fmt.Sprint(qm), "attn_norm"))
}

Output:
quantized: 9 Q6_K: 3 output.weight: true norms F32: true

func LlamaToGGUF 

func LlamaToGGUF(m *Llama) (map[string]any, map[string]*tensor.Tensor)

LlamaToGGUF is the inverse of LlamaFromGGUF: it serializes a Llama into the GGUF metadata + tensor maps (ggml/llama.cpp names and layout), transposing every projection back into torch [out, in]. Pass the result to gguf.Write via a gguf.File.

func MLMMask 

func MLMMask(tokens []int, maskProb float64, maskID, vocabSize int, rng *rand.Rand) (input, labels []int)

MLMMask applies BERT MLM corruption to tokens. Each position is selected with probability maskProb (BERT: 0.15); a selected position is replaced by maskID with prob 0.8, by a uniform random id in [0,vocabSize) with prob 0.1, and left unchanged with prob 0.1. It returns the corrupted input and the labels: labels[i] is the ORIGINAL token at a selected position and MLMIgnoreLabel elsewhere (so the loss is computed only on the 15%). Exclude special tokens upstream. The original tokens are not modified; the pair is losslessly reversible (MLMReconstruct).

Example 
package main

import (
	"fmt"
	"math/rand/v2"
	"slices"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	doc := []int{10, 11, 12, 13, 14, 15}
	rng := rand.New(rand.NewPCG(2, 3))
	input, labels := nlp.MLMMask(doc, 0.5, 0, 1000, rng)
	back, ok := nlp.MLMReconstruct(input, labels)
	fmt.Println(ok, slices.Equal(back, doc))
}

Output:
true true

func MLMMaskExcluding 

func MLMMaskExcluding(tokens []int, maskProb float64, maskID, vocabSize int, specialIDs []int, rng *rand.Rand) (input, labels []int)

MLMMaskExcluding is MLMMask with special-token protection: any position whose token is in specialIDs (e.g. [CLS], [SEP], padding) is NEVER selected for masking — its label stays MLMIgnoreLabel and its input is left unchanged — matching BERT/HuggingFace, which build a special-tokens mask so those positions never contribute to the MLM loss. The ~maskProb rate then applies only to the ordinary tokens. Special positions consume no randomness, so with a nil/empty specialIDs this reduces to MLMMask exactly (same RNG sequence). The pair is still losslessly reversible with MLMReconstruct.

Example 
package main

import (
	"fmt"
	"math/rand/v2"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// [CLS]=1 and [SEP]=2 are protected: with a 100% rate, only the content token is masked.
	doc := []int{1, 10, 2}
	in, lab := nlp.MLMMaskExcluding(doc, 1.0, 99, 1000, []int{1, 2}, rand.New(rand.NewPCG(5, 5)))
	fmt.Println(lab[0] == nlp.MLMIgnoreLabel, lab[2] == nlp.MLMIgnoreLabel, in[1] == 99)
}

Output:
true true true

func MLMReconstruct 

func MLMReconstruct(input, labels []int) (doc []int, ok bool)

MLMReconstruct recovers the original document from an (input, labels) MLM pair: a predicted position (labels[i] ≠ MLMIgnoreLabel) takes its label — the original token — and every other position keeps the (unmodified) input token. It is the exact inverse of MLMMask. ok is false if the lengths differ.

func MaxContextCosine 

func MaxContextCosine(candReps, contextReps [][]float64) []float64

MaxContextCosine returns, for each candidate representation, the maximum cosine similarity to any of the context (previously generated) token representations — the degeneration-penalty term of contrastive search. candReps is [numCandidates][dim], contextReps is [numContext][dim]. With no context every penalty is 0.

func MedusaTreeMask 

func MedusaTreeMask(parent []int) (mask *tensor.Tensor, depth []int)

MedusaTreeMask builds the tree attention mask and depth-based position ids for a candidate tree whose node i has parent parent[i] (a root has parent −1). Nodes must be given in an order where every parent precedes its children (parent[i] ∈ {−1} ∪ [0,i)). In Medusa's single-pass verification each candidate token may attend ONLY to its ancestors along the path back to the root (its true prefix), never to tokens in sibling branches — so one forward pass scores every root-to-node path independently. The returned mask is the additive pre-softmax bias: mask[i][j] = 0 when j is an ancestor of i or j==i, and −∞ otherwise (add it to the attention logits, as a causal mask). depth[i] is node i's depth in the tree (root 0), used as its position id. A linear chain reproduces the ordinary causal mask.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// root 0 with two children 1,2; token 2 attends its ancestor 0 and itself, but NOT the
	// sibling branch 1.
	mask, depth := nlp.MedusaTreeMask([]int{-1, 0, 0})
	fmt.Printf("2→0:%v 2→1:%v depth:%v\n", mask.AtF64(2, 0) == 0, mask.AtF64(2, 1) == 0, depth)
}

Output:
2→0:true 2→1:false depth:[0 1 1]

func NgramLookup 

func NgramLookup(seq []int, maxNgram, draftLen int) []int

NgramLookup is the prompt-lookup drafter (Saxena; the reference-copy of LLMA, Yang et al. 2023, §R89): it proposes a continuation for seq WITHOUT a draft model by copying from seq's own history. Trying the longest n-gram first (maxNgram down to 1), it takes the suffix of that length, finds an EARLIER occurrence of it in seq, and returns the up-to-draftLen tokens that FOLLOWED that occurrence — the guess that the same continuation repeats. Returns nil when no n-gram recurs. Because every draft is verified against the model, a wrong guess only costs speed, never correctness. Exported for external speculative loops (e.g. the llamagpu batched decoders, §T426).

func NoRepeatNGramBlock 

func NoRepeatNGramBlock(logits []float64, generated []int, n int)

NoRepeatNGramBlock enforces the no_repeat_ngram_size decoding constraint (HuggingFace NoRepeatNGramLogitsProcessor / fairseq; tri-gram blocking from Paulus, Xiong & Socher 2017, "A Deep Reinforced Model for Abstractive Summarization", arXiv:1705.04304). It HARD-blocks the model from ever repeating any n-gram: a next-token candidate t is forbidden if appending it would recreate an n-gram (context + t) that already occurred in the generated sequence.

The current context is the last n−1 generated tokens. For every earlier position whose (n−1)-gram equals that context, the token that FOLLOWED it is banned by setting its logit to −∞, so it can never be sampled this step. Unlike the soft repetition/frequency penalties (ApplyPenalties), this is a hard constraint. It is a no-op when n ≤ 0 or the sequence is shorter than n; n = 3 is the classic tri-gram blocking. logits is modified in place; out-of-range history tokens are ignored.

Example 
package main

import (
	"fmt"
	"math"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// "the cat sat the cat": with n=3 the context "the cat" was already followed by
	// "sat", so repeating "sat" (token 2) is blocked.
	logits := []float64{0, 0, 0, 0} // the=0 cat=1 sat=2 mat=3
	nlp.NoRepeatNGramBlock(logits, []int{0, 1, 2, 0, 1}, 3)
	fmt.Printf("sat blocked=%v mat free=%v\n", math.IsInf(logits[2], -1), logits[3] == 0)
}

Output:
sat blocked=true mat free=true

func PackSequences 

func PackSequences(seqs [][]int, maxLen int) (blocks, docIDs [][]int)

PackSequences bin-packs seqs into blocks of length ≤ maxLen using first-fit-decreasing, returning the packed token blocks and, for each block, the local document index of every position (0 for the first document in the block, 1 for the next, …). Sequences longer than maxLen are truncated to maxLen; empty sequences are skipped.

func PyramidKVBudgets 

func PyramidKVBudgets(numLayers, totalBudget, minBudget int) []int

PyramidKVBudgets returns the per-layer KV budgets: a linearly decreasing sequence from a maximum at layer 0 (bottom) down to minBudget at the top layer, with the total conserved (Σ budgets == numLayers·(totalBudget/numLayers) == totalBudget). The maximum is b_max = 2·avg − minBudget so the mean of the endpoints equals the average budget avg = totalBudget/numLayers. minBudget is clamped to [0, avg] (minBudget = avg gives a uniform allocation). Integer rounding residue is folded into the bottom (largest) layer so the sum is exact.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// 6 layers, total budget 600 (avg 100), top-layer minimum 40: a decreasing pyramid whose
	// budgets still sum to 600.
	fmt.Println(nlp.PyramidKVBudgets(6, 600, 40))
}

Output:
[160 136 112 88 64 40]

func SelfExtendPositions 

func SelfExtendPositions(seqLen, window, group int) [][]int

SelfExtendPositions builds the [seqLen, seqLen] causal matrix of effective relative positions under Self-Extend: entry [q,k] is SelfExtendRelPos(q,k,window,group) for k ≤ q and 0 above the diagonal (masked). Feed these relative positions to a RoPE/relative attention so a model trained on a short context attends correctly over a longer one.

func SelfExtendRelPos 

func SelfExtendRelPos(q, k, window, group int) int

SelfExtendRelPos returns the effective RELATIVE position a Self-Extend model uses for a query at index q attending to a key at index k (q ≥ k), with neighbor window window and group size group (group ≥ 1). Within the window it is the true distance q−k; beyond it, the grouped distance ⌊q/G⌋−⌊k/G⌋ + (w − ⌊w/G⌋), the shift keeping it continuous with the neighbor region.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// Window 4, group 8. A key 3 back is a neighbor (relpos 3); a key 800 back is grouped and
	// compressed to about 800/8 + shift.
	fmt.Println(nlp.SelfExtendRelPos(1000, 997, 4, 8), nlp.SelfExtendRelPos(1000, 200, 4, 8))
}

Output:
3 104

func SnapKVKeep 

func SnapKVKeep(obsAttn [][]float64, budget, kernel int) []int

SnapKVKeep returns the prompt token indices SnapKV retains, given obsAttn — the attention from the observation-window queries (the last len(obsAttn) prompt tokens) to all len(obsAttn[0]) prompt keys, per one head. budget is the prompt KV budget (max_capacity_prompt); kernel is the pooling kernel (default 7). It keeps the whole observation window plus the top (budget − windowSize) earlier positions by pooled importance, returned ascending (temporal order). If the prompt already fits the budget all indices are kept.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// 6-token prompt, observation window = last 2 (positions 4,5). Early key 1 is the heavy
	// hitter; with budget 3 SnapKV keeps it plus the window.
	obsAttn := [][]float64{
		{0.05, 0.80, 0.05, 0.02, 0.05, 0.03},
		{0.05, 0.75, 0.05, 0.05, 0.05, 0.05},
	}
	fmt.Println(nlp.SnapKVKeep(obsAttn, 3, 1))
}

Output:
[1 4 5]

func SnapKVPool 

func SnapKVPool(scores []float64, kernel int) []float64

SnapKVPool applies SnapKV's 1D max-pooling (kernel odd, stride 1, "same" length, edges clamped) to importance scores. Pooling spreads a peak to its neighbours so that top-k selection retains a cluster around each important position rather than isolated tokens — the ablation-critical step of SnapKV. kernel ≤ 1 returns the scores unchanged.

func SpeculativeRun 

func SpeculativeRun(target, draft [][]float64, draftTokens []int, rng *rand.Rand) []int

SpeculativeRun verifies a whole run of drafted tokens (Leviathan Algorithm 1). target holds the target distributions for positions 0..K (K+1 of them), draft the draft distributions for positions 0..K−1, and draftTokens the K proposed tokens. Tokens are accepted left-to-right until the first rejection (which emits one residual token and discards the rest); if all K are accepted a bonus token is sampled from target[K]. It returns between 1 and K+1 tokens, each distributed as the corresponding target position (§R53).

Example

Speculative decoding verifies a whole run of draft tokens in one target pass. Here the target fully agrees with the draft (matching one-hot distributions), so all K=3 proposals are accepted and a bonus 4th token is appended — up to K+1 tokens produced from a single target forward pass, with the output distributed exactly as the target model would sample. 

package main

import (
	"fmt"
	"math/rand/v2"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	oh := func(at int) []float64 { d := make([]float64, 5); d[at] = 1; return d }
	target := [][]float64{oh(1), oh(2), oh(3), oh(4)} // K+1 = 4 target positions
	draft := [][]float64{oh(1), oh(2), oh(3)}         // K = 3 draft positions
	rng := rand.New(rand.NewPCG(1, 1))

	out := nlp.SpeculativeRun(target, draft, []int{1, 2, 3}, rng)
	fmt.Println(out)
}

Output:
[1 2 3 4]

func SpeculativeSample 

func SpeculativeSample(target, draft []float64, draftToken int, rng *rand.Rand) (token int, accepted bool)

SpeculativeSample performs one modified-rejection-sampling step: given the target distribution p and the draft distribution q (both length-V probability vectors) and a token drafted from q, it accepts the draft with probability min(1, p[t]/q[t]); on rejection it resamples from the normalized residual max(0, p−q). The returned token is distributed exactly as p for ANY q — the losslessness guarantee (Leviathan Thm 1 / Chen). rng draws the accept coin and any residual sample.

Example

SpeculativeSample is the single-token core of speculative decoding: accept the drafted token with probability min(1, p/q), else resample the residual. When the draft distribution q equals the target p (here both concentrate on token 1), p/q = 1, so the draft is always accepted and the output is exactly the target's. 

package main

import (
	"fmt"
	"math/rand/v2"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	target := []float64{0.2, 0.8} // large "target" model's distribution
	draft := []float64{0.2, 0.8}  // small "draft" model happens to agree
	rng := rand.New(rand.NewPCG(1, 1))

	tok, accepted := nlp.SpeculativeSample(target, draft, 1 /*drafted token*/, rng)
	fmt.Println(tok, accepted)
}

Output:
1 true

func StreamingKeep 

func StreamingKeep(n, sink, recent int) []int

StreamingKeep returns the token indices a StreamingLLM cache retains for a cache of n tokens: the first `sink` "attention sink" tokens (kept permanently) plus the most recent `recent` tokens, evicting the middle (Xiao et al. 2023). The sinks matter because softmax must sum to 1, so a model dumps excess attention onto the always-visible initial tokens regardless of content; keeping ~4 of them restores window-attention quality once the window rolls past the start. If sink+recent ≥ n nothing is evicted (all indices returned). Indices are ascending, so gathering them preserves temporal order (and positions are then reassigned by position-within-cache, per the paper).

Example

StreamingLLM keeps a few "attention sink" tokens at the start plus a rolling window of recent tokens, so a fixed-size cache can decode indefinitely. Here a 20-token cache is bounded to 4 sinks + 6 recent = 10 entries. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	keep := nlp.StreamingKeep(20, 4 /*sinks*/, 6 /*recent*/)
	fmt.Println(keep)
}

Output:
[0 1 2 3 14 15 16 17 18 19]

func T5Corrupt 

func T5Corrupt(tokens []int, density, meanSpan float64, sentinelBase int, rng *rand.Rand) (input, target []int)

T5Corrupt applies span corruption to a document. About `density` of the tokens are dropped (T5 default 0.15), grouped into spans of mean length `meanSpan` (T5 default 3). Each dropped span is replaced in `input` by a unique sentinel (sentinelBase, sentinelBase+1, …) and appears in `target` behind that sentinel; `target` ends with one final sentinel. The sentinel ids must be reserved (≥ sentinelBase, above every document token). rng draws the span/gap segmentation. The original tokens are not modified.

Example 
package main

import (
	"fmt"
	"math/rand/v2"
	"slices"

	"github.com/jxsl13/goai/nlp"
)

const sentBase = 10000

func main() {
	// Corrupt a short document; reconstruction recovers it.
	doc := []int{10, 11, 12, 13, 14, 15}
	rng := rand.New(rand.NewPCG(5, 5))
	in, tg := nlp.T5Corrupt(doc, 0.34, 2, sentBase, rng)
	back, _ := nlp.T5Reconstruct(in, tg, sentBase)
	fmt.Println(slices.Equal(back, doc))
}

Output:
true

func T5Reconstruct 

func T5Reconstruct(input, target []int, sentinelBase int) (doc []int, ok bool)

T5Reconstruct recovers the original document from a (input, target) span-corruption pair by splicing each sentinel in the input with its span from the target. It is the exact inverse of T5Corrupt (returning ok=false on a malformed pair). A token is treated as a sentinel iff it is ≥ sentinelBase.

func TypicalAcceptance 

func TypicalAcceptance(probs []float64, token int, epsilon, delta float64) bool

TypicalAcceptance reports whether a Medusa-head candidate token should be accepted under the paper's TYPICAL ACCEPTANCE scheme (§2.3.1, after Hewitt et al. 2022 typical sampling): accept iff the ORIGINAL model's probability for the token exceeds 

min(epsilon, delta·exp(−H(probs)))

where H(probs) is the Shannon entropy (in nats) of the original model's next-token distribution at that position. The hard floor epsilon guarantees a minimum quality bar; the entropy term relaxes it where the model is uncertain (high H ⇒ lower threshold ⇒ more candidates accepted) and tightens it where the model is confident. This replaces speculative decoding's rejection sampling: it is not distribution-exact but keeps only "typical" tokens, which the paper finds preserves generation quality while accepting more tokens per step. probs is the original model's distribution over the vocabulary; token indexes it. In a decode loop the first (greedy) token of a path is always accepted and decoding takes the longest prefix of a candidate path that each pass this test.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// Under a confident distribution the top token clears the threshold; the 0.02 tail does
	// not.
	probs := []float64{0.9, 0.08, 0.02}
	fmt.Println(nlp.TypicalAcceptance(probs, 0, nlp.MedusaEpsilon, nlp.MedusaDelta),
		nlp.TypicalAcceptance(probs, 2, nlp.MedusaEpsilon, nlp.MedusaDelta))
}

Output:
true false

func UL2Denoise 

func UL2Denoise(tokens []int, d UL2Denoiser, sentinelBase int, rng *rand.Rand) (input, target []int)

UL2Denoise applies denoiser d to a document, returning the encoder input and decoder target. The input is prefixed with d's paradigm token (sentinelBase+int(d.Mode)); span sentinels start at sentinelBase+ul2NumModes. R/X delegate to T5 span corruption; S emits a single trailing prefix-LM span. The result is exactly reversible with UL2Reconstruct. The original tokens are not modified.

Example 
package main

import (
	"fmt"
	"math/rand/v2"
	"slices"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	doc := []int{10, 11, 12, 13, 14, 15, 16, 17}
	// An S-denoiser (prefix-LM): the model sees a prefix and predicts the tail.
	d := nlp.UL2Denoiser{Mode: nlp.UL2S, Rate: 0.25}
	in, tg := nlp.UL2Denoise(doc, d, 10000, rand.New(rand.NewPCG(1, 1)))
	back, mode, ok := nlp.UL2Reconstruct(in, tg, 10000)
	fmt.Println(mode == nlp.UL2S, ok, slices.Equal(back, doc))
}

Output:
true true true

Types

type BPEOption 

type BPEOption func(*BPETokenizer)

BPEOption configures a BPETokenizer (functional-options idiom, §C12).

func WithBPEUnkID 

func WithBPEUnkID(id int) BPEOption

WithBPEUnkID sets the id emitted for a symbol absent from the vocabulary (rare for a complete byte-level vocab; default: none, unknown symbols are skipped).

type BPETokenizer 

type BPETokenizer struct {
	// contains filtered or unexported fields
}

BPETokenizer is a GPT-2 / HuggingFace byte-level BPE tokenizer (Sennrich et al. 2016 BPE §R33; GPT-2 byte-level variant, Radford 2019, §R90) — the tokenizer of the Llama-3, Qwen and Mistral GGUF model families ("gpt2"/"bpe"). Text is pre-tokenized (the GPT-2 regex), each pre-token's raw bytes are mapped through bytesToUnicode, and BPE merges are applied by RANK: the vocabulary's merge list is an ordered set of symbol pairs, and at each step the adjacent pair with the lowest rank (earliest in the list) is merged, until none remains. Because all 256 byte code points are base tokens, decode∘encode is byte-exact for any input (§V15).

Example

A byte-level BPE tokenizer merges characters into subwords by rank: with the merge "h i", the two letters of "hi" combine into the single token "hi", and decoding restores the text. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	tok, _ := nlp.NewBPE([]string{"h", "i", "hi"}, []string{"h i"})
	ids := tok.Encode("hi")
	fmt.Println(ids, tok.Decode(ids))
}

Output:
[2] hi

func BPEFromGGUF 

func BPEFromGGUF(meta map[string]any) (*BPETokenizer, error)

BPEFromGGUF builds a byte-level BPE tokenizer from the metadata map of a parsed GGUF model file (§R88): tokenizer.ggml.model must be "gpt2" or "bpe", with tokenizer.ggml.tokens (the byte-mapped vocabulary) and tokenizer.ggml.merges (the ordered "left right" merge rules). This wires a real .gguf model's embedded BPE tokenizer — the Llama-3 / Qwen / Mistral family — to NewBPE, so weights and tokenizer load from one file. SentencePiece/Unigram models ("t5"/"llama") use UnigramFromGGUF.

func BPEFromJSON 

func BPEFromJSON(data []byte, opts ...BPEOption) (*BPETokenizer, error)

BPEFromJSON builds a byte-level BPETokenizer from HuggingFace tokenizer.json bytes. It errors on invalid JSON, a non-BPE model type, or an empty vocabulary.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	j := `{"model":{"type":"BPE","vocab":{"h":0,"i":1,"hi":2},"merges":["h i"]}}`
	tok, _ := nlp.BPEFromJSON([]byte(j))
	fmt.Println(tok.Encode("hi"))
}

Output:
[2]

func NewBPE 

func NewBPE(vocab []string, merges []string, opts ...BPEOption) (*BPETokenizer, error)

NewBPE builds a byte-level BPE tokenizer from a vocabulary (byte-mapped token strings indexed by id) and an ordered merges list ("left right" per entry, lowest index = highest priority) — the form GGUF stores. It errors on an empty vocabulary.

func (*BPETokenizer) Decode 

func (t *BPETokenizer) Decode(ids []int) string

Decode reconstructs the original text: token ids → byte-mapped symbols → invert the byte→Unicode map back to raw bytes. Byte-exact for any input the vocabulary covers.

func (*BPETokenizer) Encode 

func (t *BPETokenizer) Encode(text string) []int

Encode turns text into token ids: GPT-2 pre-tokenization → byte→Unicode mapping → rank-ordered BPE merges → vocab lookup.

func (*BPETokenizer) ToJSON 

func (t *BPETokenizer) ToJSON() ([]byte, error)

ToJSON serializes the tokenizer to the minimal HuggingFace tokenizer.json BPE form (model.type "BPE" with vocab and merges), so BPEFromJSON(t.ToJSON()) round-trips it.

type Beam 

type Beam struct {
	Tokens []int   // the decoded token sequence
	Score  float64 // cumulative log-probability of the sequence
}

Beam is a decoded hypothesis and its score (cumulative log-probability, length-normalized on completion).

func BeamSearch 

func BeamSearch(next NextLogits, start []int, width, maxNew, eos int, alpha float64) []Beam

BeamSearch decodes with beam width `width`. It expands each live hypothesis by every token, scores extensions by the cumulative sum of per-token log-probs, and keeps the top `width`. A hypothesis emitting `eos` (use eos<0 to disable) is completed; a hypothesis reaching `maxNew` new tokens is completed too. Completed hypotheses are collected separately so the live beam stays full until `width` have finished (early stopping). Final scores are length-normalized by the Wu 2016 penalty lp(n)=((5+n)/6)^alpha (alpha=0 → raw log-prob sum). Returns completed hypotheses, best score first (at most `width`).

Example

Beam search keeps the `width` best partial hypotheses and returns the highest-total-log-probability sequence — often better than greedy, which can be trapped by a locally-best first token. Here a tiny model (next-token logits depend on the last token) makes greedy pick [0,0,…] while beam width 2 finds the higher-scoring [0,1,0]. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	M := [][]float64{{2.0, 1.9}, {5.0, 0.0}}
	next := func(prefix []int) []float64 { return M[prefix[len(prefix)-1]] }

	beams := nlp.BeamSearch(next, []int{0}, 2 /*width*/, 2 /*maxNew*/, -1 /*no eos*/, 0 /*no length penalty*/)
	fmt.Println(beams[0].Tokens)
}

Output:
[0 1 0]

func DiverseBeamSearch 

func DiverseBeamSearch(next NextLogits, start []int, width, groups, maxNew, eos int, alpha, lambda float64) ([]Beam, error)

DiverseBeamSearch decodes width beams in `groups` groups (width must be divisible by groups), up to maxNew new tokens, completing a hypothesis on `eos` (eos<0 disables). lambda is the diversity strength and alpha the length penalty. It returns the completed/decoded hypotheses, best raw length-normalized log-prob first (at most `width`).

Example 
package main

import (
	"fmt"
	"slices"

	"github.com/jxsl13/goai/nlp"
)

// constNext returns the same logits at every step (independent of the prefix).
func constNext(logits []float64) nlp.NextLogits {
	return func([]int) []float64 { return append([]float64(nil), logits...) }
}

// lastTokens returns the final token of each beam.
func lastTokens(beams []nlp.Beam) []int {
	out := make([]int, len(beams))
	for i, b := range beams {
		out[i] = b.Tokens[len(b.Tokens)-1]
	}
	slices.Sort(out)
	return out
}

func main() {
	// Three groups, one beam each, strong diversity: the groups return the three most likely
	// yet DISTINCT first tokens instead of three near-copies.
	next := constNext([]float64{0, -1, -2, -3})
	beams, _ := nlp.DiverseBeamSearch(next, []int{7}, 3, 3, 1, -1, 0, 100)
	fmt.Println(lastTokens(beams))
}

Output:
[0 1 2]

type Block 

type Block struct {
	LN1, LN2 *nn.LayerNorm  // LN1 pre-attention, LN2 pre-MLP LayerNorm
	Attn     *MHA           // multi-head self-attention
	W1, B1   *tensor.Tensor // FFN up: [dim, 4dim], [4dim]
	W2, B2   *tensor.Tensor // FFN down: [4dim, dim], [dim]
}

Block is one pre-LN transformer block.

type FIMSentinels 

type FIMSentinels struct {
	Pre int // <PRE> sentinel, before the prefix
	Suf int // <SUF> sentinel, before the suffix
	Mid int // <MID> sentinel, before the middle
}

FIMSentinels are the special (reserved, non-document) token ids that mark the FIM sections.

type GPT 

type GPT struct {
	Config GPTConfig      // the model hyperparameters
	TokEmb *tensor.Tensor // [vocab, dim]
	PosEmb *tensor.Tensor // [ctx, dim]
	Blocks []*Block       // the stacked transformer blocks
	LNf    *nn.LayerNorm  // final pre-logits LayerNorm
	Head   *tensor.Tensor // [dim, vocab]
}

GPT is a decoder-only transformer for inference (§T23), pre-LN architecture (GPT-2 style): 

x = tokEmb[tokens] + posEmb[:seq]
per block: x += CausalMHA(LN1(x));  x += FFN(LN2(x)) with GELU
logits = LNf(x) · Head

func FromSafetensors 

func FromSafetensors(cfg GPTConfig, ts map[string]*tensor.Tensor) (*GPT, error)

FromSafetensors assembles a GPT from a tensor map using the goai naming convention: tok_emb, pos_emb, blocks.{i}.{ln1,ln2}.{gamma,beta}, blocks.{i}.attn.{wq,wk,wv,wo}, blocks.{i}.ffn.{w1,b1,w2,b2}, lnf.{gamma,beta}, head. Weight layout is [in,out] (§B19: transpose torch [out,in] on import).

func GPT2FromHF 

func GPT2FromHF(ts map[string]*tensor.Tensor, heads int) (*GPT, error)

GPT2FromHF builds a GPT from a HuggingFace GPT-2 checkpoint's tensor map (§T506 — possible since §T505 made the biased attention projections representable). It reads the HF naming (wte/wpe, h.N.ln_1, h.N.attn.c_attn …) and performs the two structural conversions GPT-2 needs:

  • the fused c_attn [d, 3d] (+bias [3d]) is SPLIT column-wise into Wq/Wk/Wv and their biases — HF's Conv1D stores weights [in, out], matching this library's layout, so no transposes are needed for the projections;
  • the LM head is TIED to the token embedding: head = wteᵀ.

Geometry (vocab, context, dim, layers) is inferred from the tensor shapes; heads must be supplied (GPT-2: 12 for 124M). Eps is GPT-2's 1e-5.

func (*GPT) DecodeStep 

func (g *GPT) DecodeStep(ctx *backend.Context, cache *KVCache, token, pos int) (*tensor.Tensor, error)

DecodeStep advances the model by one token using the KV-cache and returns the next-token logits [1,vocab]. pos is the token's absolute position (== cache length before the call). Inference-only (no tape).

func (*GPT) Embed 

func (g *GPT) Embed(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

Embed gathers token+position embeddings for the prompt through the dispatch, so gradients flow to TokEmb/PosEmb (differentiable, §T34): x = Embed(TokEmb, tokens) + Embed(PosEmb, 0..seq−1).

func (*GPT) Forward 

func (g *GPT) Forward(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

Forward computes logits [seq, vocab] for the prompt tokens.

func (*GPT) ForwardEarlyExit 

func (g *GPT) ForwardEarlyExit(ctx *backend.Context, tokens []int, layers []int) (mature *tensor.Tensor, premature []*tensor.Tensor, err error)

ForwardEarlyExit runs the model forward like Forward and ADDITIONALLY early-exits the hidden state after each requested block through the shared final LayerNorm and LM head (§T442, DoLa's premature logits): premature[i] = head(LNf(h_{layers[i]})) with h_j the residual stream after block j (0-based). The mature return is exactly Forward's logits [seq,vocab]; requesting the last block yields logits identical to mature. Block indices must be strictly increasing and in range.

func (*GPT) ForwardFromEmbed 

func (g *GPT) ForwardFromEmbed(ctx *backend.Context, x *tensor.Tensor) (*tensor.Tensor, error)

ForwardFromEmbed runs the transformer blocks and LM head on a precomputed embedding x [seq, dim], returning logits [seq, vocab]. Splitting the embedding step out lets a training loop inject NEFTune noise (nn.NEFTune) between the embedding and the blocks.

func (*GPT) ForwardHidden 

func (g *GPT) ForwardHidden(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

ForwardHidden returns the final hidden states [seq, dim] — the residual stream after all blocks and the final LayerNorm, i.e. Forward WITHOUT the LM head (§T443). This is the representation auxiliary decoding heads attach to (Medusa, early-exit probes): Forward(tokens) ≡ ForwardHidden(tokens)·Head.

func (*GPT) Generate 

func (g *GPT) Generate(prompt []int, maxNew int, s TokenSampler, opts ...GenerateOption) ([]int, error)

Generate autoregressively produces up to maxNew tokens after the prompt, using the KV-cache. Returns prompt + generated tokens. Stops at the context limit. By default the decode runs on backend.Default(); WithBackend overrides it (single-token decode is often faster on the CPU for small models, §T361).

func (*GPT) JacobiGenerate 

func (g *GPT) JacobiGenerate(prompt []int, genLen, maxIters int) ([]int, int, error)

JacobiGenerate runs Jacobi parallel decoding on the model itself (§T441, the generation entry point around JacobiDecode): each iteration is ONE full forward over prompt+guess (the model's parallel greedy pass — argmax per position), so the result is exactly the model's sequential greedy generation, reached in iters ≤ genLen parallel passes (often far fewer once tokens lock from the front). Returns prompt+generated, the iteration count, and the model's first forward error if one occurred. genLen is capped by the context window.

func (*GPT) NewCache 

func (g *GPT) NewCache() *KVCache

NewCache returns an empty cache sized for this model's blocks.

func (*GPT) Params 

func (g *GPT) Params() []*tensor.Tensor

Params returns every trainable tensor (token/pos embeddings, all block weights, final LN, and the LM head) for optimizers.

func (*GPT) Safetensors 

func (g *GPT) Safetensors() map[string]*tensor.Tensor

Safetensors returns the model's parameters under the FromSafetensors naming convention — the exact inverse of FromSafetensors — so a trained model can be checkpointed with safetensors.Save/SaveFile and reloaded bit-identically. The map holds the model's LIVE tensors (no copy): serialize before mutating.

type GPTConfig 

type GPTConfig struct {
	Vocab  int     // vocabulary size
	Ctx    int     // max context length
	Dim    int     // embedding width
	Heads  int     // number of attention heads
	Layers int     // number of transformer blocks
	Eps    float64 // LayerNorm epsilon
}

GPTConfig fixes the model geometry.

type GenerateOption 

type GenerateOption func(*genConfig)

GenerateOption configures a Generate call (functional options, §C12).

func WithBackend 

func WithBackend(be backend.Backend) GenerateOption

WithBackend runs the decode loop on the given backend instead of the default. Single-token decode is dispatch-latency-bound, so for SMALL models the CPU is faster than the GPU (measured ~2.7× at dim 512), while large models still favour the GPU — the caller, who knows the model size and hardware, chooses (§T361).

type JacobiStep 

type JacobiStep func(seq []int) []int

JacobiStep is a model's parallel greedy pass: given the full token sequence it returns, for every position j, the argmax next token the model predicts from the prefix seq[:j+1] (len(result) == len(seq)). For a causal LM this is a single forward over the whole sequence — the operation Jacobi decoding batches across all positions at once.

type KVCache 

type KVCache struct {
	K, V []*tensor.Tensor // per block; nil until the first token
}

KVCache holds the per-layer key/value tensors accumulated during autoregressive decoding (§T35), so each new token attends to cached past keys/values instead of recomputing attention over the whole prefix.

func (*KVCache) EvictStreaming 

func (c *KVCache) EvictStreaming(sink, recent int)

EvictStreaming bounds every block's K/V cache to the StreamingLLM sink+recent window in place (Xiao et al. 2023, §R73). Uniform across layers — the retained token positions are the same for every block — so a single index list applies to all. A no-op while the cache still fits in sink+recent.

func (*KVCache) Len 

func (c *KVCache) Len() int

Len returns the number of tokens currently cached.

type Llama 

type Llama struct {
	Config LlamaConfig    // the model hyperparameters
	TokEmb *tensor.Tensor // [vocab, dim] token embedding (no positional embedding)
	Blocks []*LlamaBlock  // the stacked transformer blocks
	Norm   *nn.RMSNorm    // final pre-logits RMSNorm
	Out    *tensor.Tensor // [dim, vocab] output projection (untied)
}

Llama is the LLaMA / Llama-2 decoder-only transformer (Touvron et al. 2023, arXiv:2302.13971 / 2307.09288, §R92) assembled from GoAI's verified primitives. It is the modern-LLM counterpart to the GPT-2-style GPT: pre-normalization with RMSNorm, rotary position embeddings (RoPE) on the queries and keys, grouped-query attention, a SwiGLU feed-forward, and NO biases or learned positional embeddings. Per block (pre-norm residual): 

h = x + Wo·Attn( RoPE(x̄·Wq), RoPE(x̄·Wk), x̄·Wv )   with x̄ = RMSNorm(x)
x = h + SwiGLU( RMSNorm(h) )

then a final RMSNorm and the untied output projection produce the logits.

Example

A Llama model assembled from RMSNorm, RoPE attention, GQA and SwiGLU turns a prompt into next-token logits. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/backend"
	"github.com/jxsl13/goai/nlp"

	_ "github.com/jxsl13/goai/backend/cpu"
	_ "github.com/jxsl13/goai/backend/ref"
)

func main() {
	m, _ := nlp.NewLlama(nlp.LlamaConfig{
		Vocab: 32, Ctx: 16, Dim: 16, Heads: 4, KVHeads: 2, Layers: 2, Hidden: 32, Eps: 1e-5,
	}, 1)
	logits, _ := m.Forward(backend.NewContext(), []int{5, 9, 14})
	fmt.Println(logits.Shape())
}

Output:
(3, 32)

func LlamaFromGGUF 

func LlamaFromGGUF(meta map[string]any, tensors map[string]*tensor.Tensor) (*Llama, error)

LlamaFromGGUF builds a Llama from the metadata and (dequantized) tensor maps of a parsed GGUF model file (gguf.File.Metadata / .Tensors), reading the ggml/llama.cpp convention (§R93): the config from the llama.* metadata keys and the weights from the token_embd / blk.N.* / output tensors. GGUF stores linear weights in torch [out, in] layout, so every projection is TRANSPOSED into GoAI's [in, out]; the embedding and RMSNorm gains are copied as-is. An absent output.weight ties the LM head to token_embd.

NOTE (§R93): the HF→GGUF conversion permutes the attn_q/attn_k rows into GGUF's split-half rotary pairing, and a loader reading GGUF must NOT re-permute — which is exactly what this loader does (it copies q/k as stored). GoAI's RoPE is the matching split-half convention, so the layout is consistent with upstream GGUF files; model weights round-trip exactly through LlamaToGGUF→LlamaFromGGUF (§V15). Exact bit-parity against a real llama.cpp-produced file is not verified on this host (llama.cpp is not installed, §B23).

Example

A Llama model saved to the GGUF metadata/tensor maps loads back with the same geometry, ready to run. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	m, _ := nlp.NewLlama(nlp.LlamaConfig{
		Vocab: 32, Ctx: 16, Dim: 16, Heads: 4, KVHeads: 2, Layers: 1, Hidden: 32, Eps: 1e-5,
	}, 1)
	meta, ts := nlp.LlamaToGGUF(m)
	back, _ := nlp.LlamaFromGGUF(meta, ts)
	fmt.Println(back.Config.Dim, back.Config.Heads, back.Config.KVHeads, len(back.Blocks))
}

Output:
16 4 2 1

func NewLlama 

func NewLlama(cfg LlamaConfig, seed uint64) (*Llama, error)

NewLlama builds a Llama with randomly initialized weights (Xavier for the projections, ones for the RMSNorm gains), for training or as a target to load weights into. It errors on an inconsistent geometry.

func (*Llama) DecodeStep 

func (m *Llama) DecodeStep(ctx *backend.Context, cache *LlamaCache, token, pos int) (*tensor.Tensor, error)

DecodeStep advances the Llama by one token using the KV-cache and returns the next-token logits [1,vocab]. pos is the token's absolute position (== cache length before the call), used as the RoPE offset. The token's rotated k,v are appended to the cache and the single query attends to all cached keys. Inference-only (no tape); it produces the same logits as a full Forward over the prefix.

func (*Llama) Embed 

func (m *Llama) Embed(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

Embed gathers the token embeddings for the prompt (Llama has no positional embedding — position enters through RoPE inside attention).

func (*Llama) Forward 

func (m *Llama) Forward(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

Forward computes logits [seq, vocab] for the prompt tokens.

func (*Llama) ForwardFromEmbed 

func (m *Llama) ForwardFromEmbed(ctx *backend.Context, x *tensor.Tensor) (*tensor.Tensor, error)

ForwardFromEmbed runs the transformer blocks, final RMSNorm and output projection on a precomputed embedding x [seq, dim], returning logits [seq, vocab]. Splitting the embedding step out lets a training loop inject NEFTune noise (nn.NEFTune) between the embedding and the blocks.

func (*Llama) ForwardHidden 

func (m *Llama) ForwardHidden(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

ForwardHidden returns the final hidden states [seq, dim] — the residual stream after all blocks and the final RMSNorm, i.e. Forward WITHOUT the output projection (§T447, the Llama sibling of GPT.ForwardHidden). This is the representation auxiliary decoding heads attach to (Medusa): Forward(tokens) ≡ ForwardHidden(tokens)·Out.

func (*Llama) Generate 

func (m *Llama) Generate(prompt []int, maxNew int, s TokenSampler, opts ...GenerateOption) ([]int, error)

Generate autoregressively decodes up to maxNew tokens after prompt with the sampler s, using the KV-cache (one forward per new token). Returns prompt+generated. With a greedy sampler the output is identical to argmax-ing a full Forward at each step. The decode runs on backend.Default() unless WithBackend overrides it (§T361).

func (*Llama) NewCache 

func (m *Llama) NewCache() *LlamaCache

NewCache returns an empty KV-cache sized for this model's blocks.

func (*Llama) NewStreamCache 

func (m *Llama) NewStreamCache() *StreamCache

NewStreamCache returns an empty StreamingLLM cache sized for this model's blocks.

func (*Llama) Params 

func (m *Llama) Params() []*tensor.Tensor

Params returns every trainable tensor for optimizers.

func (*Llama) SelfExtendForward 

func (m *Llama) SelfExtendForward(ctx *backend.Context, tokens []int, window, group int) (*tensor.Tensor, error)

SelfExtendForward runs the Llama forward with Self-Extend grouped attention (§T513): every block's attention merges TWO score sources under ONE softmax via OpMHASelect — source 1 is the ordinary RoPE path (true positions, used for pairs within the neighbor window q−k < window), source 2 rotates queries at ⌊q/G⌋ + window − ⌊window/G⌋ and keys at ⌊k/G⌋ (used beyond the window), so distant relative positions compress into the trained range (SelfExtendRelPos). group=1 makes both sources identical and collapses onto Forward. Analysis-scale and inference-only (OpMHASelect has no VJP); returns logits [seq, vocab].

func (*Llama) SelfExtendGenerate 

func (m *Llama) SelfExtendGenerate(ctx *backend.Context, prompt []int, maxNew, window, group int) ([]int, error)

SelfExtendGenerate greedily generates maxNew tokens with Self-Extend grouped attention (§T541): every step re-runs SelfExtendForward over the whole sequence and appends the argmax token — the generation companion of the teacher-forced evaluations (§T513/§T527), letting a model trained on short windows generate with FULL attention far beyond its training length (unlike StreamGenerate's bounded sliding-window cache, distant tokens stay visible through the grouped source). Analysis-scale: O(seq²) attention per step, no KV cache (OpMHASelect is inference-only and served by the reference kernel). group=1 degenerates to plain greedy generation.

func (*Llama) StreamGenerate 

func (m *Llama) StreamGenerate(prompt []int, maxNew, sinks, window int, s TokenSampler) ([]int, error)

StreamGenerate decodes up to maxNew tokens after prompt with StreamingLLM's bounded cache (sinks attention-sink tokens + a rolling window), so it runs at constant memory regardless of how long the stream grows. Unlike KV-cache Generate it is not bounded by the model's context length. Returns prompt+generated.

func (*Llama) StreamStep 

func (m *Llama) StreamStep(ctx *backend.Context, cache *StreamCache, token, sinks, window int) (*tensor.Tensor, error)

StreamStep advances the model by one token with the bounded StreamingLLM cache and returns the next-token logits [1,vocab]. It keeps the first `sinks` "attention sink" tokens (default 4 in the paper) and a rolling window of the last `window` tokens, evicting the middle. Keys/values are cached pre-RoPE; each step re-applies RoPE using positions within the current cache (0..cacheLen−1), so relative distances stay bounded by sinks+window and generation never runs out of positional range. Inference-only.

type LlamaBlock 

type LlamaBlock struct {
	AttnNorm       *nn.RMSNorm    // RMSNorm before attention
	Wq, Wk, Wv, Wo *tensor.Tensor // attention projections (no bias); Wk/Wv are [dim, KVHeads·headDim]
	FFNNorm        *nn.RMSNorm    // RMSNorm before the FFN
	FFN            *nn.SwiGLU     // SwiGLU feed-forward
}

LlamaBlock is one pre-norm Llama transformer block.

type LlamaCache 

type LlamaCache struct {
	K, V []*tensor.Tensor // per block; nil until the first token
}

LlamaCache holds the per-layer key/value tensors accumulated during autoregressive Llama decoding, so each new token attends to the cached past instead of recomputing attention over the whole prefix. The cached keys already carry their RoPE rotation (applied at the position each token entered the cache).

func (*LlamaCache) Len 

func (c *LlamaCache) Len() int

Len returns the number of tokens currently cached.

type LlamaConfig 

type LlamaConfig struct {
	Vocab    int     // vocabulary size
	Ctx      int     // max context length
	Dim      int     // embedding width (Heads·headDim)
	Heads    int     // number of query heads
	KVHeads  int     // key/value heads (GQA); 0 → Heads (standard MHA)
	Layers   int     // number of transformer blocks
	Hidden   int     // SwiGLU inner width (Llama uses ≈ (2/3)·4·Dim rounded)
	Eps      float64 // RMSNorm epsilon
	RopeBase float64 // RoPE frequency base θ; 0 → 10000
}

LlamaConfig fixes the model geometry.

type MHA 

type MHA struct {
	Heads          int            // number of attention heads
	Wq, Wk, Wv, Wo *tensor.Tensor // each [dmodel, dmodel]
	// Causal masks position i from attending to j > i (decoder self-attention):
	// masked scores are −∞ before softmax, so their weight is exactly 0.
	Causal bool
	// LoRA holds optional low-rank adapters for the projections, keyed "q","k","v","o"
	// (§T504, attach with nlp.ApplyLoRAGPT or manually). Each adapter wraps the SAME
	// frozen weight tensor; with LoRA's zero-initialized B the forward is numerically
	// identical to the plain projection until the adapter trains. nil = plain paths.
	LoRA map[string]*nn.LoRALinear
	// Bias holds optional per-projection biases [dmodel], keyed "q","k","v","o"
	// (§T505): q = x·Wq + Bias["q"], etc. GPT-2-family checkpoints carry these
	// (c_attn.bias, c_proj.bias); Llama-family models do not. nil map or missing
	// key = no bias, the exact pre-§T505 path.
	Bias map[string]*tensor.Tensor
	// Mask, when set, replaces the fused causal attention with OpMHAMasked using
	// this [seq,seq] additive mask (§T508 — tree attention for Medusa tree
	// verification; −Inf excludes a pair). INFERENCE-ONLY (the masked op has no
	// VJP) and per-call state: set it, run Forward, clear it. nil = the exact
	// fused path.
	Mask *tensor.Tensor
}

MHA is multi-head scaled dot-product attention (Vaswani et al. 2017, "Attention Is All You Need") for a single sequence x[seq, dmodel]: 

Q = x·Wq; K = x·Wk; V = x·Wv                     (weights [dmodel,dmodel])
per head h: Aₕ = softmax(QₕKₕᵀ/√dₖ)·Vₕ            (dₖ = dmodel/heads)
out = concat(A₀..A_{h−1})·Wo

Head splits are zero-copy Slice views; matmul/softmax run through the backend dispatch, so accel backends (metal GEMM) apply automatically. Inference-focused (§T21); training VJPs for layernorm pending §B22.

func NewMHA 

func NewMHA(heads int, wq, wk, wv, wo *tensor.Tensor) (*MHA, error)

NewMHA builds an MHA block with the given weights. dmodel must divide by heads.

func (*MHA) Forward 

func (m *MHA) Forward(ctx *backend.Context, x *tensor.Tensor) (*tensor.Tensor, error)

Forward computes attention for x[seq, dmodel]. Fully differentiable: the projections are matmuls and the multi-head scaled-dot-product core is the single fused OpMHA op (head split/concat/mask internal, §T32) — so gradients flow to Wq/Wk/Wv/Wo and x through the standard dispatch, with no view ops on the tape.

func (*MHA) Params 

func (m *MHA) Params() []*tensor.Tensor

Params returns the projection weights.

func (*MHA) StepKV 

func (m *MHA) StepKV(ctx *backend.Context, h, kc, vc *tensor.Tensor) (out, kNew, vNew *tensor.Tensor, err error)

StepKV runs attention for a single new token h[1,dmodel] given cached K,V (each [t,dmodel] or nil). It appends the token's k,v to the cache and returns (out[1,dmodel], Knew, Vnew). The single query attends to all t+1 keys.

type MedusaHeads 

type MedusaHeads struct {
	W []*tensor.Tensor // head k: [dim, vocab]
}

MedusaHeads are K trainable decoding heads over the base model's final hidden state (§T443): head k is a linear map [dim,vocab] predicting the token at offset t+2+k from hidden state h_t (the base LM head already covers t+1). This is the simplest head form — the reference implementation uses a small residual block per head; a linear head is its first-order variant and trains with the library's own loop (frozen base: compute ForwardHidden outside the tape, then tape over Logits).

Example

MedusaHeads are extra decoding heads over a frozen base model's final hidden state (ForwardHidden): head k predicts the token at offset t+2+k, enabling multi-token drafting. Here two heads project a hidden batch into per-head logits. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/backend"
	"github.com/jxsl13/goai/nlp"
	"github.com/jxsl13/goai/tensor"
)

func main() {
	heads, err := nlp.NewMedusaHeads(2 /*K*/, 4 /*dim*/, 6 /*vocab*/, 1 /*seed*/)
	if err != nil {
		panic(err)
	}
	hidden := tensor.New(tensor.F32, tensor.Shape{3, 4}) // e.g. gpt.ForwardHidden(ctx, tokens)
	logits, err := heads.Logits(backend.NewContext(), hidden)
	if err != nil {
		panic(err)
	}
	fmt.Println(len(logits), logits[0].Shape(), logits[1].Shape())
}

Output:
2 (3, 6) (3, 6)

func NewMedusaHeads 

func NewMedusaHeads(k, dim, vocab int, seed uint64, opts ...MedusaHeadsOption) (*MedusaHeads, error)

NewMedusaHeads builds K linear heads with a small deterministic random init.

func (*MedusaHeads) Logits 

func (m *MedusaHeads) Logits(ctx *backend.Context, hidden *tensor.Tensor) ([]*tensor.Tensor, error)

Logits projects the hidden states [seq,dim] through every head, returning K logit tensors [seq,vocab] (differentiable — training a frozen-base Medusa is a tape over this call with the base's ForwardHidden output as constant input).

func (*MedusaHeads) Params 

func (m *MedusaHeads) Params() []*tensor.Tensor

Params returns the head weights for optimizers.

type MedusaHeadsOption 

type MedusaHeadsOption func(*medusaHeadsCfg)

MedusaHeadsOption configures NewMedusaHeads (functional options, §C12).

func WithMedusaHeadsDtype 

func WithMedusaHeadsDtype(d tensor.Dtype) MedusaHeadsOption

WithMedusaHeadsDtype sets the head weight dtype (default F32); match the base model's hidden-state dtype.

type Mirostat 

type Mirostat struct {
	Tau float64 // target surprise τ in bits (per-token cross-entropy; default 5.0)
	Eta float64 // feedback learning rate η (default 0.1)
	Mu  float64 // running surprise threshold μ (state; initialized to 2·τ)
	// contains filtered or unexported fields
}

Mirostat is the Mirostat 2.0 decoding algorithm (Basu et al. 2020, "Mirostat: A Neural Text Decoding Algorithm that Directly Controls Perplexity", arXiv:2007.14966, ICLR 2021; the mirostat_v2 variant used in llama.cpp, §R85). Unlike top-k/top-p, which fix a truncation and let perplexity drift, Mirostat targets a fixed SURPRISE τ (the per-token cross-entropy, in bits) with a feedback loop: it truncates to the tokens whose surprise is within a running threshold μ, samples, measures the actual surprise of the chosen token, and nudges μ by the error — so the long-run average surprise (hence perplexity 2^τ) is held at the target regardless of context.

Per step, with token surprise S(i) = −log₂ p(i): 

truncate to candidates with S(i) ≤ μ   (always keep ≥ 1 — the top token)
sample X from the renormalized candidates
e = S(X) − τ ;  μ ← μ − η·e            (S(X) from the ORIGINAL prob of X)

μ starts at 2·τ (Algorithm 2). It is mutable state carried across calls within one generated sequence; call Reset (or a fresh Mirostat) per new sequence.

Example

Mirostat holds the average per-token surprise at the target τ. With τ set to 0 the threshold μ starts below every token's surprise, so it keeps only the most probable token — a deterministic arg-max here selecting token 1. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	m := nlp.NewMirostat(1, nlp.WithMirostatTau(0))
	fmt.Println(m.Sample([]float64{0.2, 2.5, 1.0, 0.3}))
}

Output:
1

func NewMirostat 

func NewMirostat(seed uint64, opts ...MirostatOption) *Mirostat

NewMirostat builds a deterministic Mirostat 2.0 sampler seeded by seed, with the canonical defaults τ=5.0 bits and η=0.1, and μ initialized to 2·τ.

func (*Mirostat) Reset 

func (m *Mirostat) Reset()

Reset restores the surprise threshold μ to 2·τ, for reuse across sequences.

func (*Mirostat) Sample 

func (m *Mirostat) Sample(logits []float64) int

Sample returns a token id for the given logits and advances the surprise threshold μ by the feedback rule. It softmaxes the logits, truncates to the tokens whose surprise −log₂p ≤ μ (keeping at least the most probable token), samples one from the renormalized survivors, then updates μ ← μ − η·(S(X) − τ).

func (*Mirostat) SampleWithHistory 

func (m *Mirostat) SampleWithHistory(logits []float64, _ []int) int

SampleWithHistory satisfies TokenSampler; Mirostat ignores the history — it controls repetition through its surprise target rather than explicit penalties (a Mirostat run whose μ has adapted never collapses to a greedy loop the way untruncated small-model sampling can).

type MirostatOption 

type MirostatOption func(*Mirostat)

MirostatOption configures a Mirostat sampler (functional-options idiom, §C12).

func WithMirostatEta 

func WithMirostatEta(eta float64) MirostatOption

WithMirostatEta sets the feedback learning rate η (default 0.1).

func WithMirostatTau 

func WithMirostatTau(tau float64) MirostatOption

WithMirostatTau sets the target surprise τ in bits (default 5.0); the controlled perplexity is 2^τ.

type NextLogits 

type NextLogits func(prefix []int) []float64

NextLogits returns the next-token logits given the tokens produced so far. Beam search applies a stable log-softmax internally, so raw logits are fine.

type QuantBlock 

type QuantBlock struct {
	AttnNorm       *nn.RMSNorm     // RMSNorm before attention (f32 gain)
	Wq, Wk, Wv, Wo *nn.QuantLinear // quantized attention projections (no bias)
	FFNNorm        *nn.RMSNorm     // RMSNorm before the FFN (f32 gain)
	FFN            *nn.QuantSwiGLU // quantized SwiGLU feed-forward
}

QuantBlock is a QuantLlama transformer block: float RMSNorm gains, quantized attention projections and a quantized SwiGLU FFN.

type QuantKVCache 

type QuantKVCache struct {
	// contains filtered or unexported fields
}

QuantKVCache is a KV cache that stores each token's key and value row in 8-bit Q8_0 block format (§R108) instead of f32, cutting long-context KV memory ~3.76× (34 bytes per 32-element block = 1.0625 B/elem vs 4 B/elem) — the dominant memory cost of long-context inference. Each appended row is quantized independently along its contiguous dimension with the same per-block Q8_0 layout used for weights (§R94, already tier-2 verified): per 32 elements one f16 scale d=amax/127 and 32 int8 quants. Keys/Values dequantize the whole store back to f32 [t,dim] tensors for attention. 8-bit per-block symmetric quantization is the near-lossless tier for KV caching — the per-channel-Key/per-token-Value asymmetry (Hooper et al. 2024, KVQuant) only matters at 4/3/2-bit, not at 8-bit (§R108). The row width dim must be a multiple of 32.

Example

A Q8_0 KV cache stores each token's key/value row in 8-bit blocks instead of f32, cutting long-context memory ~3.76×. Here two tokens of width 64 are cached: their quantized footprint is 2 tokens × (64/32 blocks) × 34 bytes × 2 (K+V) = 272 bytes, versus 2×2×64×4 = 1024 for f32. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
	"github.com/jxsl13/goai/tensor"
)

func main() {
	c, _ := nlp.NewQuantKVCache(64)
	k := tensor.New(tensor.F32, tensor.Shape{64})
	v := tensor.New(tensor.F32, tensor.Shape{64})
	_ = c.Append(k, v)
	_ = c.Append(k, v)
	fmt.Printf("%d tokens, %d bytes\n", c.Len(), c.Bytes())
}

Output:
2 tokens, 272 bytes

func NewQuantKVCache 

func NewQuantKVCache(dim int) (*QuantKVCache, error)

NewQuantKVCache creates an empty Q8_0 KV cache for rows of width dim (must be a multiple of the 32-element Q8_0 block size).

func (*QuantKVCache) Append 

func (c *QuantKVCache) Append(k, v *tensor.Tensor) error

Append quantizes one token's key and value rows (each [dim] or [1,dim]) to Q8_0 and stores them. k and v are consumed as flat dim-vectors.

func (*QuantKVCache) Bytes 

func (c *QuantKVCache) Bytes() int

Bytes returns the total quantized storage (keys + values) in bytes — the memory footprint that replaces 2·t·dim·4 bytes of f32 cache.

func (*QuantKVCache) Keys 

func (c *QuantKVCache) Keys() (*tensor.Tensor, error)

Keys dequantizes the cached keys to a fresh f32 tensor [t,dim] (t = Len). Values is its twin.

func (*QuantKVCache) Len 

func (c *QuantKVCache) Len() int

Len returns the number of cached tokens.

func (*QuantKVCache) Values 

func (c *QuantKVCache) Values() (*tensor.Tensor, error)

Values dequantizes the cached values to a fresh f32 tensor [t,dim].

type QuantLlama 

type QuantLlama struct {
	Config LlamaConfig     // model hyperparameters
	TokEmb *tensor.Tensor  // [vocab, dim] f32 token embedding
	Blocks []*QuantBlock   // stacked transformer blocks with quantized projections
	Norm   *nn.RMSNorm     // final pre-logits RMSNorm (f32 gain)
	Out    *nn.QuantLinear // quantized output projection (In=dim, Out=vocab)
}

QuantLlama is a Llama that keeps its projection weights QUANTIZED (nn.QuantLinear) and never materializes them as full-precision matrices — the memory-efficient, GPU-accelerated form of a quantized model (§T149). Its forward is identical to Llama.ForwardFromEmbed except every linear projection (attn q/k/v/o, the FFN, the output head) is an in-kernel dequantized matmul that runs on the active accelerator when it accelerates the quant type (else the CPU fallback). Everything runs in f32 — quantized-inference activations — which is also what lets the GPU ops engage (the accelerators are f32-only). Inference-only: the quantized weights are frozen and bypass the tape.

Example

A Llama can be quantized to run its linear layers on the GPU with the weights kept in quantized byte form — the same forward, a fraction of the weight memory, near-identical logits. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/backend"
	"github.com/jxsl13/goai/format/gguf"
	"github.com/jxsl13/goai/nlp"
)

func main() {
	m, _ := nlp.NewLlama(nlp.LlamaConfig{
		Vocab: 12, Ctx: 16, Dim: 32, Heads: 4, KVHeads: 2, Layers: 1, Hidden: 32, Eps: 1e-5, RopeBase: 10000,
	}, 3)
	q, _ := nlp.QuantizeLlama(m, gguf.Q8_0)
	logits, _ := q.Forward(backend.NewContext(), []int{1, 2, 3})
	fmt.Println(logits.Shape(), "blocks:", len(q.Blocks))
}

Output:
(3, 12) blocks: 1

func QuantLlamaFromGGUF 

func QuantLlamaFromGGUF(meta map[string]any, tensors map[string]gguf.QuantTensor) (*QuantLlama, error)

QuantLlamaFromGGUF builds a QuantLlama from the metadata and STILL-QUANTIZED tensor map of a GGUF file (gguf.ReadRaw, §T151) — loading a real quantized model straight onto the GPU without ever materializing full-precision weights. GGUF stores linear weights in the same [out, in] block layout QuantLinear expects, so each projection's bytes are wrapped directly — NO transpose, NO re-quantization. Only the small, precision-sensitive pieces are dequantized to f32: the RMSNorm gains and the token embedding (its lookup needs a float table). An absent output.weight ties the LM head to token_embd (which must itself be quantized). Follows the ggml/llama.cpp convention (§R93), the quantized twin of LlamaFromGGUF.

func QuantizeLlama 

func QuantizeLlama(m *Llama, qt gguf.QuantType) (*QuantLlama, error)

QuantizeLlama builds a QuantLlama from a float Llama by quantizing every linear projection to qt (the ggml block layout) — the projections carry the bulk of the weights and compute, so this is where quantization pays off. RMSNorm gains and the token embedding are kept in f32 (they are tiny and precision-sensitive). Each projection's inner dimension must be a multiple of qt's block size (32 for Q8_0/Q4_0, 256 for the k-quants).

func (*QuantLlama) Close 

func (m *QuantLlama) Close() error

Close frees every device-resident weight buffer held by the model's quantized projections (attention, FFN, output head). Idempotent; call it when done with the model to release GPU memory promptly (otherwise the buffers are reclaimed only at process exit).

func (*QuantLlama) DecodeStep 

func (m *QuantLlama) DecodeStep(ctx *backend.Context, cache *LlamaCache, token, pos int) (*tensor.Tensor, error)

DecodeStep advances the quantized model by one token at absolute position pos, appending its post-RoPE K/V to the cache and returning the next-token logits [1, vocab]. It mirrors Llama.DecodeStep (§T121) exactly — quantized projections, RoPE at PosOffset=pos, single-query attention over the whole cache — so a KV-cache decode of a quantized model matches its full Forward up to f32 reassociation, all on the GPU.

func (*QuantLlama) Forward 

func (m *QuantLlama) Forward(ctx *backend.Context, tokens []int) (*tensor.Tensor, error)

Forward runs the quantized model on the token ids, returning logits [seq, vocab]. It mirrors Llama.ForwardFromEmbed exactly — RMSNorm, RoPE, MHA, residual adds and SwiGLU gating — but every projection is a quantized in-kernel matmul, all in f32.

func (*QuantLlama) Generate 

func (m *QuantLlama) Generate(prompt []int, maxNew int, s TokenSampler) ([]int, error)

Generate autoregressively decodes up to maxNew tokens after the prompt on the quantized model, using the KV-cache (each step is one token, not a full re-forward), and returns prompt+new. The sampler s selects each token (Greedy() for deterministic argmax). Stops at the context limit.

func (*QuantLlama) NewCache 

func (m *QuantLlama) NewCache() *LlamaCache

NewCache allocates an empty KV-cache for autoregressive decoding of this QuantLlama (one post-RoPE K/V slot per block, filled as tokens are decoded). Reuses LlamaCache — the cache structure is identical to the float model's; only the projections differ.

type RegexGuide 

type RegexGuide struct {
	// contains filtered or unexported fields
}

RegexGuide implements regex/FSM-guided constrained decoding (Willard & Louf 2023, §R111, "Efficient Guided Generation for Large Language Models", arXiv:2307.09702 — the method behind Outlines). A regular expression is compiled to a finite-state machine; during generation the next-token logits are MASKED so only tokens whose characters keep the FSM in a live (non-dead) state are allowed, and after a token is chosen the FSM state is ADVANCED by consuming that token's characters. Sampling may stop (EOS) only when the FSM is in an ACCEPTING state. By construction every produced string is in the regular language — the output is guaranteed to match the regex.

The FSM is the standard-library regexp NFA (regexp/syntax): a state is the ε-closed set of NFA program counters. Per the paper's efficiency contribution (§4), the state→token transitions are precomputed lazily and memoized, so per-step masking is an O(1) lookup rather than O(vocab) regex tests. States are interned to small integer ids.

Empty-width assertions (^, $, \b) are treated permissively (as ε) — guided generation over the whole output rarely needs them, and the common structured-output patterns (character classes, literals, repetition, alternation) do not.

Example

A RegexGuide constrains generation to a regular language: from the start of `[0-9]+`, digit tokens are allowed and others masked, and the state becomes accepting (generation may stop) once at least one digit has been emitted. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	vocab := []string{"3", "x", "."}
	g, _ := nlp.NewRegexGuide(`[0-9]+`, vocab)
	st := g.Start()
	fmt.Println(g.Allowed(st, 0), g.Allowed(st, 1), g.Accepting(st)) // "3" ok, "x" no, not yet accepting
	st = g.Advance(st, 0)                                            // consume "3"
	fmt.Println(g.Accepting(st))                                     // a complete match now
}

Output:
true false false
true

func NewRegexGuide 

func NewRegexGuide(pattern string, vocab []string) (*RegexGuide, error)

NewRegexGuide compiles pattern into a guide over vocab (token id → the token's character string). The pattern must match the ENTIRE generated string.

func (*RegexGuide) Accepting 

func (g *RegexGuide) Accepting(state int) bool

Accepting reports whether state is a match/accept state — the only states from which generation may terminate (EOS).

func (*RegexGuide) Advance 

func (g *RegexGuide) Advance(state, token int) int

Advance consumes the token's characters from state and returns the resulting state id, or -1 if the token is invalid (hits a dead state). Results are memoized per (state, token).

func (*RegexGuide) Allowed 

func (g *RegexGuide) Allowed(state, token int) bool

Allowed reports whether the token with the given id is valid from state (every character has a live transition). A dead state (-1) allows nothing.

func (*RegexGuide) MaskLogits 

func (g *RegexGuide) MaskLogits(state int, logits []float64, eosID int) (eosAllowed bool)

MaskLogits sets the logit of every token disallowed from state to −Inf, in place, and returns whether EOS (ending generation) is allowed — i.e. whether state is accepting. eosID, if in range, is force-allowed exactly when the state is accepting (and masked out otherwise). A dead state masks everything. logits must be vocab-sized.

func (*RegexGuide) Sampler 

func (g *RegexGuide) Sampler(inner TokenSampler, eosID int) TokenSampler

Sampler wraps inner into a TokenSampler that enforces the guide during generation (§T439): each draw masks the disallowed tokens to −Inf from the CURRENT FSM state (MaskLogits on a copy), lets inner pick, and advances the state by the picked token — so it plugs straight into any generation loop that takes a TokenSampler. eosID (−1 for none) is only ever allowed in accepting states; picking it stops advancing. The wrapper is stateful: use a fresh Sampler per generated sequence, and choose a pattern that always has a continuation (a dead end masks everything).

func (*RegexGuide) Start 

func (g *RegexGuide) Start() int

Start returns the guide's initial FSM state id.

type Sampler 

type Sampler struct {
	Temperature float64 // softmax temperature; <1 sharper, >1 flatter (default 1)
	TopK        int     // keep the K highest-prob tokens; 0 = off
	TopP        float64 // nucleus: keep smallest set with cumulative prob ≥ P; 0 = off
	MinP        float64 // keep tokens with prob ≥ MinP·maxProb; 0 = off
	Epsilon     float64 // epsilon sampling: keep tokens with prob ≥ Epsilon; 0 = off
	Eta         float64 // eta sampling: keep tokens with prob ≥ min(Eta, √Eta·exp(−H)); 0 = off
	Typical     float64 // locally typical: keep tokens whose surprisal is nearest the entropy until cum-prob ≥ τ; 0/≥1 = off

	// Repetition penalties (applied by SampleWithHistory over the recent history).
	RepeatPenalty   float64 // CTRL (Keskar et al. 2019): divide positive logits of seen tokens by this, multiply negative; 0/1 = off, typical 1.1–1.3
	FreqPenalty     float64 // subtract FreqPenalty · count(token in window) from its logit; 0 = off
	PresencePenalty float64 // subtract PresencePenalty once if the token appears in the window at all; 0 = off
	PenaltyLastN    int     // history window the penalties look at; 0 = the entire history
	// contains filtered or unexported fields
}

Sampler turns a logit vector into a token id. Pipeline (matching HuggingFace generation): repetition penalties (via SampleWithHistory) → temperature scaling → top-k mask → softmax → top-p (nucleus) mask → min-p mask → epsilon/eta mask → multinomial sample. Temperature ≤ 0 selects greedy (argmax), which ignores the truncation filters and RNG — but not the penalties, which act on the logits before the greedy/temperature split.

top-p follows Holtzman et al. 2019 (§R34): the nucleus is the smallest set of highest-probability tokens whose cumulative probability ≥ p; the token that crosses p is included, so at least the top token always survives. min-p follows Nguyen et al. 2024 (§R63): keep tokens with probability ≥ MinP·max-prob, an adaptive threshold that tightens when the model is confident. epsilon and eta sampling follow Hewitt et al. 2022 (§R91): an absolute floor, and an entropy-adaptive floor min(ε, √ε·exp(−H)). Kept probabilities are renormalized before sampling, and the arg-max token always survives every filter.

Example

A Sampler runs the standard temperature → top-k → top-p → min-p → multinomial pipeline. Here a sharply peaked distribution plus a top-p nucleus of 0.9 leaves only the dominant token inside the nucleus, so sampling is effectively pinned to it regardless of the seed — how nucleus sampling stays on-topic when the model is confident while still allowing variety when it is not. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	s := nlp.NewSampler(42, nlp.WithTopP(0.9))
	logits := []float64{6.0, 1.0, 0.5} // softmax ≈ 0.989 on token 0 → nucleus = {0}
	fmt.Println(s.Sample(logits))
}

Output:
0
Example (MinP)

ExampleSampler_minP shows min-p sampling via the functional-options constructor. A confident distribution (token 0 ≈ 0.83) with min-p 0.2 sets the keep-threshold to 0.2·0.83 ≈ 0.17, so only the top token survives and sampling is decisive. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	s := nlp.NewSampler(1, nlp.WithMinP(0.2))
	fmt.Println(s.Sample([]float64{3, 0, 0, 0}))
}

Output:
0

func Greedy 

func Greedy() *Sampler

Greedy returns a deterministic argmax sampler. It still carries an rng (unused by argmax Sample) so callers that draw from its distribution — e.g. speculative decoding's residual/bonus draws — never hit a nil rng; the greedy outcome stays deterministic regardless.

Example

Greedy decoding always picks the highest-scoring token — deterministic, no RNG. It is the simplest decoder and the baseline the others are measured against. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	logits := []float64{0.1, 2.5, 0.3} // token 1 has the largest score
	tok := nlp.Greedy().Sample(logits)
	fmt.Println(tok)
}

Output:
1

func NewSampler 

func NewSampler(seed uint64, opts ...SamplerOption) *Sampler

NewSampler builds a deterministic sampler seeded by seed. Defaults: temperature 1, no top-k / top-p / min-p — configure with the options, e.g. NewSampler(seed, WithTemperature(0.8), WithTopP(0.95)).

func (*Sampler) Dist 

func (s *Sampler) Dist(logits []float64) []float64

Dist returns the probability distribution this sampler draws from for the given logits: temperature scaling, then top-k, top-p (nucleus) and min-p filtering, renormalized to sum 1. For greedy sampling (Temperature ≤ 0) it is a one-hot vector at the arg-max. Speculative decoding (§R53) uses it to obtain the target and draft distributions p and q that the accept/reject rule compares.

func (*Sampler) Sample 

func (s *Sampler) Sample(logits []float64) int

Sample returns a token id for the given logits.

func (*Sampler) SampleWithHistory 

func (s *Sampler) SampleWithHistory(logits []float64, history []int) int

SampleWithHistory applies the sampler's repetition penalties (RepeatPenalty, FreqPenalty, PresencePenalty over the PenaltyLastN window) to the logits of tokens present in history, then samples like Sample. The penalties act before the greedy/temperature split, so they steer greedy decoding too — the classic cure for small-model repetition loops. history is the running sequence (prompt + generated), as maintained by the Generate loops. With no penalties configured it is exactly Sample. The speculative-decoding paths (SpeculativeGenerate, prompt-lookup) do NOT apply penalties: their lossless accept/reject math compares raw model distributions.

type SamplerOption 

type SamplerOption func(*Sampler)

SamplerOption configures a Sampler via the functional-options idiom (§C12).

func WithEpsilon 

func WithEpsilon(eps float64) SamplerOption

WithEpsilon enables epsilon sampling (Hewitt et al. 2022, §R91): keep only tokens with probability ≥ eps, an absolute floor. 0 disables it; typical 3e-4–2e-3.

func WithEta 

func WithEta(eps float64) SamplerOption

WithEta enables eta sampling (Hewitt et al. 2022, §R91): an entropy-adaptive threshold η = min(eps, √eps·exp(−H)) where H is the distribution's Shannon entropy in nats — it truncates hard when the model is confident (low entropy) and keeps more when uncertain. 0 disables it; typical eps≈9e-4.

Example

Eta sampling adapts its cutoff to the model's confidence. On a peaked distribution (token 0 at 0.9) the entropy is low, so η rises to the ε floor and the low-probability tail is dropped, leaving only the confident token. 

package main

import (
	"fmt"
	"math"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	s := nlp.NewSampler(1, nlp.WithEta(0.1))
	dist := s.Dist([]float64{math.Log(0.9), math.Log(0.08), math.Log(0.02)})
	fmt.Printf("%.2f %.2f %.2f\n", dist[0], dist[1], dist[2])
}

Output:
1.00 0.00 0.00

func WithFrequencyPenalty 

func WithFrequencyPenalty(f float64) SamplerOption

WithFrequencyPenalty subtracts f · count(token in the penalty window) from each seen token's logit (the OpenAI-style frequency penalty). 0 disables it.

func WithMinP 

func WithMinP(p float64) SamplerOption

WithMinP keeps tokens with probability ≥ p·max-prob (Nguyen et al. 2024); 0 disables it. Typical 0.05–0.1.

func WithPenaltyWindow 

func WithPenaltyWindow(n int) SamplerOption

WithPenaltyWindow limits the repetition penalties to the last n history tokens (llama.cpp's repeat_last_n). 0 (the default) penalizes over the entire history.

func WithPresencePenalty 

func WithPresencePenalty(p float64) SamplerOption

WithPresencePenalty subtracts p from the logit of every token that appears in the penalty window at all (the OpenAI-style presence penalty). 0 disables it.

func WithRepeatPenalty 

func WithRepeatPenalty(p float64) SamplerOption

WithRepeatPenalty enables the CTRL repetition penalty (Keskar et al. 2019): for every token in the penalty window, a positive logit is divided by p and a negative one multiplied by it. 0 or 1 disables it; typical 1.1–1.3.

func WithTemperature 

func WithTemperature(t float64) SamplerOption

WithTemperature sets the softmax temperature (default 1; ≤0 = greedy).

func WithTopK 

func WithTopK(k int) SamplerOption

WithTopK keeps only the k highest-logit tokens (0 = disabled).

func WithTopP 

func WithTopP(p float64) SamplerOption

WithTopP keeps the smallest nucleus with cumulative probability ≥ p (Holtzman et al. 2019); 0 or ≥1 disables it.

func WithTypical 

func WithTypical(tau float64) SamplerOption

WithTypical enables locally typical sampling (Meister, Pimentel, Wiher & Cotterell 2023, "Locally Typical Sampling"): keep the smallest set of tokens whose information content −log p is closest to the distribution's entropy H(p) until their cumulative probability reaches τ, filtering out both the too-predictable and the too-surprising. 0 or ≥1 disables it; typical τ 0.9–0.95.

Example 
package main

import (
	"fmt"
	"math"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	// p=[0.4,0.3,0.2,0.1]: locally typical (τ=0.7) drops the least-typical tail token.
	z := []float64{math.Log(0.4), math.Log(0.3), math.Log(0.2), math.Log(0.1)}
	d := nlp.NewSampler(1, nlp.WithTypical(0.7)).Dist(z)
	fmt.Printf("%.3f %.3f %.3f %.3f\n", d[0], d[1], d[2], d[3])
}

Output:
0.444 0.333 0.222 0.000

type SpecStats 

type SpecStats struct {
	Proposed int // total draft tokens proposed
	Accepted int // draft tokens accepted
}

SpecStats reports a speculative run: how many draft tokens were Proposed across all rounds and how many were Accepted (the residual/bonus token each round is not counted as an accept). A high AcceptanceRate is what turns the one-target-forward- per-round into an end-to-end speedup.

Example

SpecStats reports how many of the draft model's proposed tokens the target accepted; the acceptance rate is what turns one target forward per round into an end-to-end speedup. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	s := nlp.SpecStats{Proposed: 10, Accepted: 8}
	fmt.Printf("accepted %d/%d = %.0f%%\n", s.Accepted, s.Proposed, s.AcceptanceRate()*100)
}

Output:
accepted 8/10 = 80%

func MedusaGenerate 

func MedusaGenerate(model *GPT, heads *MedusaHeads, prompt []int, maxNew int, epsilon, delta float64) ([]int, SpecStats, error)

MedusaGenerate generates up to maxNew tokens with Medusa chain drafting (§T444, the decode loop over MedusaHeads + TypicalAcceptance; the single-path variant of the paper's candidate tree — tree assembly over top-k head candidates plugs into the same verification via MedusaTreeMask). Each round: (1) one forward computes the base's next greedy token x₁ (always emitted — the paper's guarantee that a round never stalls) and the K head proposals x₂…x_{K+1} from the same last hidden state; (2) ONE verification forward over seq+candidates scores every proposal under the base model, and the longest prefix passing TypicalAcceptance(ε, δ) is emitted. ε/δ ≤ 0 select the reference defaults (MedusaEpsilon/MedusaDelta). Typical acceptance is deliberately NOT distribution-exact (see TypicalAcceptance) — the output is greedy-anchored but may take plausible non-argmax tokens where the heads propose them. Greedy-only by construction (candidates are argmax proposals). Returns prompt+generated and the proposal/acceptance stats. Every round is O(seq) full forwards (analysis-scale, like DoLaDecode); the batched GPU variant belongs to llamagpu.

func MedusaGenerateTree 

func MedusaGenerateTree(model *GPT, heads *MedusaHeads, prompt []int, maxNew int, epsilon, delta float64, topK int) ([]int, SpecStats, error)

MedusaGenerateTree is MedusaGenerate with the paper's CANDIDATE TREE (§2.3): instead of one chain of argmax proposals, head k contributes its top-topK tokens and the Cartesian product forms a tree under the base's greedy token (always emitted). ONE verification forward scores every root-to-node path simultaneously — tree rows attend only to the prefix and their ancestors (MedusaTreeMask through the MHA.Mask seam, §T508) with positions P+depth — and the deepest fully TypicalAcceptance-accepted path is emitted (leftmost, i.e. highest-ranked candidates, on ties). topK=1 collapses onto MedusaGenerate's chain exactly. Deeper levels are dropped when the tree would not fit the context window or the remaining budget. Stats count path POSITIONS offered/accepted (comparable to MedusaGenerate), not tree nodes. Analysis-scale like MedusaGenerate: O(seq) full forwards per round, with the masked attention served by the reference kernel (§T461 fallback).

func PromptLookupDecode 

func PromptLookupDecode(model *GPT, prompt []int, maxNew, maxNgram, draftLen int, s *Sampler) ([]int, SpecStats, error)

PromptLookupDecode generates up to maxNew tokens after prompt with prompt-lookup (n-gram) speculative decoding (Yang et al. 2023 "Inference with Reference", LLMA, arXiv:2304.04487; Saxena's Prompt Lookup Decoding; §R89). It needs NO draft model: each round it copies a candidate continuation from the running sequence's own history (ngramLookup — the last maxNgram-token suffix's earlier occurrence), the model scores the whole draft in ONE forward pass, and the speculative accept/reject rule (SpeculativeRun, §R53) keeps the verified prefix plus one model token. Because the drafts carry a deterministic (point-mass) distribution, the accept/reject reduces to "keep the draft tokens the model would have produced anyway": the output is distributed EXACTLY as model.Generate(prompt, …, s) — token-for-token identical under greedy sampling (lossless), the speedup coming from tasks whose output copies the input (summarization, RAG, code editing). When no n-gram recurs it falls back to a plain one-token step. maxNgram≤0 defaults to 3, draftLen≤0 to 10. Returns prompt+generated, the accept stats, and any error.

func SpeculativeDecode 

func SpeculativeDecode(target, draft *GPT, prompt []int, maxNew, lookahead int, s *Sampler) ([]int, SpecStats, error)

SpeculativeDecode generates up to maxNew tokens after prompt with speculative decoding (Leviathan et al. 2023, §R53): each round the small DRAFT model proposes `lookahead` tokens autoregressively, the large TARGET model scores them in ONE forward pass, and SpeculativeRun accepts the verified prefix plus one residual-or- bonus token. The modified rejection rule makes the result distributed EXACTLY as target.Generate(prompt, …, s) under the same sampler — lossless — so a cheap draft only changes speed, never the output distribution. With draft==target every token is accepted. Returns prompt+generated, the accept stats, and any error.

This reference version recomputes full forward passes (no KV-cache) for clarity, so it is not yet faster than single-model decode; a cache-based draft loop plus a batched target verify is the perf follow-up.

func (SpecStats) AcceptanceRate 

func (s SpecStats) AcceptanceRate() float64

AcceptanceRate is Accepted/Proposed (0 when nothing was proposed).

type StreamCache 

type StreamCache struct {
	K, V []*tensor.Tensor // per block, PRE-RoPE, ≤ sinks+window rows each
}

StreamCache is a bounded KV-cache for StreamingLLM decoding (Xiao et al. 2023, "Efficient Streaming Language Models with Attention Sinks", arXiv:2309.17453, §R97). Unlike LlamaCache it keeps at most sinks+window entries per layer and stores keys BEFORE the rotary embedding, so RoPE can be re-applied at each step using each token's position WITHIN the current cache — the trick that lets generation run far past the training context length at constant memory.

func (*StreamCache) Len 

func (c *StreamCache) Len() int

Len returns the number of tokens currently cached.

type TokenSampler 

type TokenSampler interface {
	Sample(logits []float64) int
	SampleWithHistory(logits []float64, history []int) int
}

TokenSampler is anything that turns a logit row into a token id — the sequential generation loops (GPT/Llama/QuantLlama Generate, StreamGenerate, ContrastiveDecode, and the llamagpu decoders) accept any implementation. Sampler (temperature/top-k/ top-p/… truncation) and Mirostat (adaptive surprise targeting) both satisfy it. SampleWithHistory receives the running sequence (prompt + generated) so history-aware strategies (repetition penalties) can act; implementations without history use are free to ignore it. The speculative paths (SpeculativeDecode, SpeculativeGenerate, prompt-lookup) still require a concrete *Sampler — their lossless accept/reject math needs the full distribution (Dist), not just a draw.

Example

TokenSampler is the interface every sequential generation loop accepts — both the truncation-based Sampler and the adaptive Mirostat satisfy it, so either plugs into GPT/Llama Generate (or a llamagpu decoder) interchangeably. Greedy argmax picks token 0; a repeat penalty on the same history steers it to the runner-up. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	logits := []float64{2.0, 1.5, 0.1}
	var s nlp.TokenSampler = nlp.Greedy()
	fmt.Println(s.SampleWithHistory(logits, nil))

	s = nlp.NewSampler(1, nlp.WithTemperature(0), nlp.WithRepeatPenalty(2))
	fmt.Println(s.SampleWithHistory(logits, []int{0}))

	s = nlp.NewMirostat(7) // adaptive: also a TokenSampler
	tok := s.SampleWithHistory(logits, []int{0})
	fmt.Println(tok >= 0 && tok < len(logits))
}

Output:
0
1
true

type Tokenizer 

type Tokenizer struct {
	// contains filtered or unexported fields
}

Tokenizer is a byte-level BPE tokenizer compatible with tiktoken's gpt2 encoding (§T37). The rank table (byte-sequence → id) doubles as the merge priority (tiktoken's byte_pair_merge). Since all 256 single bytes are base tokens, decode(encode(x)) == x for ANY input — the tokenizer's round-trip invariant (§V15), guaranteed by construction. Algorithm: BPE (Sennrich et al. 2016, §R33), byte-level variant.

func LoadGPT2 

func LoadGPT2(path string) (*Tokenizer, error)

LoadGPT2 reads a tiktoken-exported rank file (base64(bytes) rank per line).

func TiktokenFromBytes 

func TiktokenFromBytes(data []byte) (*Tokenizer, error)

TiktokenFromBytes parses an in-memory tiktoken rank file into a byte-level BPE Tokenizer — the byte-slice counterpart of LoadGPT2.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	tok, _ := nlp.TiktokenFromBytes([]byte("YQ== 0\nYg== 1\nYWI= 2\n"))
	fmt.Println(tok.Decode([]int{2}))
}

Output:
ab

func (*Tokenizer) Decode 

func (t *Tokenizer) Decode(ids []int) string

Decode reconstructs the exact original text (byte-level → round-trip exact).

func (*Tokenizer) Encode 

func (t *Tokenizer) Encode(text string) []int

Encode turns text into token ids: GPT-2 pre-tokenization → byte-pair merge per piece.

func (*Tokenizer) ToTiktoken 

func (t *Tokenizer) ToTiktoken() []byte

ToTiktoken serializes the tokenizer back to the tiktoken rank-file format (base64(bytes) rank per line, ascending by rank), so TiktokenFromBytes(t.ToTiktoken()) round-trips it.

type UL2Denoiser 

type UL2Denoiser struct {
	Mode     UL2Mode // R, S or X
	Rate     float64 // corruption rate r
	MeanSpan float64 // mean span μ (unused for S)
	Weight   float64 // mixing proportion (need not be normalized)
}

UL2Denoiser is one denoiser in the mixture: a (Mode, Rate, MeanSpan) regime with a sampling Weight. Rate is the fraction of tokens corrupted (for S, the fraction placed in the trailing target span); MeanSpan is the mean corrupted-span length (ignored for S).

func SampleUL2Denoiser 

func SampleUL2Denoiser(mix []UL2Denoiser, rng *rand.Rand) UL2Denoiser

SampleUL2Denoiser draws one denoiser from the mixture with probability proportional to Weight.

func UL2MixtureOfDenoisers 

func UL2MixtureOfDenoisers() []UL2Denoiser

UL2MixtureOfDenoisers returns the paper's standard 7-denoiser mixture (UL2 §3.1.2): two R-denoisers (μ∈{3,8}, r=0.15), four X-denoisers (μ∈{3,8} at r=0.5, μ=64 at r∈{0.15,0.5}) and one S-denoiser (r=0.25). The paper mixes the seven with equal probability (§3.1.3), so every Weight is 1.

type UL2Mode 

type UL2Mode int

UL2Mode identifies which denoiser produced an example; its value is also the offset of the example's paradigm token from the sentinel base (R→base, S→base+1, X→base+2).

Example

ExampleUL2Mode shows the three denoiser paradigms UL2 mixes. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	for _, m := range []nlp.UL2Mode{nlp.UL2R, nlp.UL2S, nlp.UL2X} {
		fmt.Print(int(m), " ")
	}
	fmt.Println()
}

Output:
0 1 2

const (
	UL2R UL2Mode = iota // Regular: short-span T5 corruption
	UL2S                // Sequential: prefix-LM (single trailing span)
	UL2X                // eXtreme: long spans and/or high corruption rate
)

func UL2Reconstruct 

func UL2Reconstruct(input, target []int, sentinelBase int) (doc []int, mode UL2Mode, ok bool)

UL2Reconstruct recovers the original document and the denoiser mode from a UL2Denoise pair. It strips the leading paradigm token and reverses the span corruption (T5Reconstruct). ok is false for a malformed pair (empty input, a leading token outside the paradigm range, or an irreversible span layout).

type Unigram 

type Unigram struct {
	// contains filtered or unexported fields
}

Unigram is the SentencePiece Unigram-LM subword tokenizer (Kudo 2018, "Subword Regularization", arXiv:1804.10959, §3.1, §R87) — the tokenizer of the Llama, Mistral and T5 model families. The vocabulary is a fixed set of pieces each carrying a score = log p(piece) under a unigram language model P(x)=∏ᵢ p(xᵢ). Inference-time encoding is the 1-best (most-probable) segmentation, found by Viterbi dynamic programming that MAXIMIZES the sum of piece log-probabilities — not greedy longest-match, which can pick a lower-probability tiling.

DP recurrence over the input's Unicode characters, best[i] = max score of any segmentation of the first i characters: 

best[i] = max over pieces p ending at i of  best[i−len(p)] + score(p)

with backpointers to recover the pieces. A position no piece covers takes a single character as <unk> at score min(scores)−kUnkPenalty (kUnkPenalty=10), so the DP is always feasible. Whitespace is escaped to ▁ and a leading ▁ is prepended (add_dummy_prefix) before segmentation; Decode concatenates the pieces and maps ▁ back to spaces, dropping the dummy — lossless when the vocabulary covers every character. Input is assumed already Unicode-normalized (SentencePiece applies NFKC upstream; this tokenizer segments the string it is given).

Example

The Unigram tokenizer segments by the most-probable path over its scored vocabulary and detokenizes losslessly. Here "▁hello"+"▁world" is the best segmentation and decode restores the spacing. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	u, _ := nlp.NewUnigram([]nlp.UnigramPiece{
		{"<unk>", -20}, {"▁hello", -1}, {"▁world", -1}, {"▁", -5},
	})
	ids := u.Encode("hello world")
	fmt.Println(ids, u.Decode(ids))
}

Output:
[1 2] hello world

func NewUnigram 

func NewUnigram(vocab []UnigramPiece, opts ...UnigramOption) (*Unigram, error)

NewUnigram builds a Unigram tokenizer from a scored vocabulary (piece id = index). Scores are log-probabilities (≤ 0, larger = more frequent). It errors on an empty vocabulary or duplicate pieces.

func UnigramFromGGUF 

func UnigramFromGGUF(meta map[string]any) (*Unigram, error)

UnigramFromGGUF builds a Unigram tokenizer from the metadata map of a parsed GGUF model file (gguf.File.Metadata), reading the ggml/llama.cpp tokenizer convention (§R88): tokenizer.ggml.tokens (the vocabulary, an array of strings indexed by token id) paired with tokenizer.ggml.scores (an array of float32 log-probabilities), and tokenizer.ggml.unknown_token_id for the <unk> id. This wires a real .gguf model's embedded tokenizer to the Viterbi Unigram of NewUnigram, so weights and tokenizer load from one file.

It accepts tokenizer.ggml.model == "t5" (ggml's UGM — a true Unigram tokenizer, the exact match for the Viterbi decoding here) and "llama" (SentencePiece, which carries the same per-token scores and ▁ meta-symbol). NOTE: for "llama"/SPM, llama.cpp runs a greedy best-score-bigram MERGE rather than Viterbi; this loader always applies the Viterbi 1-best over the same scores, which is the higher-likelihood segmentation but can differ token-for-token from llama.cpp's SPM output (an explicit, documented choice — §R88). Byte-level BPE models ("gpt2"/"bpe") are not handled here (they need merge rules, not scores) and return an error.

Example

A GGUF model file carries its own tokenizer; UnigramFromGGUF wires that embedded vocabulary to the Viterbi Unigram so a .gguf model tokenizes end-to-end. 

package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	meta := map[string]any{
		"tokenizer.ggml.model":            "t5",
		"tokenizer.ggml.tokens":           []any{"<unk>", "▁the", "▁cat", "▁"},
		"tokenizer.ggml.scores":           []any{float32(-20), float32(-1), float32(-1), float32(-5)},
		"tokenizer.ggml.unknown_token_id": uint32(0),
	}
	u, _ := nlp.UnigramFromGGUF(meta)
	ids := u.Encode("the cat")
	fmt.Println(ids, u.Decode(ids))
}

Output:
[1 2] the cat

func UnigramFromJSON 

func UnigramFromJSON(data []byte, opts ...UnigramOption) (*Unigram, error)

UnigramFromJSON builds a SentencePiece Unigram tokenizer from HuggingFace tokenizer.json bytes. It errors on invalid JSON, a non-Unigram model type, an empty vocabulary, or a malformed vocab entry.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	j := `{"model":{"type":"Unigram","unk_id":0,"vocab":[["<unk>",0.0],["▁",-3.0],["h",-4.0],["i",-4.0],["▁hi",-5.0]]}}`
	u, _ := nlp.UnigramFromJSON([]byte(j))
	fmt.Println(u.Encode("hi"))
}

Output:
[4]

func (*Unigram) Decode 

func (u *Unigram) Decode(ids []int) string

Decode reconstructs text from token ids: concatenate the pieces, map ▁ back to spaces, and drop the leading dummy space. Lossless when the vocabulary covered every input character (no <unk> was emitted).

func (*Unigram) Encode 

func (u *Unigram) Encode(text string) []int

Encode segments text into token ids by the 1-best Viterbi over the unigram LM.

func (*Unigram) ToJSON 

func (u *Unigram) ToJSON() ([]byte, error)

ToJSON serializes the tokenizer to the minimal HuggingFace tokenizer.json Unigram form (model.type "Unigram" with unk_id and the [piece, score] vocab), so UnigramFromJSON(u.ToJSON()) round-trips it.

type UnigramOption 

type UnigramOption func(*Unigram)

UnigramOption configures a Unigram tokenizer (functional-options idiom, §C12).

func WithUnigramDummyPrefix 

func WithUnigramDummyPrefix(on bool) UnigramOption

WithUnigramDummyPrefix toggles prepending a ▁ before encoding (add_dummy_prefix, default true) so a leading word is tokenized like a mid-sentence one.

func WithUnigramUnkID 

func WithUnigramUnkID(id int) UnigramOption

WithUnigramUnkID sets the id emitted for characters no piece covers (default: the id of a piece literally named "<unk>", else 0).

type UnigramPiece 

type UnigramPiece struct {
	Piece string  // the subword (may include the ▁ space meta-symbol)
	Score float64 // log p(piece)
}

UnigramPiece is one vocabulary entry of a Unigram tokenizer: a subword string and its score, the log-probability log p(piece) of the unigram language model.

type Watermark 

type Watermark struct {
	VocabSize int     // size of the vocabulary V
	Gamma     float64 // green-list fraction γ ∈ (0,1)
	Delta     float64 // logit bias δ added to green tokens (soft watermark)
	Key       uint64  // secret key seeding the green-list PRF
}

Watermark holds the red-green watermark parameters.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	w, _ := nlp.NewWatermark(32, nlp.WithWatermarkGamma(0.25), nlp.WithWatermarkKey(2024))
	// Generate greedily from flat logits under the bias: every picked token is green.
	toks := []int{0}
	zero := make([]float64, 32)
	for len(toks) < 40 {
		biased, _ := w.BiasLogits(zero, toks[len(toks)-1])
		best := 0
		for i := range biased {
			if biased[i] > biased[best] {
				best = i
			}
		}
		toks = append(toks, best)
	}
	fmt.Println(w.IsWatermarked(toks))
}

Output:
true

func NewWatermark 

func NewWatermark(vocabSize int, opts ...WatermarkOption) (*Watermark, error)

NewWatermark builds a watermark over a vocabulary of vocabSize with the defaults γ=0.25, δ=2.0 and key 0 (γ=0.25 is the official lm-watermarking repo's recommended default; the paper's headline analysis uses γ=0.5 — set it with WithWatermarkGamma).

func (*Watermark) BiasLogits 

func (w *Watermark) BiasLogits(logits []float64, prevToken int) ([]float64, error)

BiasLogits returns a copy of logits with δ added to every green-list token for the step whose preceding token is prevToken (the soft watermark, Algorithm 2). The caller samples from the result as usual. len(logits) must equal VocabSize.

func (*Watermark) Detect 

func (w *Watermark) Detect(tokens []int) (z float64, green, scored int)

Detect scores a token sequence: each token from index 1 on is checked against the green list seeded by its predecessor. It returns the number of green tokens, the number scored (T = len−1) and the z-statistic z = (green − γ·T)/√(T·γ·(1−γ)) (0 when T=0). Compare z to WatermarkZThreshold.

func (*Watermark) GreenMask 

func (w *Watermark) GreenMask(prevToken int) []bool

GreenMask returns a length-VocabSize boolean mask of the green list seeded by prevToken: a partial Fisher–Yates shuffle picks the first ⌊γ·|V|⌋ ids of a (Key, prevToken)-seeded permutation as green. Deterministic — generation and detection call it identically.

func (*Watermark) IsWatermarked 

func (w *Watermark) IsWatermarked(tokens []int) bool

IsWatermarked reports whether Detect's z-score exceeds WatermarkZThreshold.

func (*Watermark) Sampler 

func (w *Watermark) Sampler(inner TokenSampler) TokenSampler

Sampler wraps inner so every sampled token is watermarked (§T439): before inner draws, BiasLogits adds δ to the green list seeded by the LAST history token, so the wrapper plugs straight into any generation loop that takes a TokenSampler (GPT/Llama/QuantLlama Generate, StreamGenerate, the llamagpu decoders). The history — prompt included — is what the loops already pass to SampleWithHistory. A history-less Sample (or a logits length ≠ VocabSize) draws unbiased: with no predecessor there is no green list to seed. Detect/IsWatermarked verify the output.

type WatermarkOption 

type WatermarkOption func(*Watermark)

WatermarkOption configures a Watermark via the functional-options idiom (§C12).

func WithWatermarkDelta 

func WithWatermarkDelta(d float64) WatermarkOption

WithWatermarkDelta sets the green-list logit bias δ (default 2.0).

func WithWatermarkGamma 

func WithWatermarkGamma(g float64) WatermarkOption

WithWatermarkGamma sets the green-list fraction γ (default 0.25).

func WithWatermarkKey 

func WithWatermarkKey(k uint64) WatermarkOption

WithWatermarkKey sets the secret key seeding the green lists (default 0).

type WordPiece 

type WordPiece struct {
	// contains filtered or unexported fields
}

WordPiece is the subword tokenizer of BERT (Devlin, Chang, Lee & Toutanova 2019; the wordpiece model of Schuster & Nakajima 2012 and Wu et al. 2016 GNMT §4). It completes this library's tokenizer set alongside byte-level BPE (GPT-2/Llama-3, nlp.Tokenizer) and Unigram/SentencePiece (Llama/Mistral/T5, nlp.Unigram §R87). Given a pre-tokenized word, WordPiece encodes it by GREEDY LONGEST-MATCH-FIRST (MaxMatch): from the start of the word it takes the LONGEST prefix present in the vocabulary, emits it, and continues from the remainder — distinct from BPE's learned merges and Unigram's probabilistic Viterbi. Continuation pieces (any not at the word start) are looked up with a "##" prefix (e.g. "playing" with {play, ##ing} → [play, ##ing]); if at some position no substring down to a single character matches, the WHOLE word becomes one [UNK]; words longer than MaxChars are [UNK] without tokenizing.

Basic pre-tokenization (punctuation splitting, lowercasing, accent stripping) and normalization are assumed applied upstream (as with Unigram's NFKC); Encode splits the input on whitespace and runs WordPiece per word.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	w, _ := nlp.NewWordPiece([]string{"[UNK]", "play", "##ing", "un", "##able"})
	fmt.Println(w.Encode("playing"))
	fmt.Println(w.Decode(w.Encode("playing")))
}

Output:
[1 2]
playing

func NewWordPiece 

func NewWordPiece(vocab []string, opts ...WordPieceOption) (*WordPiece, error)

NewWordPiece builds a WordPiece tokenizer from vocab (id = index). Continuation pieces carry the "##" prefix in vocab. Defaults: continuation "##", maxChars 100, unk = id of "[UNK]" if present.

func WordPieceFromJSON 

func WordPieceFromJSON(data []byte, opts ...WordPieceOption) (*WordPiece, error)

WordPieceFromJSON builds a WordPiece tokenizer from HuggingFace tokenizer.json bytes. It errors on invalid JSON, a non-WordPiece model type, or an empty vocabulary. The unk_token, continuing_subword_prefix ("##") and max_input_chars_per_word (100) from the file are applied; caller opts override them.

Example 
package main

import (
	"fmt"

	"github.com/jxsl13/goai/nlp"
)

func main() {
	js := `{"model":{"type":"WordPiece","unk_token":"[UNK]","continuing_subword_prefix":"##",
	        "vocab":{"[UNK]":0,"play":1,"##ing":2}}}`
	w, _ := nlp.WordPieceFromJSON([]byte(js))
	fmt.Println(w.Encode("playing"))
}

Output:
[1 2]

func (*WordPiece) Decode 

func (w *WordPiece) Decode(ids []int) string

Decode concatenates the pieces of ids: a continuation piece ("##…") is appended to the current word (prefix stripped), a non-continuation piece starts a new space-separated word.

func (*WordPiece) Encode 

func (w *WordPiece) Encode(text string) []int

Encode tokenizes text into ids: it splits on whitespace and applies greedy longest-match-first WordPiece to each word.

func (*WordPiece) ToJSON 

func (w *WordPiece) ToJSON() ([]byte, error)

ToJSON serializes the tokenizer to the minimal HuggingFace tokenizer.json WordPiece form, so WordPieceFromJSON(w.ToJSON()) round-trips it.

type WordPieceOption 

type WordPieceOption func(*WordPiece)

WordPieceOption configures a WordPiece tokenizer (functional-options idiom, §C12).

func WithWordPieceContinuation 

func WithWordPieceContinuation(p string) WordPieceOption

WithWordPieceContinuation sets the continuation-subword prefix (default "##").

func WithWordPieceMaxChars 

func WithWordPieceMaxChars(n int) WordPieceOption

WithWordPieceMaxChars sets the max characters (runes) per word before it is emitted as unk (default 100, the BERT/HF value); non-positive is ignored.

func WithWordPieceUnk 

func WithWordPieceUnk(id int) WordPieceOption

WithWordPieceUnk sets the id emitted for an unmatchable or over-long word (default: the id of the "[UNK]" piece if present, else 0).

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.