llama_cpp_4/

sampling.rs

//! Safe wrapper around `llama_sampler`.

use std::borrow::Borrow;
use std::ffi::{c_char, CString};
use std::fmt::{Debug, Formatter};
use std::ptr::NonNull;

use llama_cpp_sys_4::{
    common::common_sampler_params, llama_logit_bias, llama_sampler, llama_sampler_chain_add,
    llama_sampler_chain_default_params, llama_sampler_chain_init, llama_sampler_chain_n,
    llama_sampler_chain_remove, llama_sampler_clone, llama_sampler_copy, llama_sampler_free,
    llama_sampler_get_seed, llama_sampler_init_adaptive_p, llama_sampler_init_dist,
    llama_sampler_init_dry, llama_sampler_init_grammar, llama_sampler_init_grammar_lazy_patterns,
    llama_sampler_init_greedy, llama_sampler_init_infill, llama_sampler_init_logit_bias,
    llama_sampler_init_min_p, llama_sampler_init_mirostat, llama_sampler_init_mirostat_v2,
    llama_sampler_init_penalties, llama_sampler_init_temp, llama_sampler_init_temp_ext,
    llama_sampler_init_top_k, llama_sampler_init_top_n_sigma, llama_sampler_init_top_p,
    llama_sampler_init_typical, llama_sampler_init_xtc, llama_sampler_name, llama_sampler_reset,
};

use crate::context::LlamaContext;
use crate::model::LlamaModel;
use crate::token::data_array::LlamaTokenDataArray;
use crate::token::LlamaToken;

/// Why a sampler could not be constructed.
///
/// The panicking constructors — [`LlamaSampler::grammar`] and friends — are
/// kept for the common case where the grammar is a literal or came from
/// [`ChatParams::grammar`](crate::chat::ChatParams), and a failure really is a
/// bug. These exist for the case they handle badly: a grammar, rule name or
/// trigger pattern that arrived from a request, where an interior NUL is bad
/// *input* rather than a broken invariant.
#[derive(Debug, PartialEq, Eq, thiserror::Error)]
pub enum SamplerInitError {
    /// A caller-supplied string contained an interior NUL, so it cannot cross
    /// the FFI boundary. `field` names which argument.
    #[error("{field} contains an interior NUL byte")]
    InteriorNul {
        /// Name of the offending argument.
        field: &'static str,
        /// The underlying conversion error.
        #[source]
        source: std::ffi::NulError,
    },
    /// llama.cpp returned null. For a grammar sampler this is what an
    /// unparseable GBNF grammar looks like; it logs the parse error itself.
    #[error("llama.cpp could not build the {0} sampler (see the llama.cpp log)")]
    NullReturn(&'static str),
    /// More distinct DRY sequence breakers than llama.cpp will accept. It
    /// silently drops the surplus, so this is reported rather than passed on.
    #[error("{distinct} distinct DRY sequence breakers, but llama.cpp accepts at most {max}")]
    TooManyDrySequenceBreakers {
        /// How many distinct breakers the caller supplied, counted the way
        /// llama.cpp counts them.
        distinct: usize,
        /// The upstream limit, [`DRY_MAX_SEQ_BREAKERS`].
        max: usize,
    },
}

/// Longest DRY sequence breaker llama.cpp keeps; it truncates anything longer.
///
/// Truncation happens *before* de-duplication, so two breakers that differ only
/// past this point count as one.
pub const DRY_MAX_BREAKER_LEN: usize = 40;

/// Most distinct DRY sequence breakers llama.cpp will process.
///
/// Mirrors `LLAMA_DRY_MAX_SEQ_BREAKERS`. Past this, llama.cpp logs a warning
/// and ignores the rest — the sampler is then quietly not the one that was
/// asked for, which is why [`LlamaSampler::try_dry`] refuses instead. Each
/// breaker scans the whole vocabulary during init, so the cost is
/// `O(breakers x n_vocab)`.
pub const DRY_MAX_SEQ_BREAKERS: usize = llama_cpp_sys_4::LLAMA_DRY_MAX_SEQ_BREAKERS as usize;

/// Count breakers the way llama.cpp does: truncate to [`DRY_MAX_BREAKER_LEN`]
/// bytes, drop empties, then de-duplicate.
fn distinct_dry_breakers(breakers: &[CString]) -> usize {
    breakers
        .iter()
        .map(|b| {
            let bytes = b.as_bytes();
            &bytes[..bytes.len().min(DRY_MAX_BREAKER_LEN)]
        })
        .filter(|b| !b.is_empty())
        .collect::<std::collections::BTreeSet<_>>()
        .len()
}

/// Convert a caller string, naming the argument so the error says which one.
fn c_string(field: &'static str, value: &str) -> Result<CString, SamplerInitError> {
    CString::new(value).map_err(|source| SamplerInitError::InteriorNul { field, source })
}

/// A safe wrapper around `llama_sampler`.
pub struct LlamaSampler {
    pub(crate) sampler: NonNull<llama_sampler>,
}

impl Debug for LlamaSampler {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LlamaSamplerChain").finish()
    }
}
#[derive(Debug, Clone)]
#[allow(
    missing_docs,
    clippy::struct_excessive_bools,
    clippy::module_name_repetitions,
    dead_code
)]
pub struct LlamaSamplerParams {
    top_k: i32,
    top_p: f32,
    temp: f32,
    seed: u32,
}

