Struct TopKSampler 

pub struct TopKSampler { /* private fields */ }

Top-k sampler for text generation.

Samples from the k most likely tokens.

Example

let sampler = TopKSampler::new(50);
let token = sampler.sample(&logits);

Implementations

impl TopKSampler

pub fn new(top_k: usize) -> Self

Create a new top-k sampler.

Arguments

• top_k - Number of top tokens to keep

pub fn with_temperature(self, temperature: f32) -> Self

Set temperature for sampling.

pub fn filter(&self, logits: &Tensor) -> Tensor

Filter logits to keep only top-k tokens.

pub fn sample(&self, logits: &Tensor) -> usize

Sample a token from the filtered distribution.

pub fn top_k(&self) -> usize

Get top_k value.

Trait Implementations

impl Debug for TopKSampler

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.