API Reference

---

Data

Container classes for raw LLM outputs and token-level data returned by the llama_interface layer.

Data container for LLM output.

class problm_solver.data.LLMOutputData(prompt, data)

Bases: object

Stores LLM output data and handles serialization.

Initialize.

Parameters:

• prompt (str) – The prompt used to generate the data.

• data (NDArray[Any]) – Array of string responses from the LLM.

write(fname)

Save data to a file in JSONL format.

Parameters:

fname (str) – File name to write to, absolute path.

Return type:

None

read(fname)

Read data to this object from a JSONL file.

Parameters:

fname (str) – File name to read from, absolute path.

Return type:

None

exception problm_solver.data.TokenProbError(data_class)

Bases: ValueError

Raise an error when number of tokens in data isn’t equal to probs.

Initialize error message.

Error is intended for dataclasses and other containers that hold tokens and their probabilities.

Parameters:

data_class (type) – Class type to be passed to error message.

class problm_solver.data.LLMTokenData(prompt, tokens, probs)

Bases: object

Store a single tokenized LLM response paired with per-token probabilities.

Initialize.

Parameters:

• prompt (str) – The prompt used to generate the response.

• tokens (list[str]) – The response as an ordered list of token strings.

• probs (list[float]) – The probability of each token at its position in the response. Must be the same length as tokens.

write(fname)

Save data to a file in JSON format.

Parameters:

fname (str) – File name to write to, absolute path.

Return type:

None

read(fname)

Read data from JSON file.

Parameters:

fname (str) – File name to read from, absolute path.

Return type:

None

class problm_solver.data.LLMNextTokenData(prompt, output_vec, top_k_tokens)

Bases: object

Store M most likely next tokens given context and output vector.

Initialize context and next token data.

Parameters:

• prompt (str) – The original user prompt.

• output_vec (list[int]) – The current token ID sequence (prompt + generated tokens so far).

• top_k_tokens (dict[str, float]) – Mapping of token string to log-probability for the top K candidate next tokens at this position.

write(fname)

Save data to a file in JSON format.

Parameters:

fname (str) – File name to write to, absolute path.

Return type:

None

read(fname)

Read from JSON file.

Parameters:

fname (str) – File name to read from, absolute path.

Return type:

None

class problm_solver.data.Hyperparams(alpha, top_k, top_p, max_tokens)

Bases: object

An assistant dataclass to LLMOutputDataFull with hyperparameter data included.

alpha: float

top_k: int | None

top_p: float | None

max_tokens: int

class problm_solver.data.LLMOutputDataFull(context, hyperparams, response_probabilities, response_topk, sampling_method, branch_sampler, _written=False)

Bases: object

A fuller dataclass than LLMOutputData.

Includes various parameters as strings, as well as top-k finalised token probabilities.

Parameters:

• context (list[str]) – Tokenized context provided to the model.

• hyperparams (Hyperparams) – Nested dataclass for hyperparameter storage.

• response_probabilities (tuple[list[str], list[float]]) – response tokens and their context-dependent selection probabilities after adjustment.

• response_topk (tuple[list[str], list[dict[str, float]]]) – For each token sampled, supplies the top K tokens and their probabilities.

• sampling_method (str) – User label for sampling method used.

• branch_sampler (str | None) – User label for branch sampler used, if sampling from the power distribution.

context: list[str]

hyperparams: Hyperparams

response_probabilities: tuple[list[str], list[float]]

response_topk: tuple[list[str], list[dict[str, float]]]

sampling_method: str

branch_sampler: str | None

write(fname)

Write data to JSON file.

Return type:

None

read(fname)

Read data from JSON file.

Return type:

None

---

Model interface

Thin wrapper around llama-cpp-python that loads a GGUF model and exposes generation / probability-query methods used by the rest of the library.

llama.cpp python interface for running local models.

class problm_solver.llama_interface.Model(fname, context, n_ctx=4096, logits_all=False, n_gpu_layers=0, c_api_copy_logits=True, *, rng=None, raw_completion=False, kv_unified=True)

Bases: object

Keeps a model instance and its context, with methods for querying the Llama instance.

Initialize Llama instance and store context.

The RAM cache capacity is derived from the model’s own metadata so that it can hold four full KV-cache states. One state covers the entire context window for all layers and KV heads at fp16 precision:

bytes_per_state = n_ctx × 2 × n_layers × n_kv_heads × head_dim × 2

