decoding.generators

  1"""
  2Generators are the user's highest-level interface to the decoding library.
  3By composing instances of `decoding.models.LanguageModel`, `decoding.scorers.Scorer`,
  4and control flow parameters that specify sync, stop, and search conditions, users can
  5implement a wide variety of decoding algorithms with very little code.
  6
  7The `BestOfN` and `TreeSearch` generators are currently fully supported. There is also
  8experimental support for `RolloutTreeSearch` in the `decoding.experimental` module,
  9which supports a simple wrapper interface for a more standard Monte Carlo Tree Search
 10(MCTS) algorithm. It is also on the roadmap to bring twisted `SMC` to the
 11`decoding` library.
 12
 13**NB**: The examples below are illustrative of the API, but not particularly useful.
 14See the [`examples`](https://github.com/benlipkin/decoding/tree/main/examples)
 15directory for more interesting examples.
 16"""
 17
 18from collections.abc import Callable
 19from dataclasses import dataclass
 20
 21from vllm.sampling_params import LogitsProcessor, SamplingParams
 22from vllm.transformers_utils.tokenizers import MistralTokenizer
 23
 24from decoding.models import LanguageModel
 25from decoding.pmf import LogPMF, ScoredItem, sort_scored_items, topk_scored_items
 26from decoding.scorers import Scorer
 27
 28
 29@dataclass(frozen=True, kw_only=True)
 30class _SearchParams:
 31    n: int
 32    width: int
 33    max_steps: int
 34    stop_pass: Callable[[str], bool]
 35    stop_fail: Callable[[str], bool]
 36
 37
 38def BestOfN(  # noqa: PLR0913
 39    *,
 40    prompt: str,
 41    llm: LanguageModel,
 42    scorer: Scorer,
 43    n: int = 1,
 44    min_tokens: int = 0,
 45    max_tokens: int | None = None,
 46    stop_str: list[str] | str | None = None,
 47    stop_token_ids: list[int] | str | None = None,
 48    include_stop_str_in_output: bool = True,
 49    track_logprobs: bool = False,
 50    temperature: float = 1.0,
 51    logits_processors: list[LogitsProcessor] | None = None,
 52    seed: int | None = None,
 53) -> list[ScoredItem[str]]:
 54    """
 55    Generate `n` samples from the language model `llm` using the `scorer` to rank them.
 56    See the [`vLLM.SamplingParams`](https://docs.vllm.ai/en/latest/dev/sampling_params.html)
 57    docs to learn more about some of these parameters such as `logits_processors`.
 58
 59    Args:
 60        prompt: The input prompt string.
 61        llm: The language model to generate samples from.
 62        scorer: The scorer to rank the samples.
 63        n: The number of samples to generate.
 64        min_tokens: The minimum number of tokens in each sample.
 65        max_tokens: The maximum number of tokens in each sample.
 66        stop_str: A string or list of strings that, if generated, will stop decoding.
 67        stop_token_ids: A list of token IDs that, if generated, will stop decoding.
 68            A string can also be passed, which will specify all token IDs that contain
 69            that substring.
 70        include_stop_str_in_output: Whether to include the stop string in the output.
 71        track_logprobs: Whether to track log probabilities. This comes at a performance
 72            cost, so it is off by default. In most cases, as you are alrady sampling
 73            from the model, you do not want to double count the probabilities in the
 74            scorer anyways.
 75        temperature: The temperature for sampling.
 76        logits_processors: A list of logits processors.
 77        seed: The random seed.
 78
 79    Returns:
 80        A list of `decoding.pmf.ScoredItem` objects sorted by the `scorer`.
 81
 82    Raises:
 83        ValueError: If any of the argument configurations are invalid.
 84
 85    Examples:
 86        ```python
 87        from decoding.generators import BestOfN
 88        from decoding.models import LanguageModel
 89        from decoding.scorers import Scorer
 90
 91        llm = LanguageModel.from_id("gpt2")
 92        scorer = Scorer.from_f_str_to_num(lambda x: -len(x))
 93        samples = BestOfN(
 94            prompt="The",
 95            llm=llm,
 96            scorer=scorer,
 97            n=20,
 98            stop_str=".",
 99            seed=42,
100        )
101        assert len(samples) == 20
102        assert all(s.item.endswith(".") for s in samples)
103        assert all(s.score == -len(s.item) for s in samples)
104        assert samples[0].score >= samples[-1].score
105        ```
106
107    """
108    sampling_params = SamplingParams(
109        n=_guard_positive_int(n),
110        min_tokens=min_tokens,
111        max_tokens=max_tokens,
112        stop=stop_str,
113        stop_token_ids=_prepare_token_ids(stop_token_ids, llm=llm),
114        include_stop_str_in_output=include_stop_str_in_output,
115        logprobs=_prepare_track_logprobs(track_logprobs),
116        prompt_logprobs=_prepare_track_logprobs(track_logprobs),
117        temperature=temperature,
118        logits_processors=logits_processors,
119        seed=seed,
120        **_default_sampling_kwargs,  # type: ignore[reportArgumentType]
121    )
122    samples = _BestOfN([prompt], llm, scorer, sampling_params)
123    return sort_scored_items(samples)
124
125
126def TreeSearch(  # noqa: PLR0913
127    *,
128    prompt: str,
129    llm: LanguageModel,
130    step_scorer: Scorer,
131    final_scorer: Scorer | None = None,
132    stop_cond_pass: Callable[[str], bool],
133    stop_cond_fail: Callable[[str], bool] | None = None,
134    n: int = 1,
135    beam_width: int = 1,
136    beam_factor: int = 1,
137    max_steps: int | None = None,
138    min_tokens_per_step: int = 0,
139    max_tokens_per_step: int | None = None,
140    sync_str: list[str] | str | None = None,
141    sync_token_ids: list[int] | str | None = None,
142    include_sync_str_in_output: bool = True,
143    track_logprobs: bool = False,
144    temperature: float = 1.0,
145    logits_processors: list[LogitsProcessor] | None = None,
146    seed: int | None = None,
147) -> list[ScoredItem[str]]:
148    """
149    Generate `n` samples from the language model `llm` using the `step_scorer` to
150    rank them at each sync step and the `final_scorer` to rank the final beam.
151
152    Args:
153        prompt: The input prompt string.
154        llm: The language model to generate samples from.
155        step_scorer: The scorer to rank the samples at each sync step.
156        final_scorer: The scorer to rank the final beam.
157        stop_cond_pass: A function that returns `True` if the sample should pass.
158            This stops the sample from being extended.
159        stop_cond_fail: A function that returns `True` if the sample should fail.
160            This filters the sample from the live beam.
161        n: The number of passing samples to generate before returning.
162        beam_width: The width of the beam. This is the number of samples to
163            keep at each step.
164        beam_factor: The branching factor of the beam. This is the number of
165            new samples to generate from each live sample at each sync step.
166        max_steps: The maximum number of sync steps to take.
167        min_tokens_per_step: The minimum number of tokens in each step's extension.
168        max_tokens_per_step: The maximum number of tokens in each step's extension.
169        sync_str: A string or list of strings that, if generated, will stop extending
170            each sample in the live beam and await scoring, ranking, and filtering.
171        sync_token_ids: A list of token IDs that, if generated, will stop extending
172            each sample in the live beam and await scoring, ranking, and filtering.
173            A string can also be passed, which will specify all token IDs that contain
174            that substring.
175        include_sync_str_in_output: Whether to include the stop string in the output.
176        track_logprobs: Whether to track log probabilities. This comes at a performance
177            cost, so it is off by default. In most cases, as you are already sampling
178            from the model, you do not want to double count the probabilities in the
179            scorer anyways.
180        temperature: The temperature for sampling.
181        logits_processors: A list of logits processors.
182            NB: This is applied within each step as opposed to globally.
183        seed: The random seed.
184
185    Returns:
186        A list of `decoding.pmf.ScoredItem` objects sorted by the `final_scorer`.
187
188    Raises:
189        ValueError: If any of the argument configurations are invalid
190        RuntimeError: if all live samples in the beam fail,
191            or if max steps is reached before any samples pass.
192
193    Examples:
194        ```python
195        from decoding.generators import TreeSearch
196        from decoding.models import LanguageModel
197        from decoding.pmf import ScoredItem
198        from decoding.scorers import Scorer
199
200        def f(x):
201            if "." in x:
202                x = x.split(".")[0] + "."
203            return ScoredItem(item=x, score=-len(x))
204
205        llm = LanguageModel.from_id("gpt2")
206        scorer = Scorer.from_f_str_to_sample(f)
207        samples = TreeSearch(
208            prompt="The",
209            sync_token_ids=" ",
210            stop_cond_pass=lambda x: x.endswith("."),
211            llm=llm,
212            step_scorer=scorer,
213            final_scorer=scorer,
214            n=3,
215            beam_width=50,
216            beam_factor=5,
217            seed=42,
218        )
219        assert len(samples) == 3
220        assert all(s.item.endswith(".") for s in samples)
221        assert all(s.score == -len(s.item) for s in samples)
222        assert samples[0].score >= samples[-1].score
223        ```
224
225    """
226    if final_scorer is None:
227        final_scorer = step_scorer
228    search_params = _SearchParams(
229        n=_guard_positive_int(n),
230        width=_guard_positive_int(beam_width),
231        max_steps=_prepare_max_steps(max_steps),
232        stop_pass=_prepare_stop(stop_cond_pass),
233        stop_fail=_prepare_stop(stop_cond_fail),
234    )
235    _validate_search_params(search_params)
236    sampling_params = SamplingParams(
237        n=_guard_positive_int(beam_factor),
238        min_tokens=min_tokens_per_step,
239        max_tokens=max_tokens_per_step,
240        stop=sync_str,
241        stop_token_ids=_prepare_token_ids(sync_token_ids, llm=llm),
242        include_stop_str_in_output=include_sync_str_in_output,
243        logprobs=_prepare_track_logprobs(track_logprobs),
244        prompt_logprobs=_prepare_track_logprobs(track_logprobs),
245        temperature=temperature,
246        logits_processors=logits_processors,
247        seed=seed,
248        **_default_sampling_kwargs,  # type: ignore[reportArgumentType]
249    )
250    samples = _TreeSearch([prompt], llm, step_scorer, search_params, sampling_params)
251    return sort_scored_items(final_scorer(LogPMF.from_samples(samples)))
252
253
254def _BestOfN(
255    prompts: list[str],
256    llm: LanguageModel,
257    scorer: Scorer,
258    sampling_params: SamplingParams,
259) -> list[ScoredItem[str]]:
260    return scorer(llm(prompts=prompts, params=sampling_params))
261
262
263def _TreeSearch(
264    prompts: list[str],
265    llm: LanguageModel,
266    scorer: Scorer,
267    search_params: _SearchParams,
268    sampling_params: SamplingParams,
269) -> list[ScoredItem[str]]:
270    beam = [ScoredItem(item=p, score=-float("inf")) for p in prompts]
271    passing = []
272    for _ in range(search_params.max_steps):
273        stop_pass = [search_params.stop_pass(s.item) for s in beam]
274        stop_fail = [search_params.stop_fail(s.item) for s in beam]
275        passing = []
276        prompts = []
277        for sample, passed, failed in zip(beam, stop_pass, stop_fail, strict=True):
278            if passed and not failed:
279                passing.append(sample)
280            elif not failed:
281                prompts.append(sample.item)
282            else:  # failed
283                pass
284        if len(passing) >= search_params.n:
285            return passing
286        if len(prompts) == 0:
287            return _handle_failed_beam(passing)
288        live = _BestOfN(prompts, llm, scorer, sampling_params)
289        beam = passing + live
290        if len(beam) > search_params.width:
291            beam = topk_scored_items(beam, search_params.width)
292    return _handle_maxsteps(passing)
293
294
295def _prepare_token_ids(
296    token_ids: list[int] | str | None, *, llm: LanguageModel
297) -> list[int] | None:
298    if isinstance(token_ids, str):
299        return _get_token_ids_from_delimiter(llm=llm, delimiter=token_ids)
300    return token_ids
301
302
303def _get_token_ids_from_delimiter(*, llm: LanguageModel, delimiter: str) -> list[int]:
304    _validate_delimiter(delimiter)
305    tokenizer = llm.tokenizer
306    if isinstance(tokenizer, MistralTokenizer):
307        msg = "vLLM Mistral tokenizer does not currently support `batch_decode`."
308        raise NotImplementedError(msg)
309    tokens = list(tokenizer.get_vocab().values())
310    strs = tokenizer.batch_decode(tokens)
311    return [tokens[i] for i, s in enumerate(strs) if delimiter in s]
312
313
314def _validate_search_params(params: _SearchParams) -> None:
315    if params.n > params.width:
316        msg = "`beam_width` cannot be less than `n`."
317        raise ValueError(msg)
318
319
320def _validate_delimiter(delimiter: str) -> None:
321    if len(delimiter) != 1:
322        msg = f"Delimiter must be a single character, got: {delimiter}."
323        raise ValueError(msg)
324
325
326def _prepare_stop(
327    stop: Callable[[str], bool] | None,
328) -> Callable[[str], bool]:
329    if stop is None:
330
331        def _dont_stop(_: str) -> bool:
332            return False
333
334        return _dont_stop
335    return stop
336
337
338def _prepare_max_steps(max_steps: int | None) -> int:
339    if max_steps is None:
340        return 2**32
341    return _guard_positive_int(max_steps)
342
343
344def _prepare_track_logprobs(track_logprobs: bool) -> int | None:  # noqa: FBT001
345    return 0 if track_logprobs else None
346
347
348def _guard_positive_int(n: int) -> int:
349    if n < 1:
350        msg = f"Expected a positive integer, got: {n}."
351        raise ValueError(msg)
352    return n
353
354
355def _handle_failed_beam(passing: list[ScoredItem[str]]) -> list[ScoredItem[str]]:
356    if len(passing) == 0:
357        msg = "All live samples failed before any passed stop conditions."
358        msg += " Check compatibility of stop conditions or expand search."
359        raise RuntimeError(msg)
360    import warnings
361
362    msg = "All live samples failed before completing search,"
363    msg += " but some completed samples have already passed stopping conditions."
364    msg += " Returning available passing samples."
365    warnings.warn(msg, stacklevel=2)
366    return passing
367
368
369def _handle_maxsteps(passing: list[ScoredItem[str]]) -> list[ScoredItem[str]]:
370    if len(passing) == 0:
371        msg = "Max steps reached, and no samples passed stop conditions."
372        raise RuntimeError(msg)
373    import warnings
374
375    msg = "Max steps reached before completing search,"
376    msg += "but some samples have already passed stopping conditions."
377    msg += " Returning available passing samples."
378    warnings.warn(msg, stacklevel=2)
379    return passing
380
381
382_default_sampling_kwargs = {
383    "detokenize": True,
384    "ignore_eos": False,
385    "truncate_prompt_tokens": None,
386}