A path to adding more types of samplers, including Mirostat 1 and 2

[KerfuffleV2]

I've been working on a Rust implementation of the samplers llama.cpp supports, referencing their implementation. You can see the project here: https://github.com/KerfuffleV2/llm-samplers / https://crates.io/crates/llm-samplers

This is very lightly tested at the moment, all I can really say is my versions of samplers pass the existing llama.cpp sampler tests and when I tried some of them (including Mirostat 1 and 2) in my little RWKV project it seemed to produce pretty decent results.

The advanced samplers like Mirostat can really make a noticeable difference in generation quality.

Here's what using some of these samplers can look like:

edit: Updated for v0.0.3

use std::sync::{Arc, RwLock};

use anyhow::Result;
use rand::rngs::StdRng;

use llm_samplers::prelude::*;

pub fn test_chain() -> Result<()> {
    let testvals = vec![0.1f32, 0.2, 0.3, 0.4];
    let mut logits = Logits::try_from_iter(testvals.clone())?;
    let mut logits2 = logits.clone();

    // This is an `Arc<RwLock<Vec>>` so the samplers that need it can hold a copy of the `Arc`
    // to read and we can also update the `Vec` in between sampling.
    let last_tokens = Arc::new(RwLock::new(vec![]));

    let mut sc = SamplerChain::new();
    sc
        // Apply a repetition penalty.
        .push_sampler(SampleRepetition::new(1.1, 64, last_tokens.clone()))

        // Apply frequency and presence penalities.
        .push_sampler(SampleFreqPresence::new(0.05, 0.1, 64, last_tokens.clone()))

        // Apply temperature of 0.8 to the logits.
        .push_sampler(SampleTemperature::new(0.8))

        // Bias token ID 3 to minus infinity. In other words: never select it.
        .push_sampler(SampleFlatBias::new(&[(3, f32::NEG_INFINITY)]))
        // Use a Mirostat v1 sampler to select a token at the end.
        .push_sampler(SampleMirostat1::<u32, f32, StdRng>::new(
            4,
            5.0,
            0.1,
            60,
            10.0,
            Box::new(RngBox::new_seedable(Some(123))),
        ));

    // Put a value into `last_tokens`, this simulates us having already picked
    // that token (3) previously.
    last_tokens.write().unwrap().push(3);
    assert_eq!(sc.sample_token(&mut logits)?, Some(2));

    // Now add the last selected token to the list.
    last_tokens.write().unwrap().push(2);

    // And pick the next one. *Important*: Note that we don't reuse `logits`.
    // This is because `logits` already has all the filtering/sorting/permutation
    // from the previous sample call applied to it.
    assert_eq!(sc.sample_token(&mut logits2)?, Some(1));
    Ok(())
}

Taken from a test branch in smolrsrwkv.

If there's interest, one thing I could really use help with is tests. Either PRs adding tests or just telling me given a set of logits + certain parameters for a sampler, the result should be something. I'm really not very confident in the tests I "ported" from llama.cpp (and they don't include all samplers either).

[KerfuffleV2]

In v0.0.3, I've updated the interface to be a lot easier to use by adding in sampler chains (which also implement the Sampler) trait. It's now also possible to create a clonable thread-safe RNG.

It should now also be a lot more practical to provide last tokens to the samplers that need it (i.e. repetition penalty).