Four states comfortably accommodates the save/restore pattern used by SamplePowerDist: the saved pre-branch snapshot, the current working state, and spare capacity for shared prefix entries.

Parameters:

• fname (str) – absolute path of the model .gguf file.

• context (str) – query that the model is initialized with.

• n_ctx (int) – context window size in tokens. Must be large enough to hold the formatted prompt plus max_tokens of generated output. Defaults to 4096, which comfortably fits MATH500 problems (~300 prompt tokens) and up to 2048 generated tokens.

• logits_all (bool) – whether or not probability logging is necessary in the Llama instance.

• n_gpu_layers (int) – number of GPU layers to pass to Llama. Required for GPU accelerated jobs, set to 0 otherwise.

• c_api_copy_logits (bool) – Controls whether ModelCBackend.last_logits() returns a copied array (default) or a zero-copy view of the C logits buffer.

• rng (RandomManager | Generator | int | None) – Optional random source. May be a numpy.random.Generator, integer seed, or RandomManager.

• raw_completion (bool) – When True, bypass the model’s chat template and feed self.context as a raw completion prompt. Required for base/pretrained models, which do not follow chat instructions and must be prompted with plain completion text.

• kv_unified (bool) – When True (default), construct the llama.cpp context with a unified KV-cache pool so that multi-sequence batched decoding works — required by query_branches_from_live_batch() (and thus the batched proposal path in SamplePowerDist). With the default partitioned pool the context supports only sequence id 0, so batched branch decoding would fail with llama_decode error -1. Disable only if you do not use batched branch decoding and want the stock llama.cpp context.

query_n_times(n)

Query the LLM with the same context N times, return the output.

Parameters:

n (int) – number of times N to query the Llama instance.

Return type:

NDArray[Any]

Returns:

an array of response strings.

query(max_tokens=512, rng=None)

Query the LLM once.

Return type:

str

Returns:

the response string.

generate_data(n_samples)

Generate data by querying the LLM n_samples times.

Parameters:

n_samples (int) – the number of times to query the Llama instance.

Return type:

LLMOutputData

Returns:

A data container with all responses and the prompt.

query_log_probs(rng=None)

Query the model and return the response as tokens with probabilities.

Evaluates the formatted prompt in a single forward pass, then samples tokens autoregressively using the Gumbel-max trick until EOS or max_tokens steps are reached. Each sampled token is decoded to a string via _tokens_as_strings() and its probability is computed as exp(log_softmax(logits)[token_id]).

Return type:

LLMTokenData

Returns:

A data container holding the prompt, alongside the tokens and their probabilities.

query_log_probs_next_token_ids(context_tokens, n_tokens)

Return top-k next-token candidates in token-ID space.

Return type:

CandidateTokens

query_log_probs_next_token(context_tokens, n_tokens)

Return the top-k most likely next tokens and their log-probabilities.

Resets the model state, evaluates context_tokens in a single forward pass, applies log-softmax to the last-position logits, selects candidate IDs via CandidateGeneratorFactory, and converts only those selected IDs to token strings.

EOS detection is the caller’s responsibility: the EOS token will appear naturally in the returned distribution when the model prefers it, and generate_with_sampler() stops the loop after the EOS token ID is sampled.

Parameters:

• context_tokens (list[int]) – The current context as a list of integer token IDs.

• n_tokens (int) – Number of top candidate tokens to return.

Return type:

LLMNextTokenData

Returns:

LLMNextTokenData containing the top-K token → log-prob mapping.

query_branch(context_tokens, max_tokens, rng=None)

Generate a branch of up to max_tokens and return its total log-probability.

Evaluates context_tokens in a single forward pass, then immediately snapshots the resulting KV cache and logit state via save_live_state(). That snapshot is restored via load_live_state() before the generation loop begins, which guarantees a clean branch start regardless of any side-effects from the save itself, and lays the groundwork for future multi-branch calls where the context prefix need only be evaluated once.

At each step scores[n_tokens - 1] is the logit row for the most recently decoded position (the [n_past : n_past + n_tokens] slice written by eval() with logits_all=True). A token is sampled via the Gumbel-max trick — equivalent to ancestral sampling from the full-vocabulary categorical distribution — and its log-probability is accumulated. Generation stops at EOS or after max_tokens steps.

Parameters:

• context_tokens (list[int]) – The current context as a list of integer token IDs.

