decode

Documentation

Overview

Package decode provides decoding algorithms for sequence-to-sequence and transducer models.

Includes greedy decoding, beam search, and utility functions like top-k selection and length normalization used across different decoding strategies.

Index

• func BeamSearch(initState interface{}, ...) []int
• func GreedyDecode(firstToken, eotToken, maxSteps int, ...) []int
• func LengthNormalization(seqLen int, alpha float32) float32
• func TopKIndicesHeap(vals []float32, k int) []int
• type BeamConfig
• type BeamHypothesis

Functions

func BeamSearch

func BeamSearch(
	initState interface{},
	expandFunc func(hyp *BeamHypothesis) (logProbs []float32, newState interface{}),
	cfg BeamConfig,
) []int

BeamSearch performs generic beam search decoding.

initState:  initial decoder state
expandFunc: given a hypothesis, returns (logits, newState) for each possible next step
cfg:        beam search configuration

The expandFunc receives the current hypothesis and should return the log-probability distribution over the next token and an updated state.

func GreedyDecode

func GreedyDecode(firstToken, eotToken, maxSteps int, logitsFunc func(tokenID, pos int) []float32) []int

GreedyDecode performs simple greedy argmax decoding. logitsFunc is called for each step with the current token and position, returning logits over the vocabulary. Decoding stops when eotToken is emitted or maxSteps is reached.

func LengthNormalization

func LengthNormalization(seqLen int, alpha float32) float32

LengthNormalization computes Google NMT length penalty:

penalty = ((5 + seqLen) / 6) ^ alpha

func TopKIndicesHeap

func TopKIndicesHeap(vals []float32, k int) []int

TopKIndicesHeap returns the indices of the k largest values using a min-heap. O(n log k) complexity.

Types

type BeamConfig

type BeamConfig struct {
	BeamSize    int
	MaxSteps    int
	EOTToken    int
	LengthAlpha float32 // length normalization exponent (0 = disabled)
}

BeamConfig controls beam search behavior.

type BeamHypothesis

type BeamHypothesis struct {
	Tokens   []int
	Score    float32
	Finished bool
	State    interface{} // opaque decoder state for the hypothesis
}

BeamHypothesis represents one candidate sequence during beam search.