Beam search over a next-token scorer (Geron Ch 14, morie.fn hmbms)

model(src, prefix) returns log-probabilities over the vocabulary for the token after prefix (a 0-based integer vector). The scorer is checked for normalisation, because an unnormalised one turns beam search into an arbitrary greedy walk. Emitted tokens are 0-based.

Usage

morie_geron_beam_search(
  model,
  src,
  beam_width = 3,
  max_len = 10,
  eos = NULL,
  length_penalty = 0
)

Arguments

model

Scorer function.

src

Source, passed through.

beam_width

Hypotheses kept per step.

max_len

Token cap.

eos

Optional 0-based end-of-sequence token.

length_penalty

Exponent alpha in score / len^alpha for the final rank.

Value

List with sequence, score, beams, scores, finished.

Examples

lp <- log(c(0.5, 0.3, 0.2))
morie_geron_beam_search(function(s, prefix) lp, NULL, beam_width = 2,
    max_len = 3)
#> $sequence
#> [1] 0 0 0
#> 
#> $score
#> [1] -2.079442
#> 
#> $beams
#> $beams[[1]]
#> [1] 0 0 0
#> 
#> $beams[[2]]
#> [1] 0 0 1
#> 
#> 
#> $scores
#> [1] -2.079442 -2.590267
#> 
#> $finished
#> [1] 2
#> 
#> $estimate
#> [1] -2.079442
#> 
#> $n
#> [1] 3
#> 
#> $method
#> [1] "Beam search with width 2"
#>