• max_tokens (int) – Maximum number of tokens to generate in the branch.

Return type:

float

Returns:

Sum of per-token log-probabilities for all generated tokens, or 0.0 if EOS is sampled on the first step.

query_branch_from_live(max_tokens, rng=None)

Generate a branch from the currently loaded live state.

Unlike query_branch(), this method does not call reset() and does not evaluate any prefix/context tokens. It assumes the live model state is already positioned at the branch root.

Parameters:

max_tokens (int) – Maximum number of tokens to generate in the branch.

Return type:

float

Returns:

Sum of per-token log-probabilities for all generated tokens, or 0.0 if EOS is sampled on the first step.

query_branches_from_live_batch(max_tokens, n_branches, rng=None)

Generate n_branches independent branches in parallel from live state.

Unlike query_branch_from_live() (which decodes one branch token at a time on a single sequence), this method clones the currently loaded live (sequence-0) KV cache into n_branches auxiliary sequences and decodes one token per active sequence in a single batched llama_decode call per depth step. All branches share the same root prefix; they diverge only via independent Gumbel-max samples.

Branch proposals are drawn from the base model and are independent of any Markov-chain state, so generating them all up front in a batch is statistically identical to generating them one at a time — the caller can feed the returned log-probabilities straight into a sequential Metropolis-Hastings accept/reject loop without changing the target distribution.

Branches that sample EOS stop accumulating (their partial sum is retained, matching query_branch_from_live()’s 0.0 for immediate EOS). The live sequence-0 state and its n_tokens count are left unchanged; all auxiliary branch sequences are torn down on exit via kv_cache_seq_keep(0).

Parameters:

• max_tokens (int) – Maximum number of tokens to generate in each branch.

• n_branches (int) – Number of independent branches to generate in parallel.

• rng (RandomManager | Generator | int | None) – Optional random source override.

Return type:

NDArray[float64]

Returns:

(n_branches,) float64 array of per-branch total log-probabilities.

save_live_state()

Return LlamaState object.

Return type:

LlamaState

load_live_state(state)

Restore LlamaState object.

Return type:

None

change_context(ctx)

Update the provided context.

Making this change requires resetting the Llama instance state.

Parameters:

ctx (str) – New context string provided to the model.

Return type:

None

generate_with_sampler(top_k, top_p, adjust_fn, max_tokens, *, alpha=1.0, sampling_method=None, branch_sampler=None, rng=None)

Generate a response token-by-token with adjusted next-token probabilities.

At each step the candidate generator scales the raw logits by alpha (inverse temperature) and applies top-k/top-p truncation; the resulting candidates are passed to adjust_fn, which may modify the log-probability distribution. A single token is then sampled from the adjusted distribution and appended to the context before the next step.

Parameters:

• top_k (int) – Number of top candidate tokens to retrieve at each step.

• top_p (float) – Threshold total probability of retrieved tokens.

• adjust_fn (AdjustFn) – Callable that receives a SamplerContext and returns adjusted candidate token IDs/log-probabilities as CandidateTokens. Values do not need to be normalized.

• max_tokens (int) – Maximum number of tokens to generate.

• alpha (float) – Inverse temperature (temperature = 1 / alpha) applied by the candidate generator to the raw logits.

Return type:

LLMOutputDataFull

Returns:

LLMOutputDataFull containing the model’s response, candidate tokens at each step, and logprobs.

sample_token(top_k, top_p, alpha, adjust_fn, *, use_live_state=True, context_tokens=None, prev_probs=None, commit_token=True, rng=None)

Sample exactly one token from an adjusted next-token distribution.

This method is optimized for iterative decoding: when use_live_state is True and the model already has decoded tokens (self._llm_backend.n_tokens > 0), it does not rebuild prompt/KV state and samples directly from the current live logits.

Fallback behavior:

• if use_live_state=False: always rebuild state from context_tokens (if provided) or the formatted prompt.

• if use_live_state=True but no live state exists: rebuild from context_tokens or the formatted prompt.

Parameters:

• top_k (int) – Number of top-k tokens to consider.

• top_p (float) – Nucleus threshold in (0, 1].

• alpha (float) – Inverse temperature (temperature = 1 / alpha). applied by the candidate generator to the raw logits.

• adjust_fn (AdjustFn) – Function that adjusts candidate token log-probabilities.

• use_live_state (bool) – Prefer sampling from current live model state.

