token-sampling 0.1.0

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use token_sampling::{
    unit, Candidates, ChainConfig, Dist, Dry, DynTemperature, Error, Greedy, MirostatV1,
    MirostatV2, Penalties, Rng, Sampler, SamplerChain, SamplerCtx, SplitMix64, TailFree, TapeRng,
    Temperature, TopK, Xtc,
};

#[test]
fn defect_d1_dry_penalizes_repeated_extensions() {
    let mut candidates = Candidates::new();
    candidates.load(&[-1.386_294_4; 4]).unwrap();
    let mut rng = TapeRng::new(&[1_u64 << 62]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[0, 1, 2, 0, 1],
    };
    Dry {
        multiplier: 1.0,
        base: 1.1,
        allowed_length: 2,
        last_n: 5,
        breakers: vec![],
    }
    .apply(&mut candidates, &mut context)
    .unwrap();
    Dist.apply(&mut candidates, &mut context).unwrap();
    let penalized = candidates
        .as_slice()
        .iter()
        .find(|candidate| candidate.id == 2)
        .unwrap();
    assert!((penalized.p - 0.109_232).abs() < 1e-5);
}

#[test]
fn defect_d2_xtc_excludes_leading_choices() {
    let logits = [0.4_f32.ln(), 0.3_f32.ln(), 0.2_f32.ln(), 0.1_f32.ln()];
    let mut candidates = Candidates::new();
    candidates.load(&logits).unwrap();
    let mut rng = TapeRng::new(&[1_u64 << 63]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[],
    };
    Xtc {
        probability: 0.99,
        threshold: 0.19,
        min_keep: 0,
    }
    .apply(&mut candidates, &mut context)
    .unwrap();
    assert_eq!(
        candidates
            .as_slice()
            .iter()
            .map(|candidate| candidate.id)
            .collect::<Vec<_>>(),
        vec![2, 3]
    );
}

#[test]
fn xtc_zero_threshold_preserves_the_only_unmasked_candidate() {
    fn sample(logits: &mut [f32]) -> Result<u32, Error> {
        let mut chain = SamplerChain::new()
            .push(Xtc {
                probability: 1.0,
                threshold: 0.0,
                min_keep: 0,
            })
            .push(Dist);
        let mut rng = SplitMix64::new(7);
        chain.sample(logits, &[], &mut rng)
    }

    let outcomes = [
        sample(&mut [2.0, f32::NEG_INFINITY]),
        sample(&mut [f32::NEG_INFINITY, 2.0]),
        sample(&mut [f32::NEG_INFINITY, 2.0, f32::NEG_INFINITY]),
    ];
    assert_eq!(outcomes, [Ok(0), Ok(1), Ok(1)]);
}

#[test]
fn defect_d3_zero_dynamic_range_matches_fixed_temperature() {
    let logits = [2.0, 1.0, -0.5];
    let mut fixed = Candidates::new();
    let mut dynamic = Candidates::new();
    fixed.load(&logits).unwrap();
    dynamic.load(&logits).unwrap();
    let mut rng = TapeRng::new(&[]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[],
    };
    Temperature { t: 0.7 }
        .apply(&mut fixed, &mut context)
        .unwrap();
    DynTemperature {
        t: 0.7,
        delta: 0.0,
        exponent: 1.0,
    }
    .apply(&mut dynamic, &mut context)
    .unwrap();
    assert_eq!(fixed.as_slice(), dynamic.as_slice());
}

#[test]
fn defect_d4_top_k_keeps_the_largest_prefix() {
    let logits = (0..4096).map(|value| value as f32).collect::<Vec<_>>();
    let mut candidates = Candidates::new();
    candidates.load(&logits).unwrap();
    let mut rng = TapeRng::new(&[]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[],
    };
    TopK { k: 40 }.apply(&mut candidates, &mut context).unwrap();
    assert_eq!(candidates.as_slice().len(), 40);
    assert_eq!(candidates.as_slice()[0].id, 4095);
    assert_eq!(candidates.as_slice()[39].id, 4056);
}

#[test]
fn defect_d6_repeat_below_one_rewards_history_tokens() {
    let mut candidates = Candidates::new();
    candidates.load(&[2.0, 1.0]).unwrap();
    let mut rng = TapeRng::new(&[]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[0],
    };
    Penalties {
        last_n: 1,
        repeat: 0.5,
        freq: 0.0,
        presence: 0.0,
    }
    .apply(&mut candidates, &mut context)
    .unwrap();
    assert_eq!(candidates.as_slice()[0].logit, 4.0);
}

