problm-solver 1.5.1__py3-none-any.whl

problm_solver/__init__.py

@@ -0,0 +1,9 @@
+"""probLM-solver."""
+
+from problm_solver.random import RandomManager
+
+PSRandom = RandomManager
+
+__all__ = [
+    'PSRandom'
+]

problm_solver/adjust_probs.py

@@ -0,0 +1,583 @@
+"""Implement several adjustment functions for generate_adjusted."""
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+from tqdm import tqdm
+
+from problm_solver.candidates import CandidateTokens
+from problm_solver.random import RNGLike, resolve_rng
+
+# -- Module-wide setup -- #
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GenerationContext:
+    """All information available to an :data:`AdjustFn` at each generation step.
+
+    Injected by ``generate_adjusted()`` so that adjustment functions can
+    access model-querying capabilities without a direct dependency on
+    ``ModelInstance``. All mutable fields are defensive copies.
+
+    :param token_id_probs: Current top-k candidate token IDs and log-probabilities.
+    :param prev_probs: Normalised probabilities of all previously selected
+        tokens in this generation. Empty on the first step.
+    :param context_tokens: The current token ID sequence (prompt + generated
+        tokens so far).
+    :param query_next_id: Queries the model for top-k next-token candidates in
+        token-ID space given a context token ID list.
+    :param query_branch: 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.
+    :param query_branch_from_live: 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.
+    """
+
+    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
+
+
+# Callable that receives a GenerationContext and returns adjusted
+# token-ID candidates with log-probabilities.
+type AdjustFn = Callable[[GenerationContext], CandidateTokens]
+
+
+def candidate_tokens_to_id_logprobs(candidates: CandidateTokens) -> dict[int, float]:
+    """Convert CandidateTokens to an insertion-ordered token-id logprob map."""
+    return {
+        int(tid): float(lp)
+        for tid, lp in zip(
+            candidates.candidate_ids,
+            candidates.candidate_logprobs,
+            strict=True,
+        )
+    }
+
+
+def id_logprobs_to_candidate_tokens(id_logprobs: dict[int, float]) -> CandidateTokens:
+    """Convert token-id logprob map to CandidateTokens preserving insertion order."""
+    items = list(id_logprobs.items())
+    if not items:
+        return CandidateTokens(
+            candidate_ids=np.empty(0, dtype=np.int32),
+            candidate_logprobs=np.empty(0, dtype=np.float64),
+        )
+
+    ids = np.array([tid for tid, _ in items], dtype=np.int32)
+    lps = np.array([lp for _, lp in items], dtype=np.float64)
+    return CandidateTokens(candidate_ids=ids, candidate_logprobs=lps)
+
+
+def adjust_identity(context: GenerationContext) -> CandidateTokens:
+    """Return token log-probabilities unchanged in token-ID space."""
+    return context.token_id_probs
+
+
+class SampleLowTemp:
+    """Adjust token log-probabilities by power-scaling with selection history.
+
+    At each generation step, the current token probabilities are raised to
+    ``alpha`` and multiplied by the product of all previously selected token
+    probabilities each also raised to ``alpha``. The result is returned as
+    log-probabilities for downstream renormalisation and sampling.
+
+    :param alpha: Scaling exponent. Values greater than 1 sharpen the
+        distribution (favouring already-likely tokens); values between 0
+        and 1 flatten it.
+
+    Example usage::
+
+        adjust_fn = SampleLowTemp(alpha=2)
+        result = adjust_fn(context)
+    """
+
+    def __init__(self, alpha: float) -> None:
+        """Initialise with scaling exponent.
+
+        :param alpha: Exponent applied to current and previous token
+            probabilities when computing the adjustment.
+        """
+        self.alpha = alpha
+        self._prev_len = 0
+        self._prev_logprob_sum = 0.0
+
+    def reset(self) -> None:
+        """Reset rolling history state for a new generation."""
+        self._prev_len = 0
+        self._prev_logprob_sum = 0.0
+
+    def _history_log_shift(self, prev_probs: list[float]) -> float:
+        """Return rolling ``alpha * sum(log(prev_probs))`` for history scaling."""
+        cur_len = len(prev_probs)
+        if cur_len == 0:
+            self.reset()
+            return 0.0
+
+        if cur_len < self._prev_len:
+            self._prev_logprob_sum = float(np.sum(np.log(np.array(prev_probs, dtype=float))))
+        elif cur_len == self._prev_len + 1:
+            self._prev_logprob_sum += float(np.log(prev_probs[-1]))
+        elif cur_len == self._prev_len:
+            self._prev_logprob_sum = float(np.sum(np.log(np.array(prev_probs, dtype=float))))
+        else:
+            self._prev_logprob_sum = float(np.sum(np.log(np.array(prev_probs, dtype=float))))
+
+        self._prev_len = cur_len
+        return self.alpha * self._prev_logprob_sum
+
+    def __call__(self, context: GenerationContext) -> CandidateTokens:
+        """Apply power-scaling adjustment to the current token-ID distribution."""
+        candidate_ids = context.token_id_probs.candidate_ids
+        lp = context.token_id_probs.candidate_logprobs.astype(np.float64, copy=True)
+        lp -= lp.max()
+
+        prev_probs = context.prev_probs if context.prev_probs is not None else []
+        history_log_shift = self._history_log_shift(prev_probs)
+        new_logprobs: npt.NDArray[np.float64] = self.alpha * lp + history_log_shift
+        return CandidateTokens(
+            candidate_ids=candidate_ids.astype(np.int32, copy=False),
+            candidate_logprobs=new_logprobs,
+        )
+
+
+class BranchSampler(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
+    :meth:`future_logprob_from_context`.
+    """
+
+    supports_token_beam = False
+
+    def reset(self) -> None: # noqa: B027
+        """Reset internal state at the start of each candidate-token chain.
+
+        No-op for stateless samplers. Stateful samplers (e.g.
+        :class:`MetropolisSampler`) should override this.
+        """
+
+    @abstractmethod
+    def step(
+        self,
+        proposed_log_prob: float,
+        alpha: float = 1.0,
+        forward_log_q: float = 0.0,
+        reverse_log_q: float = 0.0,
+    ) -> float:
+        """Process one proposed branch and return the accepted chain state."""
+
+    @abstractmethod
+    def should_continue(self, branch_log_probs: npt.NDArray[np.float64]) -> bool:
+        """Return ``True`` if more branch proposals should be sampled."""
+
+    @abstractmethod
+    def future_logprob(self, alpha: float, branch_log_probs: npt.NDArray[np.float64]) -> np.float64:
+        """Calculate weighting to token probability from sampled branches."""
+
+    def future_logprob_from_context(
+        self,
+        alpha: float,
+        base_live_state: Any,
+        branch_token_ids: list[int],
+        lookahead_depth: int,
+        query_next_ids_from_live: Callable[[int], list[tuple[int, float]]],
+        save_live_state: Callable[[], Any],
+        load_live_state: Callable[[Any], None],
+        eval_tokens: Callable[[list[int]], None],
+    ) -> np.float64:
+        """Optional token-by-token beam expansion hook.
+
+        Subclasses that implement token-level beam search should override this
+        method and set ``supports_token_beam = True``.
+        """
+        raise NotImplementedError
+
+
+class MetropolisSampler(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``.
+
+    :param equil_branches: Number of accepted samples treated as burn-in
+        (equilibration); discarded before checking SEM.
+    :param max_branches: Hard upper limit on accepted samples.
+    :param tolerance: SEM threshold below which sampling is considered
+        converged.
+    """
+
+    def __init__(
+        self,
+        equil_branches: int = 5,
+        max_branches: int = 30,
+        tolerance: float = 1e-1,
+        rng: RNGLike = None
+    ) -> None:
+        """Initialise with convergence parameters."""
+        self._current_log_prob: float | None = None
+        self._equil_branches = equil_branches
+        self._max_branches = max_branches
+        self._tolerance = tolerance
+        self._rng = resolve_rng(rng, stream='adjust.metropolis')
+
+    def reset(self) -> None:
+        """Clear chain state before starting a new candidate-token chain."""
+        self._current_log_prob = None
+
+    def step(
+        self,
+        proposed_log_prob: float,
+        alpha: float = 1.0,
+        forward_log_q: float = 0.0,
+        reverse_log_q: float = 0.0,
+    ) -> float:
+        """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)``.
+
+        :param proposed_log_prob: Proposed branch log-probability under ``p``.
+        :param alpha: Power-distribution exponent.
+        :param forward_log_q: ``log q(x'|x)`` for the proposal.
+        :param reverse_log_q: ``log q(x|x')`` for the reverse proposal.
+        :returns: Accepted chain state's log-probability.
+        """
+        if self._current_log_prob is None:
+            self._current_log_prob = proposed_log_prob
+            return self._current_log_prob
+
+        log_accept_ratio = (
+            (alpha - 1) * (proposed_log_prob - self._current_log_prob)
+            + reverse_log_q
+            - forward_log_q
+        )
+        if np.log(self._rng.random()) < min(0.0, log_accept_ratio):
+            self._current_log_prob = proposed_log_prob
+
+        return self._current_log_prob
+
+    def should_continue(self, branch_log_probs: npt.NDArray[np.float64]) -> bool:
+        """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 len(branch_log_probs) < self._max_branches
+
+    def future_logprob(self, alpha: float, branch_log_probs: npt.NDArray[np.float64]) -> np.float64:
+        """Monte Carlo mean weight."""
+        post_eq = branch_log_probs[self._equil_branches:]
+        scaled = alpha * post_eq
+        max_lp = np.float64(scaled.max())
+        return np.log(np.mean(np.exp(scaled - max_lp))) + max_lp
+
+
+class BeamSampler(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.
+
+    :param beam_width: Number of active beams retained per depth.
+    :param branch_top_k: Number of next-token candidates considered for each
+        active beam during expansion.
+    """
+
+    supports_token_beam = True
+
+    def __init__(self, beam_width: int = 3, branch_top_k: int = 5) -> None:
+        """Initialise beam-search width and per-beam expansion width."""
+        if beam_width < 1:
+            msg = f'beam_width must be >= 1, got {beam_width}'
+            raise ValueError(msg)
+        if branch_top_k < 1:
+            msg = f'branch_top_k must be >= 1, got {branch_top_k}'
+            raise ValueError(msg)
+
+        self.beam_width = beam_width
+        self.branch_top_k = branch_top_k
+
+    def reset(self) -> None:
+        """No-op: beam expansion is stateless across candidates."""
+
+    def step(
+        self,
+        proposed_log_prob: float,
+        alpha: float = 1.0, #noqa:ARG002
+        forward_log_q: float = 0.0, #noqa:ARG002
+        reverse_log_q: float = 0.0, #noqa:ARG002
+    ) -> float:
+        """Compatibility no-op; token-beam mode does not use MH transitions."""
+        return proposed_log_prob
+
+    def should_continue(self, branch_log_probs: npt.NDArray[np.float64]) -> bool:
+        """Compatibility no-op; token-beam mode controls depth directly."""
+        return False
+
+    def future_logprob(self, alpha: float, branch_log_probs: npt.NDArray[np.float64]) -> np.float64:
+        """Return log-mean-exp over supplied branch scores.
+
+        This method is retained for compatibility, but token-beam mode
+        normally uses :meth:`future_logprob_from_context`.
+        """
+        if len(branch_log_probs) == 0:
+            msg = 'branch_log_probs cannot be empty'
+            raise ValueError(msg)
+
+        scaled = alpha * branch_log_probs
+        max_lp = np.float64(scaled.max())
+        return np.log(np.mean(np.exp(scaled - max_lp))) + max_lp
+
+    def future_logprob_from_context(
+        self,
+        alpha: float,
+        base_live_state: Any,
+        branch_token_ids: list[int],
+        lookahead_depth: int,
+        query_next_ids_from_live: Callable[[int], list[tuple[int, float]]],
+        save_live_state: Callable[[], Any],
+        load_live_state: Callable[[Any], None],
+        eval_tokens: Callable[[list[int]], None],
+    ) -> np.float64:
+        """Run token-level beam expansion with KV-cache state reuse."""
+        load_live_state(base_live_state)
+        if branch_token_ids:
+            eval_tokens(branch_token_ids)
+        root_state = save_live_state()
+
+        beams: list[tuple[Any, float]] = [(root_state, 0.0)]
+
+        for _ in range(lookahead_depth):
+            expanded: list[tuple[Any, float]] = []
+
+            for beam_state, cum_lp in beams:
+                load_live_state(beam_state)
+                top_next = query_next_ids_from_live(self.branch_top_k)
+
+                for token_id, token_lp in top_next:
+                    load_live_state(beam_state)
+                    eval_tokens([token_id])
+                    child_state = save_live_state()
+                    expanded.append((child_state, cum_lp + float(token_lp)))
+
+            if not expanded:
+                break
+
+            expanded.sort(key=lambda item: item[1], reverse=True)
+            beams = expanded[: self.beam_width]
+
+        if not beams:
+            return np.float64(-np.inf)
+
+        beam_log_probs = np.array([lp for _, lp in beams], dtype=np.float64)
+        scaled = alpha * beam_log_probs
+        max_lp = np.float64(scaled.max())
+        return np.log(np.mean(np.exp(scaled - max_lp))) + max_lp
+
+
+class SamplePowerDist:
+    """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 :class:`BranchSampler` (e.g. Metropolis-Hastings), continuing
+    until :meth:`~BranchSampler.should_continue` signals convergence. Each
+    branch is evaluated in a single model call via
+    :attr:`~GenerationContext.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.
+
+    :param alpha: Scaling exponent applied to the current token log-probability.
+    :param lookahead_depth: Maximum number of steps to sample in each branch.
+    :param branch_sampler: Strategy used to sample tokens within each branch
+        and to determine when enough branches have been collected.
+        Must be a :class:`BranchSampler` instance; its
+        :meth:`~BranchSampler.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)
+    """
+
+    def __init__(
+        self,
+        alpha: float,
+        lookahead_depth: int,
+        branch_sampler: BranchSampler,
+    ) -> None:
+        """Initialise with lookahead parameters and a branch sampler.
+
+        :param alpha: Scaling exponent for the current token log-probability.
+        :param lookahead_depth: Maximum depth of each branch.
+        :param branch_sampler: The :class:`BranchSampler` to use within
+            branches and for convergence decisions.
+        """
+        self.alpha = alpha
+        self.lookahead_depth = lookahead_depth
+        self.branch_sampler = branch_sampler
+
+    def __call__(self, context: GenerationContext) -> CandidateTokens:
+        """Apply power-distribution adjustment using lookahead branch sampling.
+
+        :param context: The current generation context in token-ID space.
+        :returns: Adjusted candidate token IDs with log-probabilities.
+        """
+        result: dict[int, float] = {}
+
+        if self.branch_sampler.supports_token_beam:
+            if (
+                context.base_live_state is None
+                or context.query_next_ids_from_live is None
+                or context.save_live_state is None
+                or context.load_live_state is None
+                or context.eval_tokens is None
+            ):
+                msg = 'Token-beam sampler requires live-state callables in GenerationContext'
+                raise ValueError(msg)
+
+            base_live_state = context.base_live_state
+            query_next_ids_from_live = context.query_next_ids_from_live
+            save_live_state = context.save_live_state
+            load_live_state = context.load_live_state
+            eval_tokens = context.eval_tokens
+
+            def score_future(branch_token_ids: list[int]) -> np.float64:
+                return self.branch_sampler.future_logprob_from_context(
+                    alpha=self.alpha,
+                    base_live_state=base_live_state,
+                    branch_token_ids=branch_token_ids,
+                    lookahead_depth=self.lookahead_depth,
+                    query_next_ids_from_live=query_next_ids_from_live,
+                    save_live_state=save_live_state,
+                    load_live_state=load_live_state,
+                    eval_tokens=eval_tokens,
+                )
+        else:
+            has_live_branch = (
+                context.query_branch_from_live is not None
+                and context.base_live_state is not None
+                and context.save_live_state is not None
+                and context.load_live_state is not None
+                and context.eval_tokens is not None
+            )
+
+            if has_live_branch:
+                base_live_state = context.base_live_state
+                query_branch_from_live = context.query_branch_from_live
+                save_live_state = context.save_live_state
+                load_live_state = context.load_live_state
+                eval_tokens = context.eval_tokens
+
+                def score_future_from_candidate_id(candidate_id: int) -> np.float64:
+                    load_live_state(base_live_state)
+                    eval_tokens([candidate_id])
+                    candidate_root_state = save_live_state()
+
+                    branch_log_probs_list: list[float] = []
+                    self.branch_sampler.reset()
+
+                    while True:
+                        load_live_state(candidate_root_state)
+                        proposed_branch_log_prob = query_branch_from_live(self.lookahead_depth)
+
+                        accepted_log_prob = self.branch_sampler.step(
+                            proposed_log_prob=proposed_branch_log_prob,
+                            alpha=self.alpha,
+                        )
+                        branch_log_probs_list.append(accepted_log_prob)
+
+                        if not self.branch_sampler.should_continue(
+                            np.array(branch_log_probs_list, dtype=np.float64)
+                        ):
+                            break
+
+                    branch_log_probs = np.array(branch_log_probs_list, dtype=np.float64)
+                    return self.branch_sampler.future_logprob(self.alpha, branch_log_probs)
+            else:
+                def score_future_from_candidate_id(candidate_id: int) -> np.float64:
+                    branch_ctx = list(context.context_tokens) + [candidate_id]
+                    branch_log_probs_list: list[float] = []
+                    self.branch_sampler.reset()
+
+                    while True:
+                        proposed_branch_log_prob = context.query_branch(
+                            branch_ctx,
+                            self.lookahead_depth,
+                        )
+
+                        accepted_log_prob = self.branch_sampler.step(
+                            proposed_log_prob=proposed_branch_log_prob,
+                            alpha=self.alpha,
+                        )
+                        branch_log_probs_list.append(accepted_log_prob)
+
+                        if not self.branch_sampler.should_continue(
+                            np.array(branch_log_probs_list, dtype=np.float64)
+                        ):
+                            break
+
+                    branch_log_probs = np.array(branch_log_probs_list, dtype=np.float64)
+                    return self.branch_sampler.future_logprob(self.alpha, branch_log_probs)
+
+        candidate_ids = context.token_id_probs.candidate_ids
+        candidate_logprobs = context.token_id_probs.candidate_logprobs
+
+        candidate_bar = tqdm(
+            zip(candidate_ids, candidate_logprobs, strict=True),
+            desc='candidates',
+            total=len(candidate_ids),
+            unit='tok',
+            leave=False,
+        )
+        for token_id, log_prob in candidate_bar:
+            tid = int(token_id)
+            if self.branch_sampler.supports_token_beam:
+                future_lp = score_future([tid])
+            else:
+                future_lp = score_future_from_candidate_id(tid)
+            result[tid] = self.alpha * float(log_prob) + float(future_lp)
+
+        return id_logprobs_to_candidate_tokens(result)

problm_solver/analysis/__init__.py

@@ -0,0 +1,8 @@
+"""Statistical analysis tools for LLM output data."""
+
+from problm_solver.analysis.probabilities import prob_of_token, sample_from_logprobs
+
+__all__ = [
+    'prob_of_token',
+    'sample_from_logprobs',
+]

problm_solver/analysis/probabilities.py

@@ -0,0 +1,51 @@
+"""Utilities for sampling tokens and getting token probabilities."""
+
+
+import numpy as np
+
+from problm_solver.random import RNGLike, resolve_rng
+
+
+def prob_of_token(token: str, log_probs: dict[str, float]) -> float:
+    """Return the normalised probability of a specific token from a log-prob dict.
+
+    Applies the same shift-exp-normalise procedure as
+    :func:`sample_from_logprobs`, then returns the scalar probability for
+    the named token rather than sampling.
+
+    :param token: The token string to look up. Must be a key in ``log_probs``.
+    :param log_probs: Mapping of token string to log-probability.
+    :returns: The normalised probability of ``token`` in the distribution,
+        in the range (0, 1].
+    :raises KeyError: If ``token`` is not present in ``log_probs``.
+    """
+    tokens = list(log_probs.keys())
+    lp = np.array([log_probs[t] for t in tokens], dtype=np.float64)
+    lp -= lp.max()
+    probs = np.exp(lp)
+    probs /= probs.sum()
+    return float(probs[tokens.index(token)])
+
+
+def sample_from_logprobs(
+    log_probs: dict[str, float],
+    rng: RNGLike = None
+) -> str:
+    """Sample a token from a log-probability distribution.
+
+    Converts log-probabilities to probabilities via ``exp()``, renormalises,
+    and returns a single sampled token string.
+
+    :param log_probs: Mapping of token string to log-probability. Values do
+        not need to correspond to a normalised distribution — renormalisation
+        is applied before sampling.
+    :returns: A single sampled token string drawn from the distribution.
+    """
+    tokens = list(log_probs.keys())
+    lp = np.array([log_probs[t] for t in tokens], dtype=np.float64)
+    lp -= lp.max()  # shift for numerical stability before exp
+    probs = np.exp(lp)
+    probs /= probs.sum()
+    method_rng = resolve_rng(rng, stream='analysis.sample_from_logprobs')
+    idx: int = int(method_rng.choice(len(tokens), p=probs))
+    return tokens[idx]