• context_tokens (list[int] | None) – Optional explicit context token IDs to evaluate when rebuilding state.

• prev_probs (list[float] | None) – Optional previously sampled token probabilities, passed through to SamplerContext.

• commit_token (bool) – Whether to append the sampled token to the live model state via decode(token_ids) when non-terminal.

Return type:

dict[str, Any]

Returns:

A dictionary containing candidate distributions before/after adjustment and details of the sampled token.

test_dataset_adjusted(dataset, top_k, top_p, adjust_fn, max_tokens)

Generate answers to a series of questions in a provided dataset.

Return type:

list[str]

---

Probability adjustment

Adjustment functions and MCMC samplers used by the adjusted-generation pipeline, including the BranchSampler abstract base class and its MetropolisSampler implementation.

Implement several sampling functions for generate_with_sampler and sample_token.

class problm_solver.samplers.SamplerContext(token_id_probs, prev_probs, context_tokens, query_next_id, query_branch, query_branch_from_live=None, base_live_state=None, query_next_ids_from_live=None, save_live_state=None, load_live_state=None, eval_tokens=None, query_branches_from_live_batch=None)

Bases: object

All information available to an AdjustFn at each generation step.

Injected by generate_with_sampler() so that adjustment functions can access model-querying capabilities without a direct dependency on Model. All mutable fields are defensive copies.

Parameters:

• token_id_probs (CandidateTokens) – Current top-k candidate token IDs and log-probabilities.

• prev_probs (list[float]) – Normalized probabilities of all previously selected tokens in this generation. Empty on the first step.

• context_tokens (list[int]) – The current token ID sequence (prompt + generated tokens so far).

• query_next_id (Callable[[list[int]], CandidateTokens]) – Queries the model for top-k next-token candidates in token-ID space given a context token ID list.

• query_branch (Callable[[list[int], int], float]) – Generates a complete branch of up to depth tokens from the given context in a single model call and returns the sum of per-token log-probabilities. Returns 0.0 on immediate EOS.

• query_branch_from_live (Callable[[int], float] | None) – Generates a complete branch of up to depth tokens from the model’s currently loaded live state and returns the sum of per-token log-probabilities. Optional.

• query_branches_from_live_batch (Callable[[int, int], NDArray[float64]] | None) – Generates n_branches independent branches of up to depth tokens in a single batched pass from the live state, returning an (n_branches,) array of per-branch total log-probabilities. Optional; when present, SamplePowerDist uses it to vectorise proposal generation.

token_id_probs: CandidateTokens

prev_probs: list[float]

context_tokens: list[int]

query_next_id: Callable[[list[int]], CandidateTokens]

query_branch: Callable[[list[int], int], float]

query_branch_from_live: Callable[[int], float] | None = None

base_live_state: Any | None = None

query_next_ids_from_live: Callable[[int], list[tuple[int, float]]] | None = None

save_live_state: Callable[[], Any] | None = None

load_live_state: Callable[[Any], None] | None = None

eval_tokens: Callable[[list[int]], None] | None = None

query_branches_from_live_batch: Callable[[int, int], NDArray[float64]] | None = None

type problm_solver.samplers.AdjustFn = Callable[[SamplerContext], CandidateTokens]

problm_solver.samplers.candidate_tokens_to_id_logprobs(candidates)

Convert CandidateTokens to an insertion-ordered token-id logprob map.

Return type:

dict[int, float]

problm_solver.samplers.id_logprobs_to_candidate_tokens(id_logprobs)

Convert token-id logprob map to CandidateTokens preserving insertion order.

Return type:

CandidateTokens

problm_solver.samplers.adjust_identity(context)

Return token log-probabilities unchanged in token-ID space.

Return type:

CandidateTokens

class problm_solver.samplers.SampleLowTempNucleus(alpha)

Bases: object

Labelled pass-through adjustment for low-temperature (power-scaled) sampling.

Temperature scaling itself now happens inside the candidate generator: CandidateGeneratorFactory scales the raw logits by alpha (inverse temperature, temperature = 1 / alpha) BEFORE top-k/top-p truncation, matching HuggingFace transformers’ warper order (TemperatureLogitsWarper, then TopKLogitsWarper, then TopPLogitsWarper). The candidates received by this adjustment have therefore already been temperature-scaled and truncated, and this class returns them unchanged.

