llm-samplers 0.0.3

Token samplers for large language models
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

use crate::types::*;

/// Temperature sampling
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SampleTemperature<T> {
    temperature: T,
}

impl<T: CanLogit> SampleTemperature<T> {
    pub fn new(temperature: T) -> Self {
        Self { temperature }
    }
}

impl<TID: CanTokenId, L: CanLogit> Sampler<TID, L> for SampleTemperature<L> {
    fn sample<'a>(
        &mut self,
        logits: &'a mut Logits<TID, L>,
    ) -> Result<&'a mut Logits<TID, L>, SamplerError> {
        let temp = self.temperature;
        logits.iter_mut().for_each(|l| l.logit = l.logit / temp);
        Ok(logits)
    }
}