impl LlamaSamplerParams {
    /// Set the seed of the context
    ///
    /// # Examples
    ///
    /// ```rust
    /// use llama_cpp_4::sampling::LlamaSamplerParams;
    /// let params = LlamaSamplerParams::default();
    /// let params = params.with_seed(1234);
    /// assert_eq!(params.seed(), 1234);
    /// ```
    #[must_use]
    pub fn with_seed(mut self, seed: u32) -> Self {
        self.seed = seed;
        self
    }

    /// Get the seed of the context
    ///
    /// # Examples
    ///
    /// ```rust
    /// use llama_cpp_4::sampling::LlamaSamplerParams;
    /// let params = LlamaSamplerParams::default()
    ///     .with_seed(1234);
    /// assert_eq!(params.seed(), 1234);
    /// ```
    #[must_use]
    pub fn seed(&self) -> u32 {
        self.seed
    }
}

impl Default for LlamaSamplerParams {
    fn default() -> Self {
        Self {
            top_k: 50,
            top_p: 0.9,
            temp: 0.8,
            seed: 1234,
        }
    }
}

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

impl LlamaSampler {
    /// Create new sampler with default params.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn new() -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sparams = unsafe { llama_sampler_chain_default_params() };

        Self {
            // SAFETY: the arguments are either scalars or pointers that outlive the call, and
            // the library owns whatever it returns.
            sampler: NonNull::new(unsafe { llama_sampler_chain_init(sparams) }).unwrap(),
        }
    }

    /// Sample and accept a token from the idx-th output of the last evaluation
    #[must_use]
    /// # Panics
    ///
    /// Panics if llama.cpp's grammar code rejects the state — see
    /// [`Self::try_sample`], which returns the error instead. This *panics*
    /// rather than aborting because the call is routed through a guard; calling
    /// `llama_sampler_sample` directly would let a C++ exception unwind into
    /// Rust and kill the process.
    pub fn sample(&self, ctx: &LlamaContext, idx: i32) -> LlamaToken {
        self.try_sample(ctx, idx)
            .unwrap_or_else(|e| panic!("sampling failed: {e}"))
    }

    /// Sample a token, reporting llama.cpp's failures instead of panicking.
    ///
    /// The realistic failure is a grammar that can no longer accept anything:
    /// llama.cpp raises *"Unexpected empty grammar stack"* when a model's
    /// vocabulary cannot satisfy the constraint — a JSON schema against a
    /// vocabulary with no `{`, say. That is a property of the model/grammar
    /// pair, not a bug, so it is worth handling.
    ///
    /// # Errors
    ///
    /// Returns [`ShimError::Failed`](crate::shim::ShimError::Failed) carrying
    /// llama.cpp's message.
    pub fn try_sample(
        &self,
        ctx: &LlamaContext,
        idx: i32,
    ) -> Result<LlamaToken, crate::shim::ShimError> {
        let mut status = llama_cpp_sys_4::LLAMA_SHIM_OK;
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let token = unsafe {
            llama_cpp_sys_4::common_shim_sampler_sample_raw(
                self.sampler.as_ptr(),
                ctx.context.as_ptr(),
                idx,
                &raw mut status,
            )
        };
        crate::shim::check_status(status)?;
        Ok(LlamaToken(token))
    }

    /// Applies this sampler to a [`LlamaTokenDataArray`].
    pub fn apply(&mut self, data_array: &mut LlamaTokenDataArray) {
        data_array.apply_sampler(self);
    }

    /// Accepts a token from the sampler, possibly updating the internal state of certain samplers
    /// (e.g. grammar, repetition, etc.)
    /// # Panics
    ///
    /// Panics if the token leaves a grammar with nothing it can accept — see
    /// [`Self::try_accept`].
    pub fn accept(&mut self, token: LlamaToken) {
        self.try_accept(token)
            .unwrap_or_else(|e| panic!("accepting token {}: {e}", token.0));
    }

    /// Accept a token, reporting llama.cpp's failures instead of panicking.
    ///
    /// # Errors
    ///
    /// Returns [`ShimError::Failed`](crate::shim::ShimError::Failed) when the
    /// grammar cannot accept this token — see [`Self::try_sample`].
    pub fn try_accept(&mut self, token: LlamaToken) -> Result<(), crate::shim::ShimError> {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            llama_cpp_sys_4::common_shim_sampler_accept_raw(self.sampler.as_ptr(), token.0)
        };
        crate::shim::check_status(status)
    }

    /// Accepts several tokens from the sampler or context, possibly updating the internal state of
    /// certain samplers (e.g. grammar, repetition, etc.)
    /// # Panics
    ///
    /// Panics on the first token the grammar cannot accept — see
    /// [`Self::try_accept`].
    pub fn accept_many(&mut self, tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>) {
        for token in tokens {
            self.accept(*token.borrow());
        }
    }

    /// Accepts several tokens from the sampler or context, possibly updating the internal state of
    /// certain samplers (e.g. grammar, repetition, etc.)
    #[must_use]
    pub fn with_tokens(
        mut self,
        tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>,
    ) -> Self {
        self.accept_many(tokens);
        self
    }

    /// Combines a list of samplers into a single sampler that applies each component sampler one
    /// after another.
    ///
    /// If you are using a chain to select a token, the chain should always end with one of
    /// [`LlamaSampler::greedy`], [`LlamaSampler::dist`], [`LlamaSampler::mirostat`], and
    /// [`LlamaSampler::mirostat_v2`].
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn chain(samplers: impl IntoIterator<Item = Self>, no_perf: bool) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        unsafe {
            let mut params = llama_sampler_chain_default_params();
            params.no_perf = no_perf;
            let chain = llama_sampler_chain_init(params);

            for sampler in samplers {
                llama_sampler_chain_add(chain, sampler.sampler.as_ptr());

                // Do not call `llama_sampler_free` on the sampler, as the internal sampler is now
                // owned by the chain
                std::mem::forget(sampler);
            }

            Self {
                sampler: NonNull::new(chain).unwrap(),
            }
        }
    }

    /// Same as [`Self::chain`] with `no_perf = false`.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Example
    /// ```rust
    /// use llama_cpp_4::token::{
    ///    LlamaToken,
    ///    data::LlamaTokenData,
    ///    data_array::LlamaTokenDataArray
    /// };
    /// use llama_cpp_4::sampling::LlamaSampler;
    ///
    /// let mut data_array = LlamaTokenDataArray::new(vec![
    ///     LlamaTokenData::new(LlamaToken(0), 0., 0.),
    ///     LlamaTokenData::new(LlamaToken(1), 1., 0.),
    ///     LlamaTokenData::new(LlamaToken(2), 2., 0.),
    /// ], false);
    ///
    /// data_array.apply_sampler(&mut LlamaSampler::chain_simple([
    ///     LlamaSampler::temp(0.5),
    ///     LlamaSampler::greedy(),
    /// ]));
    ///
    /// assert_eq!(data_array.data[0].logit(), 0.);
    /// assert_eq!(data_array.data[1].logit(), 2.);
    /// assert_eq!(data_array.data[2].logit(), 4.);
    ///
    /// assert_eq!(data_array.data.len(), 3);
    /// assert_eq!(data_array.selected_token(), Some(LlamaToken(2)));
    /// ```
    #[must_use]
    pub fn chain_simple(samplers: impl IntoIterator<Item = Self>) -> Self {
        Self::chain(samplers, false)
    }

    /// Updates the logits `l_i`' = `l_i/t`. When `t <= 0.0`, the maximum logit is kept at its original
    /// value, the rest are set to -inf.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Example:
    /// ```rust
    /// use llama_cpp_4::token::{
    ///    LlamaToken,
    ///    data::LlamaTokenData,
    ///    data_array::LlamaTokenDataArray
    /// };
    /// use llama_cpp_4::sampling::LlamaSampler;
    ///
    /// let mut data_array = LlamaTokenDataArray::new(vec![
    ///     LlamaTokenData::new(LlamaToken(0), 0., 0.),
    ///     LlamaTokenData::new(LlamaToken(1), 1., 0.),
    ///     LlamaTokenData::new(LlamaToken(2), 2., 0.),
    /// ], false);
    ///
    /// data_array.apply_sampler(&mut LlamaSampler::temp(0.5));
    ///
    /// assert_eq!(data_array.data[0].logit(), 0.);
    /// assert_eq!(data_array.data[1].logit(), 2.);
    /// assert_eq!(data_array.data[2].logit(), 4.);
    /// ```
    #[must_use]
    pub fn temp(t: f32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_temp(t) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Dynamic temperature implementation (a.k.a. entropy) described in the paper
    /// <https://arxiv.org/abs/2309.02772>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn temp_ext(t: f32, delta: f32, exponent: f32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_temp_ext(t, delta, exponent) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration"
    /// <https://arxiv.org/abs/1904.09751>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Example:
    /// ```rust
    /// use llama_cpp_4::token::{
    ///    LlamaToken,
    ///    data::LlamaTokenData,
    ///    data_array::LlamaTokenDataArray
    /// };
    /// use llama_cpp_4::sampling::LlamaSampler;
    ///
    /// let mut data_array = LlamaTokenDataArray::new(vec![
    ///     LlamaTokenData::new(LlamaToken(0), 0., 0.),
    ///     LlamaTokenData::new(LlamaToken(1), 1., 0.),
    ///     LlamaTokenData::new(LlamaToken(2), 2., 0.),
    ///     LlamaTokenData::new(LlamaToken(3), 3., 0.),
    /// ], false);
    ///
    /// data_array.apply_sampler(&mut LlamaSampler::top_k(2));
    ///
    /// assert_eq!(data_array.data.len(), 2);
    /// assert_eq!(data_array.data[0].id(), LlamaToken(3));
    /// assert_eq!(data_array.data[1].id(), LlamaToken(2));
    /// ```
    #[must_use]
    pub fn top_k(k: i32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_top_k(k) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Locally Typical Sampling implementation described in the paper <https://arxiv.org/abs/2202.00666>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn typical(p: f32, min_keep: usize) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_typical(p, min_keep) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration"
    /// <https://arxiv.org/abs/1904.09751>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn top_p(p: f32, min_keep: usize) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_top_p(p, min_keep) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Minimum P sampling as described in <https://github.com/ggerganov/llama.cpp/pull/3841>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn min_p(p: f32, min_keep: usize) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_min_p(p, min_keep) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// XTC sampler as described in <https://github.com/oobabooga/text-generation-webui/pull/6335>.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn xtc(p: f32, t: f32, min_keep: usize, seed: u32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_xtc(p, t, min_keep, seed) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Grammar sampler
    ///
    /// # Panics
    /// - If either of `grammar_str` or `grammar_root` contain null bytes.
    /// - If llama.cpp returns a null pointer, which is also what an
    ///   unparseable GBNF grammar looks like.
    ///
    /// Prefer [`Self::try_grammar`] when the grammar did not come from your own
    /// source — a grammar built from a request is input, and both failures above
    /// are then reachable without a bug.
    #[must_use]
    pub fn grammar(model: &LlamaModel, grammar_str: &str, grammar_root: &str) -> Self {
        Self::try_grammar(model, grammar_str, grammar_root)
            .unwrap_or_else(|e| panic!("building grammar sampler: {e}"))
    }

    /// [`Self::grammar`], reporting llama.cpp's and the conversion's failures
    /// instead of panicking.
    ///
    /// # Errors
    ///
    /// [`SamplerInitError::InteriorNul`] if either string contains a NUL, or
    /// [`SamplerInitError::NullReturn`] if llama.cpp rejected the grammar.
    pub fn try_grammar(
        model: &LlamaModel,
        grammar_str: &str,
        grammar_root: &str,
    ) -> Result<Self, SamplerInitError> {
        let grammar_str = c_string("grammar_str", grammar_str)?;
        let grammar_root = c_string("grammar_root", grammar_root)?;

        // SAFETY: both pointers come from live `CString`s that outlast the call,
        // and llama.cpp only reads them. The vocab reference is borrowed from
        // `model` for the duration.
        let sampler = unsafe {
            llama_sampler_init_grammar(
                model.get_vocab().vocab.as_ref(),
                grammar_str.as_ptr(),
                grammar_root.as_ptr(),
            )
        };
        NonNull::new(sampler)
            .map(|sampler| Self { sampler })
            .ok_or(SamplerInitError::NullReturn("grammar"))
    }

    /// DRY sampler, designed by p-e-w, as described in:
    /// <https://github.com/oobabooga/text-generation-webui/pull/5677>, porting Koboldcpp
    /// implementation authored by pi6am: <https://github.com/LostRuins/koboldcpp/pull/982>
    ///
    /// # Panics
    /// - If any string in `seq_breakers` contains null bytes.
    /// - If llama.cpp returns a null pointer.
    #[allow(clippy::too_many_arguments)]
    #[must_use]
    pub fn dry(
        &self,
        model: &LlamaModel,
        multiplier: f32,
        base: f32,
        allowed_length: i32,
        penalty_last_n: i32,
        seq_breakers: impl IntoIterator<Item = impl AsRef<[u8]>>,
    ) -> Self {
        // A NUL in a sequence breaker is the same class of input error the
        // grammar constructors have a `try_` variant for; `dry` keeps the panic
        // because its breakers are conventionally literals (`"\n"`, `":"`).
        let seq_breakers: Vec<CString> = seq_breakers
            .into_iter()
            .map(|s| CString::new(s.as_ref()).expect("sequence breaker contains an interior NUL"))
            .collect();
        // CString::as_ptr() returns *const c_char, which matches what the binding
        // expects on every platform (signed on macOS/x86 Linux, unsigned on musl ARM).
        let mut seq_breaker_pointers: Vec<*const c_char> =
            seq_breakers.iter().map(|s| s.as_ptr()).collect();

        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe {
            llama_sampler_init_dry(
                model.get_vocab().vocab.as_ref(),
                multiplier,
                base,
                allowed_length,
                penalty_last_n,
                seq_breaker_pointers.as_mut_ptr(),
                seq_breaker_pointers.len(),
            )
        };

        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// [`Self::dry`], but reporting the two things it cannot.
    ///
    /// `dry` panics on a breaker containing an interior NUL, and — worse —
    /// cannot react at all when there are too many breakers: llama.cpp keeps
    /// the first [`DRY_MAX_SEQ_BREAKERS`] distinct ones, logs a warning, and
    /// returns a perfectly valid sampler that simply is not the one you asked
    /// for. Prefer this whenever the breakers come from configuration or user
    /// input rather than from literals.
    ///
    /// # Errors
    ///
    /// - [`SamplerInitError::InteriorNul`] if a breaker contains a NUL byte.
    /// - [`SamplerInitError::TooManyDrySequenceBreakers`] if more than
    ///   [`DRY_MAX_SEQ_BREAKERS`] remain after llama.cpp's own truncation to
    ///   [`DRY_MAX_BREAKER_LEN`] bytes and de-duplication.
    /// - [`SamplerInitError::NullReturn`] if llama.cpp returns null.
    pub fn try_dry(
        &self,
        model: &LlamaModel,
        multiplier: f32,
        base: f32,
        allowed_length: i32,
        penalty_last_n: i32,
        seq_breakers: impl IntoIterator<Item = impl AsRef<[u8]>>,
    ) -> Result<Self, SamplerInitError> {
        let seq_breakers: Vec<CString> = seq_breakers
            .into_iter()
            .map(|s| {
                CString::new(s.as_ref()).map_err(|source| SamplerInitError::InteriorNul {
                    field: "sequence breaker",
                    source,
                })
            })
            .collect::<Result<_, _>>()?;

        let distinct = distinct_dry_breakers(&seq_breakers);
        if distinct > DRY_MAX_SEQ_BREAKERS {
            return Err(SamplerInitError::TooManyDrySequenceBreakers {
                distinct,
                max: DRY_MAX_SEQ_BREAKERS,
            });
        }

        let mut seq_breaker_pointers: Vec<*const c_char> =
            seq_breakers.iter().map(|s| s.as_ptr()).collect();

        // SAFETY: the vocab outlives the call, and the pointer array is valid
        // for `len` entries whose backing `CString`s are still alive here.
        let sampler = unsafe {
            llama_sampler_init_dry(
                model.get_vocab().vocab.as_ref(),
                multiplier,
                base,
                allowed_length,
                penalty_last_n,
                seq_breaker_pointers.as_mut_ptr(),
                seq_breaker_pointers.len(),
            )
        };

        NonNull::new(sampler)
            .map(|sampler| Self { sampler })
            .ok_or(SamplerInitError::NullReturn("dry"))
    }

    /// Penalizes tokens for being present in the context.
    ///
    /// Parameters:
    /// - `n_vocab`: [`LlamaModel::n_vocab`]
    /// - `penalty_last_n`: last n tokens to penalize (0 = disable penalty)
    /// - `penalty_repeat`: repetition penalty (must be > 0.0, 1.0 = disabled)
    /// - `penalty_freq`: frequency penalty (must be finite, 0.0 = disabled)
    /// - `penalty_present`: presence penalty (must be finite, 0.0 = disabled)
    ///
    /// If `penalty_last_n` is `0`, or every penalty sits at its disabled value,
    /// llama.cpp returns a no-op sampler named `"?penalties"` — see
    /// [`Self::name`].
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[allow(clippy::too_many_arguments)]
    #[must_use]
    pub fn penalties(
        n_vocab: i32,
        penalty_last_n: i32,
        penalty_repeat: f32,
        penalty_freq: f32,
        penalty_present: f32,
    ) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe {
            llama_sampler_init_penalties(
                n_vocab,
                penalty_last_n,
                penalty_repeat,
                penalty_freq,
                penalty_present,
            )
        };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Same as [`Self::penalties`] with sensible defaults:
    /// `penalty_freq = 0.0` and `penalty_present = 0.0`.
    ///
    /// Parameters:
    /// - `n_vocab`: [`LlamaModel::n_vocab`]
    /// - `penalty_last_n`: last n tokens to penalize (0 = disable)
    /// - `penalty_repeat`: repetition penalty (must be > 0.0, 1.0 = disabled)
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn penalties_simple(n_vocab: i32, penalty_last_n: i32, penalty_repeat: f32) -> Self {
        Self::penalties(
            n_vocab,
            #[allow(clippy::cast_precision_loss)]
            {
                penalty_last_n
            },
            #[allow(clippy::cast_precision_loss)]
            {
                penalty_repeat
            },
            #[allow(clippy::cast_precision_loss)]
            {
                0.0_f32
            },
            #[allow(clippy::cast_precision_loss)]
            {
                0.0_f32
            },
        )
    }

    /// Mirostat 1.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Parameters:
    /// - `n_vocab`: [`LlamaModel::n_vocab`]
    /// - `seed`: Seed to initialize random generation with.
    /// - `tau`: The target cross-entropy (or surprise) value you want to achieve for the
    ///   generated text. A higher value corresponds to more surprising or less predictable text,
    ///   while a lower value corresponds to less surprising or more predictable text.
    /// - `eta`: The learning rate used to update `mu` based on the error between the target and
    ///   observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
    ///   updated more quickly, while a smaller learning rate will result in slower updates.
    /// - `m`: The number of tokens considered in the estimation of `s_hat`. This is an arbitrary
    ///   value that is used to calculate `s_hat`, which in turn helps to calculate the value of `k`.
    ///   In the paper, they use `m = 100`, but you can experiment with different values to see how
    ///   it affects the performance of the algorithm.
    #[must_use]
    pub fn mirostat(n_vocab: i32, seed: u32, tau: f32, eta: f32, m: i32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Mirostat 2.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Parameters:
    /// - `seed`: Seed to initialize random generation with.
    /// - `tau`: The target cross-entropy (or surprise) value you want to achieve for the
    ///   generated text. A higher value corresponds to more surprising or less predictable text,
    ///   while a lower value corresponds to less surprising or more predictable text.
    /// - `eta`: The learning rate used to update `mu` based on the error between the target and
    ///   observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
    ///   updated more quickly, while a smaller learning rate will result in slower updates.
    #[must_use]
    pub fn mirostat_v2(seed: u32, tau: f32, eta: f32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_mirostat_v2(seed, tau, eta) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Selects a token at random based on each token's probabilities.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn dist(seed: u32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_dist(seed) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Selects the most likely token.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Example:
    /// ```rust
    /// use llama_cpp_4::token::{
    ///    LlamaToken,
    ///    data::LlamaTokenData,
    ///    data_array::LlamaTokenDataArray
    /// };
    /// use llama_cpp_4::sampling::LlamaSampler;
    ///
    /// let mut data_array = LlamaTokenDataArray::new(vec![
    ///     LlamaTokenData::new(LlamaToken(0), 0., 0.),
    ///     LlamaTokenData::new(LlamaToken(1), 1., 0.),
    /// ], false);
    ///
    /// data_array.apply_sampler(&mut LlamaSampler::greedy());
    ///
    /// assert_eq!(data_array.data.len(), 2);
    /// assert_eq!(data_array.selected_token(), Some(LlamaToken(1)));
    /// ```
    #[must_use]
    pub fn greedy() -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_greedy() };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Top-N sigma sampling.
    ///
    /// Keeps tokens within N standard deviations of the maximum logit.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn top_n_sigma(n: f32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_top_n_sigma(n) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Adaptive P sampling.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Parameters
    /// - `target`: Target probability.
    /// - `decay`: Decay rate.
    /// - `seed`: Random seed.
    #[must_use]
    pub fn adaptive_p(target: f32, decay: f32, seed: u32) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_adaptive_p(target, decay, seed) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Logit bias sampler.
    ///
    /// Applies additive bias to specific token logits before sampling.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Parameters
    /// - `n_vocab`: Number of tokens in the vocabulary ([`LlamaModel::n_vocab`]).
    /// - `biases`: Slice of `(token_id, bias)` pairs.
    #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
    #[must_use]
    pub fn logit_bias(n_vocab: i32, biases: &[(LlamaToken, f32)]) -> Self {
        let logit_biases: Vec<llama_logit_bias> = biases
            .iter()
            .map(|(token, bias)| llama_logit_bias {
                token: token.0,
                bias: *bias,
            })
            .collect();

        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe {
            llama_sampler_init_logit_bias(n_vocab, logit_biases.len() as i32, logit_biases.as_ptr())
        };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Infill sampler.
    ///
    /// Reorders token probabilities for fill-in-the-middle tasks.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn infill(model: &LlamaModel) -> Self {
        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe { llama_sampler_init_infill(model.get_vocab().vocab.as_ref()) };
        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }

    /// Get the seed of the sampler.
    ///
    /// Returns `LLAMA_DEFAULT_SEED` if the sampler is not seeded.
    #[must_use]
    pub fn get_seed(&self) -> u32 {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_sampler_get_seed(self.sampler.as_ptr()) }
    }

    /// Get the name of the sampler.
    ///
    /// # Disabled samplers
    ///
    /// When a constructor is handed parameters that make it a no-op — e.g.
    /// [`Self::temp`] with `1.0`, or [`Self::penalties`] with `penalty_last_n =
    /// 0` — llama.cpp does not build that sampler. It substitutes an identity
    /// sampler whose name carries a `?` prefix (`"?temp"`, `"?penalties"`).
    /// Construction still succeeds, so this name is the only signal that the
    /// sampler will not do anything. Affects `temp`, `temp_ext`, `top_k`,
    /// `top_p`, `min_p`, `typical`, `xtc`, `top_n_sigma`, `dry`, and
    /// `penalties`.
    ///
    /// # Panics
    ///
    /// Panics if the name is not valid UTF-8.
    #[must_use]
    pub fn name(&self) -> String {
        // SAFETY: the pointer is non-null and points at a NUL-terminated string the
        // library owns, which outlives the borrow taken here.
        let c_str = unsafe { llama_sampler_name(self.sampler.as_ptr()) };
        // SAFETY: the pointer is non-null and points at a NUL-terminated string the
        // library owns, which outlives the borrow taken here.
        let c_str = unsafe { std::ffi::CStr::from_ptr(c_str) };
        c_str
            .to_str()
            .expect("sampler name is not valid UTF-8")
            .to_owned()
    }

    /// Reset the sampler state (e.g. grammar, repetition penalties).
    pub fn reset(&mut self) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_sampler_reset(self.sampler.as_ptr()) }
    }

    /// Get the number of samplers in a chain.
    ///
    /// Returns 0 if this sampler is not a chain.
    #[must_use]
    pub fn chain_n(&self) -> i32 {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_sampler_chain_n(self.sampler.as_ptr()) }
    }

    /// Remove and return the sampler at position `i` from a chain.
    ///
    /// The returned sampler is owned by the caller and will be freed on drop.
    ///
    /// # Panics
    ///
    /// Panics if `i` is out of range or if llama.cpp returns a null pointer.
    #[must_use]
    pub fn chain_remove(&mut self, i: i32) -> Self {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let sampler = unsafe { llama_sampler_chain_remove(self.sampler.as_ptr(), i) };
        Self {
            sampler: NonNull::new(sampler).expect("chain_remove returned null"),
        }
    }

    /// Grammar sampler with lazy activation via regex patterns.
    ///
    /// The grammar is only activated when one of the trigger patterns or trigger tokens matches.
    ///
    /// # Panics
    /// - If `grammar_str` or `grammar_root` contain null bytes.
    /// - If any trigger pattern contains null bytes.
    /// - If llama.cpp returns a null pointer.
    ///
    /// Prefer [`Self::try_grammar_lazy_patterns`] when any of those strings came
    /// from a request rather than your own source.
    #[must_use]
    pub fn grammar_lazy_patterns(
        model: &LlamaModel,
        grammar_str: &str,
        grammar_root: &str,
        trigger_patterns: &[&str],
        trigger_tokens: &[LlamaToken],
    ) -> Self {
        Self::try_grammar_lazy_patterns(
            model,
            grammar_str,
            grammar_root,
            trigger_patterns,
            trigger_tokens,
        )
        .unwrap_or_else(|e| panic!("building lazy grammar sampler: {e}"))
    }

    /// [`Self::grammar_lazy_patterns`], reporting failures instead of panicking.
    ///
    /// # Errors
    ///
    /// [`SamplerInitError::InteriorNul`] for a NUL in the grammar, the rule name
    /// or any trigger pattern, or [`SamplerInitError::NullReturn`] if llama.cpp
    /// rejected the grammar.
    pub fn try_grammar_lazy_patterns(
        model: &LlamaModel,
        grammar_str: &str,
        grammar_root: &str,
        trigger_patterns: &[&str],
        trigger_tokens: &[LlamaToken],
    ) -> Result<Self, SamplerInitError> {
        let grammar_str = c_string("grammar_str", grammar_str)?;
        let grammar_root = c_string("grammar_root", grammar_root)?;
        let pattern_cstrings = trigger_patterns
            .iter()
            .map(|pattern| c_string("trigger_patterns", pattern))
            .collect::<Result<Vec<CString>, _>>()?;
        let mut pattern_ptrs: Vec<*const c_char> =
            pattern_cstrings.iter().map(|s| s.as_ptr()).collect();

        // SAFETY: every pointer is borrowed from a `CString`/slice that outlives
        // the call, and the two lengths describe exactly those buffers.
        // `trigger_tokens` is `#[repr(transparent)]` over `llama_token`, so the
        // cast is a layout no-op. llama.cpp copies what it needs.
        let sampler = unsafe {
            llama_sampler_init_grammar_lazy_patterns(
                model.get_vocab().vocab.as_ref(),
                grammar_str.as_ptr(),
                grammar_root.as_ptr(),
                pattern_ptrs.as_mut_ptr(),
                pattern_ptrs.len(),
                trigger_tokens.as_ptr().cast(),
                trigger_tokens.len(),
            )
        };
        NonNull::new(sampler)
            .map(|sampler| Self { sampler })
            .ok_or(SamplerInitError::NullReturn("lazy grammar"))
    }

    /// Clone this sampler.
    ///
    /// Creates an independent copy of this sampler with the same state.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub fn clone_sampler(&self) -> Self {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let sampler = unsafe { llama_sampler_clone(self.sampler.as_ptr()) };
        Self {
            sampler: NonNull::new(sampler).expect("sampler_clone returned null"),
        }
    }

    /// Copy mutable state from `src` into this sampler, in place.
    ///
    /// Unlike [`Self::clone_sampler`], which allocates a new sampler, this
    /// overwrites the state of an existing one and so is the cheap way to
    /// rewind a sampler to a checkpoint in a loop. Added upstream in
    /// llama.cpp `b10470` (`llama_sampler_copy`).
    ///
    /// # Safety and preconditions
    ///
    /// llama.cpp requires `src` and `self` to be **the same sampler type with
    /// the same configuration** — e.g. two `dist` samplers, or two chains built
    /// the same way. Copying between mismatched samplers is undefined
    /// behaviour upstream and is not checked here, so treat the pairing as the
    /// caller's contract. A sampler produced by `src.clone_sampler()` always
    /// satisfies it.
    pub fn copy_state_from(&mut self, src: &Self) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_sampler_copy(src.sampler.as_ptr(), self.sampler.as_ptr()) }
    }

    /// Print sampler performance data.
    pub fn perf_print(&self) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_cpp_sys_4::llama_perf_sampler_print(self.sampler.as_ptr()) }
    }

    /// Reset sampler performance counters.
    pub fn perf_reset(&mut self) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_cpp_sys_4::llama_perf_sampler_reset(self.sampler.as_ptr()) }
    }

    /// Get sampler performance data.
    #[must_use]
    pub fn perf_data(&self) -> llama_cpp_sys_4::llama_perf_sampler_data {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { llama_cpp_sys_4::llama_perf_sampler(self.sampler.as_ptr()) }
    }

    /// Get a non-owning reference to the `i`th sampler in a chain.
    ///
    /// # Safety
    ///
    /// The returned pointer is owned by the chain. Do not free it or use it
    /// after the chain is dropped or modified.
    #[must_use]
    pub unsafe fn chain_get_ptr(&self, i: i32) -> *mut llama_sampler {
        llama_cpp_sys_4::llama_sampler_chain_get(self.sampler.as_ptr(), i)
    }

    /// Create a sampler from a raw interface and context.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `iface` and `ctx` are valid and that the
    /// interface functions properly manage the context lifetime.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    #[must_use]
    pub unsafe fn from_raw(
        iface: *mut llama_cpp_sys_4::llama_sampler_i,
        ctx: llama_cpp_sys_4::llama_sampler_context_t,
    ) -> Self {
        let sampler = llama_cpp_sys_4::llama_sampler_init(iface, ctx);
        Self {
            sampler: NonNull::new(sampler).expect("sampler_init returned null"),
        }
    }

    /// Adopt a raw `llama_sampler *`, taking ownership.
    ///
    /// # Safety
    ///
    /// `raw` must come from a llama.cpp entry point documented as returning a
    /// sampler the caller owns and releases with `llama_sampler_free`, and must
    /// not be owned by anything else — [`Drop`] frees it.
    pub(crate) unsafe fn from_raw_ptr(raw: NonNull<llama_sampler>) -> Self {
        Self { sampler: raw }
    }

    /// The underlying `llama_sampler *`. Borrowed; still owned by `self`.
    pub(crate) fn as_ptr(&self) -> *mut llama_sampler {
        self.sampler.as_ptr()
    }

    /// Creates a new instance of `LlamaSampler` with common sampling parameters.
    ///
    /// This function initializes a `LlamaSampler` using default values from `common_sampler_params`
    /// and configures it with common settings such as `top_k`, `top_p`, `temperature`, and `seed` values.
    ///
    /// # Panics
    ///
    /// Panics if llama.cpp returns a null pointer.
    ///
    /// # Returns
    /// A `LlamaSampler` instance configured with the common sampling parameters.
    #[must_use]
    pub fn common() -> Self {
        let params = common_sampler_params::default();

        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let sampler = unsafe {
            let mut sparams = llama_sampler_chain_default_params();
            sparams.no_perf = false;

            let smpl = llama_sampler_chain_init(sparams);

            llama_sampler_chain_add(smpl, llama_sampler_init_top_k(params.top_k));
            llama_sampler_chain_add(
                smpl,
                #[allow(clippy::cast_sign_loss)]
                llama_sampler_init_top_p(params.top_p, params.min_keep as usize),
            );
            llama_sampler_chain_add(smpl, llama_sampler_init_temp(params.temp));
            #[allow(clippy::cast_sign_loss)]
            llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.seed));

            smpl
        };

        Self {
            sampler: NonNull::new(sampler).unwrap(),
        }
    }
}

impl Drop for LlamaSampler {
    fn drop(&mut self) {
        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
        // afterwards.
        unsafe {
            llama_sampler_free(self.sampler.as_ptr());
        }
    }
}