It is retained as a named class so that generate_with_sampler records a meaningful sampling_method label (‘SampleLowTempNucleus’) for low-temperature runs, and so the low-temperature API keeps a stable surface. It must not apply any scaling itself: doing so would apply the temperature a second time.

The selection history (prev_probs) is deliberately not folded into the output. An earlier version added alpha * sum(log(prev_probs)) to every candidate’s log-probability, modelling a joint-sequence probability. Because that term is identical across candidates at a given step, it is shift-invariant under the softmax used for sampling and therefore has no effect on which token is selected or on the stored per-token probabilities. Its only observable effect was to inject a steadily growing negative offset into the top-k log-probability map stored in LLMOutputDataFull.response_topk. Omitting it keeps the stored log-probabilities on a stable, interpretable scale.

Parameters:

alpha (float) – The inverse temperature used by the run. Recorded for reference only; the scaling itself is performed by the candidate generator.

Example usage:

adjust_fn = SampleLowTempNucleus(alpha=2)
result = adjust_fn(context)

Initialize with the run’s inverse temperature.

Parameters:

alpha (float) – Inverse temperature of the run. Recorded for reference; applied by the candidate generator, not here.

class problm_solver.samplers.BranchSampler

Bases: ABC

Abstract base class for branch-level sampling strategies.

A BranchSampler typically runs over complete branch proposals. Some subclasses may additionally provide token-by-token beam expansion via future_logprob_from_context().

supports_token_beam = False

property max_proposals: int

Max branch proposals a batched path should generate upfront.

Returns 0 by default to signal that this sampler does not support batched (vectorised) proposal generation. Samplers that do (e.g. MetropolisSampler) override this so SamplePowerDist knows how many proposals to request in a single batched call.

reset()

Reset internal state at the start of each candidate-token chain.

No-op for stateless samplers. Stateful samplers (e.g. MetropolisSampler) should override this.

Return type:

None

abstractmethod step(proposed_log_prob, alpha=1.0, forward_log_q=0.0, reverse_log_q=0.0)

Process one proposed branch and return the accepted chain state.

Return type:

float

abstractmethod should_continue(branch_log_probs)

Return True if more branch proposals should be sampled.

Return type:

bool

abstractmethod future_logprob(alpha, branch_log_probs)

Calculate weighting to token probability from sampled branches.

Return type:

float64

future_logprob_from_context(alpha, base_live_state, branch_token_ids, lookahead_depth, query_next_ids_from_live, save_live_state, load_live_state, eval_tokens)

Optional token-by-token beam expansion hook.

Subclasses that implement token-level beam search should override this method and set supports_token_beam = True.

Return type:

float64

class problm_solver.samplers.MetropolisSampler(equil_branches=5, max_branches=30, tolerance=0.1, rng=None)

Bases: BranchSampler

Metropolis-Hastings sampler over complete branch proposals.

Given a sequence of proposed branch log-probabilities, maintains an MCMC chain and accepts a proposal with probability