#[test]
fn defect_d7_mirostat_v1_requires_a_vocabulary_size() {
    let sampler = MirostatV1::new(32);
    assert_eq!(sampler.n_vocab, 32);
    assert_eq!(sampler.mu(), 10.0);
}

#[test]
fn defect_d8_custom_rng_needs_one_method() {
    struct Fixed(u64);
    impl Rng for Fixed {
        fn next_u64(&mut self) -> u64 {
            self.0
        }
    }
    let mut rng = Fixed(0);
    assert_eq!(rng.next_u64(), 0);
}

#[test]
fn defect_d9_iterator_logits_fit_a_shared_dynamic_chain() {
    fn assert_send<T: Send>() {}
    assert_send::<SamplerChain>();
    let mut chain = SamplerChain::new();
    chain.push_dyn(Box::new(Temperature { t: 0.8 }));
    let sampler = Arc::new(Mutex::new(chain.push(Dist)));
    let mut logits = [0.1, 0.8, -0.2].into_iter().collect::<Vec<f32>>();
    let mut rng = SplitMix64::new(3);
    let token = sampler
        .lock()
        .unwrap()
        .sample(&mut logits, &[1, 2], &mut rng)
        .unwrap();
    assert!(token < 3);
}

#[test]
fn defect_d10_invalid_logits_return_typed_errors_at_load() {
    let mut candidates = Candidates::new();
    assert_eq!(candidates.load(&[]), Err(Error::EmptyLogits));
    assert_eq!(
        candidates.load(&[0.0, f32::NAN]),
        Err(Error::InvalidLogit { index: 1 })
    );
    assert_eq!(candidates.load(&[f32::NEG_INFINITY]), Err(Error::AllMasked));
}

#[test]
fn defect_d11_fixture_counts_are_pinned() {
    let value: serde_json::Value =
        serde_json::from_str(include_str!("data/conformance-vectors.json")).unwrap();
    assert_eq!(value["vector_count"], 140);
    assert_eq!(value["chain_count"], 10);
    assert_eq!(value["mirostat_sequence_count"], 8);
    assert_eq!(value["mirostat_step_count"], 48);
    assert_eq!(value["llama_queue_count"], 24);
}

#[test]
fn rng_conversion_and_consumption_follow_the_fixed_contract() {
    let mut maximum = TapeRng::new(&[u64::MAX]);
    assert_eq!(unit(&mut maximum), 1.0 - 1.0 / 9_007_199_254_740_992.0);

    let tape = [0, 1_u64 << 63];
    let mut rng = TapeRng::new(&tape);
    let mut candidates = Candidates::new();
    candidates.load(&[1.0, 0.0]).unwrap();
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[],
    };
    Xtc {
        probability: 0.0,
        threshold: 0.1,
        min_keep: 0,
    }
    .apply(&mut candidates, &mut context)
    .unwrap();
    assert_eq!(context.rng.next_u64(), 0);
    Dist.apply(&mut candidates, &mut context).unwrap();
    assert_eq!(rng.remaining(), 0);

    let singleton_tape = [17];
    let mut singleton_rng = TapeRng::new(&singleton_tape);
    let mut singleton = Candidates::new();
    singleton.load(&[1.0]).unwrap();
    Dist.apply(
        &mut singleton,
        &mut SamplerCtx {
            rng: &mut singleton_rng,
            history: &[],
        },
    )
    .unwrap();
    assert_eq!(singleton_rng.remaining(), 0);
}

#[test]
fn reset_restores_mirostat_targets() {
    let mut version_one = MirostatV1::new(32);
    version_one.tau = 3.0;
    version_one.reset();
    assert_eq!(version_one.mu(), 6.0);

    let mut version_two = MirostatV2::new(4.0, 0.1);
    version_two.tau = 2.0;
    version_two.reset();
    assert_eq!(version_two.mu(), 4.0);
}

