API Reference

Base Sampler

class llm_samplers.base.BaseSampler

Bases: ABC

Base class for all samplers.

__init__()

abstract sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Sample from the model using the specific sampling strategy.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments specific to the sampler

Returns:

Generated token IDs

Return type:

torch.Tensor

Beam Search Sampler

class llm_samplers.beam.BeamSearchSampler

Bases: BaseSampler

Beam search sampler that maintains the top k most promising sequences.

__init__(beam_width=5)

Initialize the beam search sampler.

Parameters:

beam_width (int) – Number of sequences to keep at each step

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Sample from the model using beam search.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments specific to the sampler

Returns:

Generated token IDs

Return type:

torch.Tensor

Temperature Sampler

class llm_samplers.temperature.TemperatureSampler

Bases: BaseSampler

Temperature sampling adjusts the “sharpness” of the probability distribution.

• Low temperature (<1.0): More deterministic, picks high-probability tokens

• High temperature (>1.0): More random, flatter distribution

__init__(temperature=1.0)

Initialize the temperature sampler.

Parameters:

temperature (float) – Temperature value to control sampling sharpness

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using temperature sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

Top-K Sampler

class llm_samplers.top_k.TopKSampler

Bases: BaseSampler

Top-K sampling considers only the ‘k’ most probable tokens. All other tokens are set to have zero probability.

__init__(k=50)

Initialize the Top-K sampler.

Parameters:

k (int) – Number of top tokens to consider

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using Top-K sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

Top-P (Nucleus) Sampler

class llm_samplers.top_p.TopPSampler

Bases: BaseSampler

Top-P (Nucleus) sampling selects the smallest set of tokens whose cumulative probability exceeds threshold ‘p’.

__init__(p=0.9)

Initialize the Top-P sampler.

Parameters:

p (float) – Probability threshold for nucleus sampling

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using Top-P sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

Min-P Sampler

class llm_samplers.min_p.MinPSampler

Bases: BaseSampler

Min-P sampling dynamically adjusts the sampling pool size based on the probability of the most likely token.

__init__(min_p=0.05)

Initialize the Min-P sampler.

Parameters:

min_p (float) – Minimum probability threshold for tokens

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using Min-P sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

Anti-Slop Sampler

class llm_samplers.anti_slop.AntiSlopSampler

Bases: BaseSampler

Anti-Slop sampling down-weights probabilities at word & phrase level. Uses backtracking to retry with adjusted token probabilities if it encounters a disallowed word/phrase.

__init__(disallowed_tokens=None, disallowed_phrases=None, penalty=0.5, max_retries=3)

Initialize the Anti-Slop sampler.

Parameters:

• disallowed_tokens (Optional[Set[int]]) – Set of token IDs to penalize

• disallowed_phrases (Optional[List[List[int]]]) – List of token ID sequences to penalize

• penalty (float) – Probability penalty factor (0-1)

• max_retries (int) – Maximum number of retries for backtracking

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using Anti-Slop sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

XTC (Exclude Top Choices) Sampler

class llm_samplers.xtc.XTCSampler

Bases: BaseSampler

XTC (Exclude Top Choices) sampling “turns truncation on its head.” Instead of pruning low-probability tokens, XTC targets the most probable ones under certain conditions to enhance creativity.

__init__(top_k=5, exclusion_threshold=0.3, min_probability=0.1)

Initialize the XTC sampler.

Parameters:

• top_k (int) – Number of top tokens to potentially exclude

• exclusion_threshold (float) – Probability threshold for exclusion

• min_probability (float) – Minimum probability to consider for exclusion

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Generate text using XTC sampling.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments

Returns:

Generated token IDs

Return type:

torch.Tensor

QAlign Sampler

class llm_samplers.qalign.QAlignSampler

Bases: BaseSampler

QAlign sampler that uses MCMC to improve model outputs based on a reward model.

This implementation is based on the paper: “Sample, Don’t Search: Rethinking Test-Time Alignment for Language Models” by Faria et al. (2024) Paper: https://arxiv.org/abs/2504.03790

QAlign uses Markov Chain Monte Carlo (MCMC) to align model outputs at test time without requiring model fine-tuning. It converges to sampling from the optimal aligned distribution as test-time compute scales.

__init__(reward_model, num_steps=10, temperature=1.0, beta=1.0, proposal_temp=1.0)

Initialize QAlign sampler.

Parameters:

• reward_model (Module) – The reward model to use for scoring outputs

• num_steps (int) – Number of MCMC steps to perform

• temperature (float) – Temperature for sampling from the base model

• beta (float) – Inverse temperature for the reward model (higher = more emphasis on reward)

• proposal_temp (float) – Temperature for the proposal distribution

sample(model, input_ids, max_length=100, num_return_sequences=1, **kwargs)

Sample from the model using QAlign.

Parameters:

• model (Any) – The language model to sample from

• input_ids (Tensor) – Input token IDs

• max_length (int) – Maximum length of the generated sequence

• num_return_sequences (int) – Number of sequences to return

• **kwargs – Additional arguments specific to the sampler

Returns:

Generated token IDs

Return type:

torch.Tensor