Tutorials#

Overview#

This page provides a quick-start guide to scripting with probLM-solver. Please see the other tutorial pages listed below if they are relevant to you.

Tutorial Pages#

CLI#

To get started with using the CLI for already-implemented samplers and generating large numbers of outputs, to serve the app straight from the command line, run

$ docker compose --profile app up app

You will be asked to provide a .gguf model file in the directory ~/.problm-solver/models

Scripting with problm_solver#

The main class that probLM-solver uses for interacting with an LLM is problm_solver.llama_interface.Model. A simple script for setting up a model, and generating a response with low-temperature sampling, can be seen in the code block below.

 1from problm_solver.samplers import SampleLowTemp
 2from problm_solver.llama_interface import Model
 3
 4temp = 0.5
 5top_k = 8
 6top_p = 0.9
 7model = Model('path/to/model.gguf', 'Why is the sky blue?', logits_all=True)
 8
 9sampling_func = SampleLowTemp(alpha=1./temp)
10data = model.generate_with_sampler(
11    top_k=top_k,
12    top_p=top_p,
13    adjust_fn=sampling_func,
14    max_tokens=128
15)

Let’s explain in detail what this code is doing.

The key idea here is that sampling functions can be defined or freely switched between when generating data, as they aren’t specified until Model.generate_with_sampler() is called.

A couple of hyperparameters are specified in advance, namely the temperature (temp), of the model and the number of most-probable tokens considered at each step (top_k). If you have already heard of low-temperature sampling or top-k sampling, these concepts will be familiar.

Next, the model stored at a given path is instanced with some context. Importantly, the parameter logits_all=True is set. If this is left blank, i.e. False, then custom sampling functions will fail when trying to access the model’s native token probabilities in order to adjust them.

Then, a sampling function is generated from one of the built-in classes. SampleLowTemp() is a callable class, with relevant hyperparameters set when it is instanced.

Finally, data from the response are collected into a dataclass (data), as an output of generate_with_sampler(). This dictionary is an instance of problm_solver.data.LLMOutputDataFull, and a list of the values contained within is as follows:

 1from problm_solver.data import LLMOutputDataFull
 2
 3...
 4
 5data.context # Is a tokenised version of the context provided to the model.
 6data.hyperparams # A nested dataclass containing the hyperparameters used.
 7data.response_probabilities # A tuple of chosen tokens and their (post-adjustment) probability of being chosen.
 8data.response_topk # The top-k candidate tokens at each step, and their logarithmic probabilities.
 9data.sampling_method # User defined label for the method used.
10data.branch_sampler # User defined label for the branch sampler used. (More on that later.)

Sampling Token-by-token#

The previous section introduced the function generate_with_sampler() to generate outputs based on a provided context and sampling function. It may sometimes be desirable for the user to generate one token at a time, to either switch sampling functions out or analyze specific probabilities. This is where the method Model.sample_token() is used. See the following example:

 1from problm_solver.samplers import SampleLowTemp
 2from problm_solver.llama_interface import Model
 3
 4temp = 0.5
 5top_k = 30
 6top_p = 0.9
 7model = Model('path/to/model.gguf', 'Why is the sky blue?', logits_all=True)
 8
 9sampling_func = SampleLowTemp(alpha=1./temp)
10token_data = model.sample_token(
11    top_k=top_k,
12    top_p=top_p,
13    adjust_fn=sampling_func,
14)

Where token_data is a dictionary of useful quantities, including candidate distributions before and after adjustment (token_data['candidates_before_adjustment'] and token_data['candidates_after_adjustment'] respectively) and the sampled token, its ID, log-probability and probability in token_data['sampled_token'].

Sampling from the Power Distribution#

Recent developments in LLM sampling methods have pointed out that sharper distributions yield more favorable outputs. Whilst low-temperature sampling goes a way towards solving this issue, it only considers previous tokens outputted when sampling, rather than favoring tokens that may have potential future minima.

At considerable computational expense, sampling from the power distribution can be performed by performing look-ahead runs, and utilizing their probabilities before sampling.

This code implements a general framework for sampling from the power distribution. Let’s start with the code snippet below. In this example we use the built-in Metropolis method to sample potential future branches, inspired by the implementation from Karan and Du.

 1from problm_solver.samplers import MetropolisSampler, SamplePowerDist
 2from problm_solver.llama_interface import Model
 3
 4temp = 0.5
 5top_k = 8
 6peek = 10
 7model = Model('path/to/model.gguf', 'Why is the sky blue?', logits_all=True)
 8
 9sampling_func = SamplePowerDist(
10    alpha=1./temp,
11    lookahead_depth=peek,
12    branch_sampler=MetropolisSampler(equil_branches=5, max_branches=10)
13)
14data = model.generate_with_sampler(
15    top_k=top_k,
16    top_p=0.9,
17    adjust_fn=sampling_func,
18    max_tokens=128
19)

