token-sampling 0.1.0

use core::cmp::Ordering;
use core::fmt;

/// One token candidate with its current logit and probability.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Candidate {
    /// Token identifier from the original logits index.
    pub id: u32,
    /// Current logit.
    pub logit: f32,
    /// Probability from the most recent softmax.
    pub p: f32,
}

/// Errors returned while loading or sampling logits.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Error {
    /// The logits slice is empty.
    EmptyLogits,
    /// A logit is NaN or positive infinity.
    InvalidLogit {
        /// Index of the first invalid value.
        index: usize,
    },
    /// Every input logit is negative infinity.
    AllMasked,
    /// The chain finished without a terminal stage selecting a token.
    NoSelection,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptyLogits => f.write_str("logits are empty"),
            Self::InvalidLogit { index } => write!(f, "logit at index {index} is invalid"),
            Self::AllMasked => f.write_str("all logits are masked"),
            Self::NoSelection => f.write_str("no sampling stage selected a token"),
        }
    }
}

impl std::error::Error for Error {}

/// Reusable token candidates and stage scratch storage.
#[derive(Clone, Debug)]
pub struct Candidates {
    pub(crate) buf: Vec<Candidate>,
    pub(crate) counts: Vec<(u32, u32)>,
    pub(crate) repeat_count: Vec<usize>,
    pub(crate) max_repeats: Vec<(u32, usize)>,
    pub(crate) breaker_indices: Vec<usize>,
    pub(crate) sorted: bool,
    pub(crate) selected: Option<usize>,
}

impl Candidates {
    /// Creates an empty candidate set.
    pub const fn new() -> Self {
        Self {
            buf: Vec::new(),
            counts: Vec::new(),
            repeat_count: Vec::new(),
            max_repeats: Vec::new(),
            breaker_indices: Vec::new(),
            sorted: false,
            selected: None,
        }
    }

    /// Creates an empty candidate set sized for `n_vocab` logits.
    pub fn with_capacity(n_vocab: usize) -> Self {
        Self {
            buf: Vec::with_capacity(n_vocab),
            counts: Vec::new(),
            repeat_count: Vec::new(),
            max_repeats: Vec::new(),
            breaker_indices: Vec::new(),
            sorted: false,
            selected: None,
        }
    }

    /// Validates and loads logits while retaining allocated storage.
    pub fn load(&mut self, logits: &[f32]) -> Result<(), Error> {
        if logits.is_empty() {
            return Err(Error::EmptyLogits);
        }
        let mut any_unmasked = false;
        for (index, &logit) in logits.iter().enumerate() {
            if logit.is_nan() || logit == f32::INFINITY {
                return Err(Error::InvalidLogit { index });
            }
            any_unmasked |= logit != f32::NEG_INFINITY;
        }
        if !any_unmasked {
            return Err(Error::AllMasked);
        }

        reserve_to(&mut self.buf, logits.len());
        self.buf.clear();
        self.buf
            .extend(logits.iter().enumerate().map(|(id, &logit)| Candidate {
                id: id as u32,
                logit,
                p: 0.0,
            }));
        self.counts.clear();
        self.repeat_count.clear();
        self.max_repeats.clear();
        self.breaker_indices.clear();
        self.sorted = false;
        self.selected = None;
        Ok(())
    }

    /// Returns the current candidates in stage order.
    pub fn as_slice(&self) -> &[Candidate] {
        &self.buf
    }

    /// Returns the selected token from the last terminal stage.
    pub fn selected(&self) -> Option<u32> {
        self.selected
            .and_then(|position| self.buf.get(position))
            .map(|candidate| candidate.id)
    }

    pub(crate) fn sort_desc(&mut self) {
        if !self.sorted {
            self.buf.sort_unstable_by(candidate_order);
            self.sorted = true;
        }
    }

    pub(crate) fn softmax(&mut self, sort: bool) -> Result<(), Error> {
        if self.buf.is_empty() {
            return Err(Error::AllMasked);
        }
        if !self.sanitize_logits() {
            return Err(Error::AllMasked);
        }
        if sort {
            self.sort_desc();
        }
        let max_logit = self
            .buf
            .iter()
            .map(|candidate| candidate.logit)
            .fold(f32::NEG_INFINITY, f32::max);
        let sum = self
            .buf
            .iter()
            .map(|candidate| ((candidate.logit - max_logit) as f64).exp())
            .sum::<f64>();
        if !sum.is_finite() || sum <= 0.0 {
            return Err(Error::AllMasked);
        }
        for candidate in &mut self.buf {
            candidate.p = (((candidate.logit - max_logit) as f64).exp() / sum) as f32;
        }
        Ok(())
    }

    pub(crate) fn sanitize_logits(&mut self) -> bool {
        let mut any_unmasked = false;
        let mut changed = false;
        for candidate in &mut self.buf {
            if candidate.logit.is_nan() {
                candidate.logit = f32::NEG_INFINITY;
                changed = true;
            } else if candidate.logit == f32::INFINITY {
                candidate.logit = f32::MAX;
                changed = true;
            }
            any_unmasked |= candidate.logit.is_finite();
        }
        if changed {
            self.sorted = false;
        }
        any_unmasked
    }
}

impl Default for Candidates {
    fn default() -> Self {
        Self::new()
    }
}

pub(crate) fn candidate_order(left: &Candidate, right: &Candidate) -> Ordering {
    right
        .logit
        .total_cmp(&left.logit)
        .then_with(|| left.id.cmp(&right.id))
}

fn reserve_to<T>(values: &mut Vec<T>, capacity: usize) {
    if values.capacity() < capacity {
        values.reserve_exact(capacity - values.capacity());
    }
}