#[test]
fn dry_ignores_empty_breaker_sequences() {
    let mut candidates = Candidates::new();
    candidates.load(&[0.0, 0.0, 0.0]).unwrap();
    let mut rng = TapeRng::new(&[]);
    let mut context = SamplerCtx {
        rng: &mut rng,
        history: &[0, 1, 2, 0, 1, 2, 0],
    };
    let mut sampler = Dry {
        multiplier: 1.0,
        base: 1.75,
        allowed_length: 2,
        last_n: 64,
        breakers: vec![vec![2, 3], vec![], vec![1]],
    };
    sampler.apply(&mut candidates, &mut context).unwrap();
    assert_eq!(candidates.as_slice().len(), 3);
    assert_eq!(sampler.breakers, vec![vec![2, 3], vec![], vec![1]]);
}

#[test]
fn tail_free_zero_keeps_one_candidate() {
    let mut candidates = Candidates::new();
    candidates.load(&[3.0, 2.0, 1.0]).unwrap();
    let mut rng = TapeRng::new(&[]);
    TailFree {
        z: 0.0,
        min_keep: 0,
    }
    .apply(
        &mut candidates,
        &mut SamplerCtx {
            rng: &mut rng,
            history: &[],
        },
    )
    .unwrap();
    assert_eq!(candidates.as_slice().len(), 1);
}

#[test]
fn finite_inputs_survive_extreme_logit_transforms() {
    assert_sample_succeeds(
        SamplerChain::llama_default(ChainConfig::default()),
        &[3.0e38, 3.0e38, 0.0],
        &[],
    );
    assert_sample_succeeds(
        SamplerChain::new()
            .push(Temperature { t: 1.0e-40 })
            .push(Dist),
        &[100.0, 1.0],
        &[],
    );
    assert_sample_succeeds(
        SamplerChain::new()
            .push(Temperature { t: f32::INFINITY })
            .push(Dist),
        &[1.0, f32::NEG_INFINITY, 2.0],
        &[],
    );
    assert_sample_succeeds(
        SamplerChain::new()
            .push(Penalties {
                last_n: i32::MAX,
                repeat: 1.0,
                freq: -3.0e38,
                presence: 0.0,
            })
            .push(Dist),
        &[1.0, 0.0],
        &[0, 0],
    );
}

#[test]
fn greedy_never_selects_a_masked_candidate_after_infinite_temperature() {
    let mut chain = SamplerChain::new()
        .push(Temperature { t: f32::INFINITY })
        .push(Greedy);
    let mut logits = [f32::NEG_INFINITY, 5.0];
    let mut rng = SplitMix64::new(1);
    assert_eq!(chain.sample(&mut logits, &[], &mut rng), Ok(1));
}

#[test]
fn penalty_counting_finishes_a_large_reverse_window_in_bounded_time() {
    const WINDOW: usize = 250_000;
    let history = (0..WINDOW as u32).rev().collect::<Vec<_>>();
    let mut candidates = Candidates::new();
    candidates.load(&[0.0]).unwrap();
    let mut rng = TapeRng::new(&[]);
    let start = Instant::now();
    Penalties {
        last_n: i32::MAX,
        repeat: 1.0,
        freq: 1.0,
        presence: 0.0,
    }
    .apply(
        &mut candidates,
        &mut SamplerCtx {
            rng: &mut rng,
            history: &history,
        },
    )
    .unwrap();
    assert!(
        start.elapsed() < Duration::from_secs(2),
        "large penalty window exceeded the two-second bound"
    );
}

#[test]
fn public_chain_and_context_types_implement_debug() {
    let chain = SamplerChain::new();
    assert!(format!("{chain:?}").contains("stage_count"));
    let mut rng = SplitMix64::new(1);
    let context = SamplerCtx {
        rng: &mut rng,
        history: &[1, 2],
    };
    assert!(format!("{context:?}").contains("history"));
}

#[test]
fn chain_without_a_terminal_stage_reports_no_selection() {
    let mut chain = SamplerChain::new();
    let mut logits = [0.0, 1.0];
    let mut rng = SplitMix64::new(1);
    assert_eq!(
        chain.sample(&mut logits, &[], &mut rng),
        Err(Error::NoSelection)
    );
}

fn assert_sample_succeeds(mut chain: SamplerChain, logits: &[f32], history: &[u32]) {
    let mut logits = logits.to_vec();
    let mut rng = SplitMix64::new(7);
    let token = chain.sample(&mut logits, history, &mut rng).unwrap();
    assert!((token as usize) < logits.len());
}