probLM-solver again implements the sampling function as a callable class, in problm_solver.samplers.SamplePowerDist. This takes three main parameters:

  • alpha, which you are already familiar with;

  • lookahead_depth, which controls how many future tokens the branch sampler can generate when assessing probabilities;

  • and branch_sampler, which takes a class which defines the logic for generating and accepting new branches.

We can see that MetropolisSampler is instanced with two variables. equil_branches is the number of branches to generate and then discard from probability and convergence calculations, as the method is still locating a minimum. max_branches is the maximum number of branches to sample for each of the top-k most probable tokens.

A second built-in option is token-level beam expansion:

from problm_solver.samplers import BeamSampler, SamplePowerDist

sampling_func = SamplePowerDist(
    alpha=2.0,
    lookahead_depth=10,
    branch_sampler=BeamSampler(beam_width=10, branch_top_k=5),
)

Creating a Branch Sampler#

If we take a look at the header for MetropolisSampler:

class MetropolisSampler(BranchSampler):

we can see that it inherits from an abstract base class called BranchSampler.

New branch sampling algorithms for generating tokens can be defined by the user based on BranchSampler. It implements four required methods (plus an optional beam hook):

 1class BranchSampler(ABC):
 2
 3    def reset(self) -> None:
 4        """Resets stateful samplers."""
 5
 6    @abstractmethod
 7    def step(
 8        self,
 9        proposed_log_prob: float,
10        alpha: float = 1.0,
11        forward_log_q: float = 0.0,
12        reverse_log_q: float = 0.0,
13    ) -> float:
14        """Processes sampler state and returns accepted branch log-probability."""
15        ...
16
17    @abstractmethod
18    def should_continue(self, branch_log_probs: npt.NDArray[np.float64]) -> bool:
19        """Control-flow logic for deciding when proposal sampling is finished."""
20        ...
21
22    @abstractmethod
23    def future_logprob(self, alpha: float, branch_log_probs: npt.NDArray[np.float64]) -> np.float64:
24        """Aggregates sampled branch log-probabilities into one future score."""
25        ...

  • BranchSampler.reset() ought to reset any state variables (the Metropolis method tracks the current most-favorable logarithmic probability, for instance).

  • BranchSampler.step() controls criteria for accepting/rejecting a new branch log-probability and returns the chosen probability.

  • BranchSampler.should_continue() handles logic for terminating proposal sampling.

  • BranchSampler.future_logprob() combines collected branch log-probabilities into one future score used by SamplePowerDist.

All of these methods are called at some point by Model.generate_with_sampler().

A user who has made the class MySampler, can now pass it to an instance of SamplePowerDist, to create a sampling function.

1from problm_solver.samplers import BranchSampler, SamplePowerDist
2
3class MySampler(BranchSampler):
4    ...
5
6temp = 0.5
7sampling_func = SamplePowerDist(alpha=1./temp, lookahead_depth=10, branch_sampler=MySampler())
8data = model.generate_with_sampler(top_k=8, top_p=0.9, adjust_fn=sampling_func, max_tokens=128)

Let’s explore a hypothetical greedy branch sampler. It might be defined in the following way:

 1from problm_solver.samplers import BranchSampler, SamplePowerDist
 2
 3class MyGreedySampler(BranchSampler):
 4    """My greedy sampler, always picks the most probable branch."""
 5
 6    def __init__(self, max_branches: int = 50):
 7        self._max_branches = max_branches
 8        self._current_log_prob: float | None = None
 9
10    def reset(self) -> None:
11        self._current_log_prob = None
12
13    def step(
14        self,
15        proposed_log_prob: float,
16        alpha: float = 1.0,
17        forward_log_q: float = 0.0,
18        reverse_log_q: float = 0.0,
19    ) -> float:
20        if self._current_log_prob is None or proposed_log_prob > self._current_log_prob:
21            self._current_log_prob = proposed_log_prob
22        return self._current_log_prob
23
24    def should_continue(self, branch_log_probs) -> bool:
25        return len(branch_log_probs) < self._max_branches
26
27    def future_logprob(self, alpha, branch_log_probs):
28        # Greedy: use best branch only.
29        return alpha * max(branch_log_probs)

You can then use this sampler in place of any other.