Add support for default top_p/top_k (#1267)

This PR enables setting default top_p/top_k values for a model. 

This is critical for users who use reasoning models such as DeepSeek R1,
QWEN3, etc... as using default top_p=1.0 and top_k=None will provide
sub-optimal outputs.

```python
llamacpp_lm = guidance.models.LlamaCpp(
    model="Llama-3.2-3B-Instruct-UD-Q8_K_XL.gguf", 
    echo=False, 
    n_ctx=8192,
    default_sampling_params=SamplingParams(
        top_p=0.95,
        top_k=64
    )
)

openai_lm = models.OpenAI(
    "gpt-3.5-turbo", 
    api_key=os.environ.get("OPENAI_API_KEY"),
    # you can also use dict
    default_sampling_params={
        "top_p" : 0.95
    }
)

vllm_lm = guidance.models.experimental.VLLMModel(
    model="DeepSeek-R1-0528-Qwen3-8B-fp8",
    base_url="http://localhost:5005/v1",
    api_key="NOKEY",
    default_sampling_params={
        "top_p" : 0.95,
        "top_k": 64
    }
)

model_desc = {
        "model_name": "DeepSeek-R1-0528-Qwen3-8B-fp8", 
        "litellm_params": { 
            "model": "hosted_vllm/DeepSeek-R1-0528-Qwen3-8B-fp8", 
            "api_key": os.environ.get("LITELLM_API_KEY", None),
            "api_base": "http://localhost:5005/v1"
        }
    }

litellm_lm = guidance.models.experimental.LiteLLM(
    model_description=model_desc, 
    echo=False, 
    default_sampling_params={
        "top_p" : 0.95,
        "top_k": 64
    }
)
```

Author: JC1DA

guidance/_schema.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal, Optional, Union, Dict
+from typing import Any, Literal, Optional, Union, TypedDict
 
 from pydantic import BaseModel, Field, NonNegativeInt, RootModel, model_validator, computed_field
 from typing_extensions import Annotated
@@ -166,3 +166,7 @@ class LLInterpreterResponse(BaseModel):
     progress: LLProgress
     stop: bool
     temperature: Optional[float]
+
+class SamplingParams(TypedDict):
+    top_p: Optional[float]
+    top_k: Optional[int]

guidance/_utils.py

@@ -9,13 +9,16 @@
 import functools
 import numpy as np
 import logging
-from typing import Union, cast
+from typing import Union, cast, Optional, TYPE_CHECKING
 import pathlib
 import pydantic
 import urllib
 import http
 import re
 
+if TYPE_CHECKING:
+   from ._schema import SamplingParams
+
 logger = logging.getLogger(__name__)
 
 
@@ -340,3 +343,47 @@ def to_utf8_or_bytes_string(_bytes: bytes) -> str:
         return _bytes.decode("utf-8")
     except UnicodeDecodeError:
         return str(_bytes)
+
+def apply_top_k_only(logits: np.ndarray, k: int) -> np.ndarray:
+    if k <= 0:
+        return logits
+    
+    indices_to_remove = logits.argpartition(-k)[:-k]
+    logits[indices_to_remove] = -float("inf")
+    return logits
+    
+
+def apply_top_k_and_top_p_filter(logits: np.ndarray, sampling_params: Optional['SamplingParams']) -> np.ndarray:
+    if sampling_params is None:
+        return logits
+    
+    top_p = sampling_params.get("top_p", None)
+    top_k = sampling_params.get("top_k", None)
+
+    if top_k is None and top_p is None:
+        # No filtering, return logits as is
+        return logits
+
+    if top_k is not None and top_p is None:
+        return apply_top_k_only(logits, top_k)
+
+    # try our best to sort logits one time only
+    sorted_logits = logits.argsort()
+    if top_k is not None and top_k > 0:
+        indices_to_remove = sorted_logits[:-top_k]
+        logits[indices_to_remove] = -float("inf")
+
+    if top_p is not None and top_p > 0:
+        sorted_indices = sorted_logits[::-1]
+        sorted_logits = logits[sorted_indices]
+        probs = softmax(sorted_logits)
+        cumulative_probs = np.cumsum(probs)
+        indices_to_remove = cumulative_probs > top_p
+        if np.any(indices_to_remove):
+            # +1 to keep the first token that exceeds the threshold
+            first_to_remove = np.argmax(indices_to_remove) + 1
+            # make sure we always keep at least one token
+            sorted_indices_to_remove = sorted_indices[max(1, first_to_remove):]
+            logits[sorted_indices_to_remove] = -float("inf")
+        
+    return logits

guidance/models/_base/_interpreter.py

@@ -1,5 +1,7 @@
 import base64
-from typing import Generic, Iterator, TypeVar
+from typing import Generic, Iterator, TypeVar, Optional
+
+from guidance._schema import SamplingParams
 
 from ..._ast import (
     ASTNode,
@@ -29,8 +31,9 @@
 
 
 class Interpreter(Generic[S]):
-    def __init__(self, state: S):
+    def __init__(self, state: S, default_sampling_params: Optional[SamplingParams] = None):
         self.state = state
+        self.default_sampling_params = SamplingParams() if default_sampling_params is None else default_sampling_params
 
     def run(self, node: ASTNode, **kwargs) -> Iterator[OutputAttr]:
         yield from node.simplify()._run(self, **kwargs)

guidance/models/_base/_model.py

@@ -35,6 +35,8 @@
 from ._interpreter import Interpreter
 from ._state import State
 
+from ..._schema import SamplingParams
+
 if TYPE_CHECKING:
     from ...library._block import Block
 
@@ -62,6 +64,7 @@ class Model:
     def __init__(
         self,
         interpreter: Interpreter[S],
+        default_sampling_params: Optional[SamplingParams] = None,
         echo: bool = True,
     ) -> None:
         self.echo = echo

guidance/models/_engine/_engine.py

@@ -15,10 +15,11 @@
     EngineResponse,
     GenToken,
     GuidanceEngineMetrics,
+    SamplingParams,
 )
 
 from ...registry import get_exchange
-from ..._utils import log_cleanup, log_init, softmax
+from ..._utils import log_cleanup, log_init, softmax, apply_top_k_and_top_p_filter
 from ...visual import (
     ExecutionCompletedMessage,
     ExecutionStartedMessage,
@@ -132,6 +133,8 @@ def __call__(
         grammar: str,
         ensure_bos_token: bool = True,
         echo: bool = True,
+        sampling_params: Optional[SamplingParams] = None,
+        **kwargs: Optional[dict],
     ) -> Iterator[EngineResponse]:
         """Main entry point for the inference-parser loop. Yields EngineCallResponse objects as
         the parser advances through the grammar.
@@ -148,6 +151,8 @@ def __call__(
             Grammar (RawFunction or GrammarFunction) used to extend the prompt.
         ensure_bos_token: bool
             Ensures that the prompt ends with the BOS token.
+        sampling_params: Optional[SamplingParams]
+            Additional sampling parameters to apply to the logits.
         """
         # TODO: Pass these to get_logits
         # images = state.images
@@ -331,6 +336,7 @@ def __call__(
                 temperature=ll_response.temperature,
                 k=self._top_k,
                 force_return_unmasked_probs=echo,
+                sampling_params=sampling_params,
             )
             last_temperature = ll_response.temperature
 
@@ -348,6 +354,7 @@ def get_next_token_with_top_k(
         temperature: float,
         k: int = 5,
         force_return_unmasked_probs: bool = False,
+        sampling_params: Optional[SamplingParams] = None,
     ) -> EngineOutput:
         """Get the next token and associated top-k tokens from the engine.
 
@@ -369,6 +376,8 @@ def get_next_token_with_top_k(
             The number of top-k tokens to return.
         force_return_unmasked_probs: bool
             If True, the top-k unmasked probabilities will be returned.
+        sampling_params : Optional[SamplingParams]
+            Additional sampling parameters to apply to the logits.
 
         Returns
         -------
@@ -399,9 +408,9 @@ def get_top_k(_probs: NDArray, _k: int = 5) -> list[GenToken]:
 
         # compute top-k without masking
         probs = (
-            softmax(np.array(logits))
+            softmax(apply_top_k_and_top_p_filter(np.array(logits), sampling_params))
             if temperature < _TEMPERATURE_EPSILON
-            else softmax(np.array(logits) / temperature)
+            else softmax(apply_top_k_and_top_p_filter(np.array(logits) / temperature, sampling_params))
         )
 
         top_k: list[GenToken] = []
@@ -413,11 +422,14 @@ def get_top_k(_probs: NDArray, _k: int = 5) -> list[GenToken]:
         if mask is not None:
             mask = np.frombuffer(mask, dtype=np.uint8)
             masked_logits = np.where(mask != 0, logits, -np.inf)
+            
             if temperature < _TEMPERATURE_EPSILON:
                 masked_logits = np.where(masked_logits == np.max(masked_logits), 0, -np.inf)
             else:
                 masked_logits /= temperature
 
+            masked_logits = apply_top_k_and_top_p_filter(masked_logits, sampling_params)
+
             masked_probs = softmax(masked_logits)
             masked_top_k = get_top_k(masked_probs, k)
 

guidance/models/_engine/_interpreter.py

@@ -1,6 +1,6 @@
 from base64 import b64decode, b64encode
 from io import BytesIO
-from typing import Iterator
+from typing import Iterator, Optional
 from copy import deepcopy
 import re
 
@@ -10,15 +10,15 @@
 from .._base import Interpreter
 from ._engine import Engine, Tokenizer
 from ._state import EngineState
-from ..._schema import GenTokenExtra
+from ..._schema import GenTokenExtra, SamplingParams
 
 
 class EngineInterpreter(Interpreter[EngineState]):
-    def __init__(self, engine: Engine):
-        self.state = EngineState()
+    def __init__(self, engine: Engine, default_sampling_params: Optional[SamplingParams] = None):
+        super().__init__(state=EngineState(), default_sampling_params=default_sampling_params)
         self.engine = engine
         self.chat_template = self.engine.get_chat_template()
-
+        
     def __deepcopy__(self, memo):
         """Custom deepcopy to ensure engine is not copied."""
         cls = self.__class__
@@ -71,6 +71,7 @@ def grammar(self, node: GrammarNode, **kwargs) -> Iterator[OutputAttr]:
             grammar=node.ll_grammar(),
             ensure_bos_token=True,
             echo=False,
+            sampling_params=self.default_sampling_params # NOTE: passing default sampling params for now
         )
 
         delayed_bytes = b""

guidance/models/_llama_cpp.py

@@ -5,12 +5,12 @@
 import sys
 from itertools import takewhile
 from pathlib import Path
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Union, Optional
 import ctypes
 
 import numpy as np
 
-from .._schema import GenToken, GenTokenExtra
+from .._schema import GenToken, GenTokenExtra, SamplingParams
 from .._utils import normalize_notebook_stdout_stderr, softmax
 from ..chat import ChatTemplate
 from ._base import Model
@@ -232,7 +232,8 @@ def __init__(
         enable_backtrack=True,
         enable_ff_tokens=True,
         enable_monitoring=True,
-        **llama_cpp_kwargs,
+        default_sampling_params: Optional[SamplingParams] = None,
+        **llama_cpp_kwargs
     ):
         """Build a new LlamaCpp model object that represents a model in a given state."""
 
@@ -245,5 +246,5 @@ def __init__(
             enable_monitoring=enable_monitoring,
             **llama_cpp_kwargs,
         )
-        interpreter = EngineInterpreter(engine)
+        interpreter = EngineInterpreter(engine, default_sampling_params=default_sampling_params)
         super().__init__(interpreter=interpreter, echo=echo)

guidance/models/_mock.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from .._schema import EngineOutput
+from .._schema import EngineOutput, SamplingParams
 from ..trace import TraceHandler
 from ..visual._renderer import DoNothingRenderer
 from ._base import Model
@@ -106,10 +106,12 @@ def get_next_token_with_top_k(
         temperature: float,
         k: int = 1,
         force_return_unmasked_probs: bool = False,
+        sampling_params: Optional[SamplingParams] = None,
+        
     ) -> EngineOutput:
         self.called_temperatures.append(temperature)
         return super().get_next_token_with_top_k(
-            logits, logits_lat_ms, token_ids, mask, temperature, k, force_return_unmasked_probs
+            logits, logits_lat_ms, token_ids, mask, temperature, k, force_return_unmasked_probs, sampling_params
         )
 
     def get_logits(self, token_ids: list[int], include_all_uncached_tokens: bool = False) -> np.ndarray:
@@ -161,6 +163,7 @@ class Mock(Model):
     def __init__(
         self,
         byte_patterns=[],
+        default_sampling_params: Optional[SamplingParams] = None,
         echo=False,
         force=False,
         **kwargs,

guidance/models/_openai.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from guidance._schema import SamplingParams
+
 
 from ._base import Model
 from ._openai_base import (
@@ -16,7 +18,8 @@
 class OpenAIInterpreter(OpenAIRuleMixin, OpenAIJSONMixin, OpenAIRegexMixin, BaseOpenAIInterpreter):
     def __init__(
         self,
-        model: str,
+        model: str,        
+        default_sampling_params: Optional[SamplingParams],
         api_key: Optional[str] = None,
         **kwargs,
     ):
@@ -26,15 +29,17 @@ def __init__(
             raise Exception(
                 "Please install the openai package version >= 1 using `pip install openai -U` in order to use guidance.models.OpenAI!"
             )
+                    
         client = openai.OpenAI(api_key=api_key, **kwargs)
-        super().__init__(model=model, client=OpenAIClientWrapper(client))
+        super().__init__(model=model, client=OpenAIClientWrapper(client), default_sampling_params=default_sampling_params)
 
 
 class OpenAI(Model):
     def __init__(
         self,
         model: str,
-        echo: bool = True,
+        default_sampling_params: Optional[SamplingParams] = None,
+        echo: bool = True,        
         *,
         api_key: Optional[str] = None,
         **kwargs,
@@ -66,4 +71,4 @@ def __init__(
         else:
             interpreter_cls = OpenAIInterpreter
 
-        super().__init__(interpreter=interpreter_cls(model, api_key=api_key, **kwargs), echo=echo)
+        super().__init__(interpreter=interpreter_cls(model, api_key=api_key, default_sampling_params=default_sampling_params, **kwargs), echo=echo)

guidance/models/_openai_base.py

@@ -9,6 +9,8 @@
 from pydantic import BaseModel, Discriminator, Field, TypeAdapter
 from typing_extensions import Annotated, assert_never
 
+from guidance._schema import SamplingParams
+
 from .._ast import (
     ASTNode,
     GenAudio,
@@ -30,6 +32,7 @@
     import openai
     from openai.types.chat import ChatCompletionChunk
     from openai.types.chat.chat_completion_chunk import ChoiceLogprobs
+    from openai import OpenAI
 
 
 def get_role_start(role: str) -> str:
@@ -183,6 +186,7 @@ def streaming_chat_completions(
         **kwargs,
     ) -> ContextManager[Iterator["ChatCompletionChunk"]]:
         """Streaming chat completions."""
+
         return self.client.chat.completions.create(
             model=model,
             messages=TypeAdapter(list[Message]).dump_python(messages),  # type: ignore[arg-type]
@@ -200,8 +204,10 @@ def __init__(
         self,
         model: str,
         client: BaseOpenAIClientWrapper,
+        default_sampling_params: Optional[SamplingParams] = None,
+        **kwargs
     ):
-        self.state = OpenAIState()
+        super().__init__(state=OpenAIState(), default_sampling_params=default_sampling_params)
         self.model = model
         self.client = client
 
@@ -249,6 +255,10 @@ def _run(self, **kwargs) -> Iterator[OutputAttr]:
             raise ValueError(
                 f"OpenAI models do not support pre-filled assistant messages: got data {self.state.content}."
             )
+            
+        # only process kwargs that are supported by the OpenAI API
+        if "top_p" not in kwargs and "top_p" in self.default_sampling_params:
+            kwargs["top_p"] = self.default_sampling_params["top_p"]
 
         with self.client.streaming_chat_completions(
             model=self.model,