min(1, exp(log p(x') - log p(x) + log q(x|x') - log q(x'|x))).

Convergence across accepted branch samples is assessed via the standard error of the mean (SEM): SEM = std(branch_log_probs) / sqrt(n). Sampling continues until SEM < tolerance, after at least equil_branches samples, and always stops at max_branches.

Parameters:

• equil_branches (int) – Number of accepted samples treated as burn-in (equilibration); discarded before checking SEM.

• max_branches (int) – Hard upper limit on accepted samples.

• tolerance (float) – SEM threshold below which sampling is considered converged.

Initialize with convergence parameters.

property max_proposals: int

Number of proposals to generate in one batched pass (= max_branches).

reset()

Clear chain state before starting a new candidate-token chain.

Return type:

None

step(proposed_log_prob, alpha=1.0, forward_log_q=0.0, reverse_log_q=0.0)

Apply one Metropolis-Hastings accept/reject step targeting p^α.

The log acceptance ratio is (α-1) * (log p(x') - log p(x)) + log q(x|x') - log q(x'|x). When the proposal q is the base model p the proposal terms cancel (forward_log_q = proposed_log_prob, reverse_log_q = current_log_prob), reducing to (α-1) * (proposed - current).

Parameters:

• proposed_log_prob (float) – Proposed branch log-probability under p.

• alpha (float) – Power-distribution exponent.

• forward_log_q (float) – log q(x'|x) for the proposal.

• reverse_log_q (float) – log q(x|x') for the reverse proposal.

Return type:

float

Returns:

Accepted chain state’s log-probability.

should_continue(branch_log_probs)

Return True if more proposals should be sampled.

Uses SEM-based convergence after equil_branches and before max_branches.

Param:

branch logarithmic probabilities to date.

Return type:

bool

future_logprob(alpha, branch_log_probs)

Monte Carlo mean weight.

Return type:

float64

class problm_solver.samplers.BeamSampler(beam_width=3, branch_top_k=5)

Bases: BranchSampler

Token-by-token beam expansion for future-branch scoring.

This sampler performs deterministic beam search over lookahead tokens. For each candidate token, it repeatedly expands active beams using query_next_ids_from_live and keeps only the top beam_width cumulative log-probability branches at every depth.

Parameters:

• beam_width (int) – Number of active beams retained per depth.

• branch_top_k (int) – Number of next-token candidates considered for each active beam during expansion.

Initialize beam-search width and per-beam expansion width.

supports_token_beam = True

reset()

No-op: beam expansion is stateless across candidates.

Return type:

None

step(proposed_log_prob, alpha=1.0, forward_log_q=0.0, reverse_log_q=0.0)

Compatibility no-op; token-beam mode does not use MH transitions.

Return type:

float

should_continue(branch_log_probs)

Compatibility no-op; token-beam mode controls depth directly.

Return type:

bool

future_logprob(alpha, branch_log_probs)

Return log-mean-exp over supplied branch scores.

This method is retained for compatibility, but token-beam mode normally uses future_logprob_from_context().

Return type:

float64

future_logprob_from_context(alpha, base_live_state, branch_token_ids, lookahead_depth, query_next_ids_from_live, save_live_state, load_live_state, eval_tokens)

Run token-level beam expansion with KV-cache state reuse.

Return type:

float64

class problm_solver.samplers.SamplePowerDist(alpha, lookahead_depth, branch_sampler)

Bases: object

Adjust token log-probabilities using future-branch power-distribution sampling.

For each candidate next token, repeatedly proposes complete future branches of length lookahead_depth and updates a Markov chain using the injected BranchSampler (e.g. Metropolis-Hastings), continuing until should_continue() signals convergence. Each branch is evaluated in a single model call via query_branch, rather than token-by-token. The accepted branch log-probabilities are kept as a numpy array and combined with the current token’s log-probability via log-sum-exp to produce the adjusted distribution.

Branches that reach EOS before lookahead_depth are terminated early with no penalty — their partial log-probability is used as-is.

Parameters:

• alpha (float) – Scaling exponent applied to the current token log-probability.

• lookahead_depth (int) – Maximum number of steps to sample in each branch.

• branch_sampler (BranchSampler) – Strategy used to sample tokens within each branch and to determine when enough branches have been collected. Must be a BranchSampler instance; its reset() method is called at the start of every branch.

Example usage:

sampler = SamplePowerDist(
    alpha=2.0,
    lookahead_depth=3,
    branch_sampler=MetropolisSampler(),
)
result = sampler(context)

Initialize with lookahead parameters and a branch sampler.

Parameters:

• alpha (float) – Scaling exponent for the current token log-probability.

• lookahead_depth (int) – Maximum depth of each branch.

• branch_sampler (BranchSampler) – The BranchSampler to use within branches and for convergence decisions.

---

Analysis — token probabilities

Utilities for querying per-token probabilities and sampling from log-prob dictionaries returned by the model.

Utilities for sampling tokens and getting token probabilities.

problm_solver.analysis.probabilities.prob_of_token(token, log_probs)

Return the normalized probability of a specific token from a log-prob dict.

Applies the same shift-exp-normalise procedure as sample_from_logprobs(), then returns the scalar probability for the named token rather than sampling.

Parameters:

• token (str) – The token string to look up. Must be a key in log_probs.

• log_probs (dict[str, float]) – Mapping of token string to log-probability.

Return type:

float

Returns:

The normalized probability of token in the distribution, in the range (0, 1].

Raises:

KeyError – If token is not present in log_probs.

problm_solver.analysis.probabilities.sample_from_logprobs(log_probs, rng=None)

Sample a token from a log-probability distribution.

Converts log-probabilities to probabilities via exp(), renormalizes, and returns a single sampled token string.

Parameters:

log_probs (dict[str, float]) – Mapping of token string to log-probability. Values do not need to correspond to a normalised distribution — renormalisation is applied before sampling.

Return type:

str

Returns:

A single sampled token string drawn from the distribution.

---

Random number generation

Utilities for managing deterministic random streams across the library.

Module interface for random.

type problm_solver.random.RNGLike = RandomManager | Generator | int | None

class problm_solver.random.RandomManager(seed=314159)

Bases: object

Manage deterministic pseudo-random streams from a single user seed.

A RandomManager provides named RNG streams so the caller can set one seed and still keep randomness in different subsystems independent (e.g. token sampling vs. Metropolis acceptance draws).

Initialize random manager.

Parameters:

seed (int) – Root integer seed for all derived streams.

property seed: int

Return root seed value.

get_rng(stream='global')

Return a cached RNG for a named stream.

Repeated calls with the same stream return the same generator instance so state advances naturally across calls.

Parameters:

stream (str) – Name of the deterministic stream.

Return type:

Generator

Returns:

numpy.random.Generator bound to stream.

spawn_rng(stream='global')

Return a fresh RNG in a named stream family.

Unlike get_rng(), each call returns a new generator that is deterministic for (seed, stream, spawn_index).

Parameters:

stream (str) – Name of stream family.

Return type:

Generator

Returns:

Fresh deterministic numpy.random.Generator instance.

problm_solver.random.resolve_rng(rng, *, stream, default_manager=None, fresh=False)

Take an RNGLike and returns a generator.

Parameters:

• rng (RandomManager | Generator | int | None) – RNGLike object, to be sorted.

• stream (str) – Named stream for RandomManager objects.

• default_manager (RandomManager | None) – Fallback manager, used when calling resolve_rng() in classes that manage their own RNG state.

• fresh (bool) – Whether or not to reinitialize a generator (should be False in most cases).

Return type:

Generator

Returns:

A numpy generator exhibiting specified behavior.

---

CLI

Entry point for the interactive command-line interface (problm-solver). Documented here for completeness; most users will interact with the library programmatically rather than through the CLI module directly.

Command line interface main loop contained here.

exception problm_solver.cli.UnexpectedFunctionError

Bases: Exception

Raised when the developer hasn’t extended the interface properly.

Initialize error message.

problm_solver.cli.ensure_models_dir()

Create the models directory if it doesn’t exist and return its path.

Return type:

Path

problm_solver.cli.ensure_responses_dir()

Create the responses directory if it doesn’t exist and return its path.

Return type:

Path

problm_solver.cli.ensure_probs_dir()

Create the probabilities directory if it doesn’t exist and return its path.

Return type:

Path

problm_solver.cli.list_models()

Return a sorted list of GGUF files in the models directory.

Return type:

list[Path]

problm_solver.cli.get_responses_path(model_path)

Return a timestamped .jsonl path inside RESPONSES_DIR for a response dataset.

Return type:

Path

problm_solver.cli.get_adjusted_path(model_path)

Return a timestamped .json path inside RESPONSES_DIR for an adjusted response.

Return type:

Path

problm_solver.cli.get_probs_path(model_path)

Return a timestamped .json path inside PROBS_DIR for a probability dataset.

Return type:

Path

problm_solver.cli.get_math500_results_path(model_path)

Return a timestamped .jsonl path inside RESPONSES_DIR for a MATH500 results dataset.

Return type:

Path

problm_solver.cli.ui_select_model()

Display available GGUF models and prompt the user to pick one.

Return type:

Path

problm_solver.cli.ui_gen_data(model, model_path)

Handle user interface for generating LLM data.

Return type:

None

problm_solver.cli.ui_save_data(fname, data)

Handle user interface for saving LLM response data.

Return type:

None

problm_solver.cli.ui_save_token_data(fname, data)

Handle user interface for saving LLM token probability data.

Return type:

None

problm_solver.cli.ui_get_probs(model, model_path)

Handle user interface for querying once with logprobs enabled.

Return type:

None

problm_solver.cli.ui_select_function()

Prompt the user to select a program function and return their integer choice.

Return type:

int

problm_solver.cli.ui_generate_low_temp(model, model_path)

Handle user interface for getting model response using low temp sampling.

Return type:

None

problm_solver.cli.ui_generate_power_mcmc(model, model_path)

Handle user interface for getting model response using power sampling with MCMC.

Return type:

None

problm_solver.cli.ui_math500(model, model_path)

Answer MATH500 problems and write response to file.

Return type:

None

problm_solver.cli.main()

Perform main CLI.

Return type:

None