Improve SamplingParams (#1283)

This PR:
1) moves sampling_params out of interpreter
2) allows creating new model with updated sampling_params

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

lm = llamacpp_lm.with_sampling_params(SamplingParams(top_p=0.8, top_k=128))
```

Author: JC1DA

guidance/library/_gen.py

@@ -96,7 +96,7 @@ def gen(
         raise ValueError("Cannot use `regex` with `tools`")
 
     assert n == 1, "We still need to add support for n>1! Consider putting your gen call in a loop for now."
-    assert top_p == 1, "We still need to add support for top_p != 1!"
+    assert top_p == 1, "Please use `model.with_sampling_params` to set top_p."
 
     logger.debug(f'start gen(name="{name}")')
 

guidance/models/_base/_interpreter.py

@@ -31,9 +31,8 @@
 
 
 class Interpreter(Generic[S]):
-    def __init__(self, state: S, default_sampling_params: Optional[SamplingParams] = None):
+    def __init__(self, state: S):
         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

@@ -59,7 +59,7 @@ class Model:
     def __init__(
         self,
         interpreter: Interpreter[S],
-        default_sampling_params: Optional[SamplingParams] = None,
+        sampling_params: SamplingParams,
         echo: bool = True,
     ) -> None:
         self.echo = echo
@@ -70,6 +70,7 @@ def __init__(
 
         self._interpreter = interpreter
         self._active_blocks: dict[Block, int] = {}
+        self.sampling_params: SamplingParams = sampling_params
 
         self._parent: Optional["Model"] = None
         self._parent_id: Optional[int] = None
@@ -131,7 +132,8 @@ def _apply_node(self, node: ASTNode) -> Self:
         else:
             self._update_trace_node(self._id, self._parent_id, StatelessGuidanceInput(value=node))
 
-        for i, output_attr in enumerate(self._interpreter.run(node)):
+        # NOTE: passing a copy of the sampling parameters to avoid modifying the original
+        for i, output_attr in enumerate(self._interpreter.run(node, sampling_params=self.sampling_params.copy())):
             if i != 0:
                 # On the first iteration, we already have a fresh trace node
                 # TODO: should be allowed to associate multiple output_attrs with a single input node?
@@ -291,6 +293,12 @@ def log_prob(self, key: str, default: Optional[D] = None) -> Union[float, list[U
         else:
             return captures["log_prob"]
 
+    def with_sampling_params(self, sampling_params: SamplingParams) -> Self:
+        """Return a new model with the given sampling parameters set."""
+        self = self.copy()
+        self.sampling_params = sampling_params
+        return self
+
     def __getattribute__(self, name):
         if name == "engine":
             # For legacy model.engine access (mostly for tests...)

guidance/models/_engine/_interpreter.py

@@ -14,8 +14,8 @@
 
 
 class EngineInterpreter(Interpreter[EngineState]):
-    def __init__(self, engine: Engine, default_sampling_params: Optional[SamplingParams] = None):
-        super().__init__(state=EngineState(), default_sampling_params=default_sampling_params)
+    def __init__(self, engine: Engine):
+        super().__init__(state=EngineState())
         self.engine = engine
         self.chat_template = self.engine.get_chat_template()
 
@@ -67,7 +67,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
+            sampling_params=kwargs.pop("sampling_params", None),
         )
 
         delayed_bytes = b""

guidance/models/_llama_cpp.py

@@ -230,7 +230,7 @@ def __init__(
         enable_backtrack=True,
         enable_ff_tokens=True,
         enable_monitoring=True,
-        default_sampling_params: Optional[SamplingParams] = None,
+        sampling_params: Optional[SamplingParams] = None,
         **llama_cpp_kwargs,
     ):
         """Build a new LlamaCpp model object that represents a model in a given state."""
@@ -244,5 +244,9 @@ def __init__(
             enable_monitoring=enable_monitoring,
             **llama_cpp_kwargs,
         )
-        interpreter = EngineInterpreter(engine, default_sampling_params=default_sampling_params)
-        super().__init__(interpreter=interpreter, echo=echo)
+        interpreter = EngineInterpreter(engine)
+        super().__init__(
+            interpreter=interpreter,
+            sampling_params=SamplingParams() if sampling_params is None else sampling_params,
+            echo=echo,
+        )

guidance/models/_mock.py

@@ -160,7 +160,7 @@ class Mock(Model):
     def __init__(
         self,
         byte_patterns=[],
-        default_sampling_params: Optional[SamplingParams] = None,
+        sampling_params: Optional[SamplingParams] = None,
         echo=False,
         force=False,
         **kwargs,
@@ -178,6 +178,7 @@ def __init__(
         super().__init__(
             interpreter=EngineInterpreter(engine),
             echo=echo,
+            sampling_params=SamplingParams() if sampling_params is None else sampling_params,
         )
 
 

guidance/models/_openai.py

@@ -18,7 +18,6 @@ class OpenAIInterpreter(OpenAIRuleMixin, OpenAIJSONMixin, OpenAIRegexMixin, Base
     def __init__(
         self,
         model: str,
-        default_sampling_params: Optional[SamplingParams],
         api_key: Optional[str] = None,
         **kwargs,
     ):
@@ -30,16 +29,14 @@ def __init__(
             )
 
         client = openai.OpenAI(api_key=api_key, **kwargs)
-        super().__init__(
-            model=model, client=OpenAIClientWrapper(client), default_sampling_params=default_sampling_params
-        )
+        super().__init__(model=model, client=OpenAIClientWrapper(client))
 
 
 class OpenAI(Model):
     def __init__(
         self,
         model: str,
-        default_sampling_params: Optional[SamplingParams] = None,
+        sampling_params: Optional[SamplingParams] = None,
         echo: bool = True,
         *,
         api_key: Optional[str] = None,
@@ -69,8 +66,7 @@ def __init__(
             interpreter_cls = OpenAIInterpreter
 
         super().__init__(
-            interpreter=interpreter_cls(
-                model, api_key=api_key, default_sampling_params=default_sampling_params, **kwargs
-            ),
+            interpreter=interpreter_cls(model, api_key=api_key, **kwargs),
+            sampling_params=SamplingParams() if sampling_params is None else sampling_params,
             echo=echo,
         )

guidance/models/_openai_base.py

@@ -201,14 +201,8 @@ class BaseOpenAIInterpreter(Interpreter[OpenAIState]):
 
     log_probs: bool = True
 
-    def __init__(
-        self,
-        model: str,
-        client: BaseOpenAIClientWrapper,
-        default_sampling_params: Optional[SamplingParams] = None,
-        **kwargs,
-    ):
-        super().__init__(state=OpenAIState(), default_sampling_params=default_sampling_params)
+    def __init__(self, model: str, client: BaseOpenAIClientWrapper, **kwargs):
+        super().__init__(state=OpenAIState())
         self.model = model
         self.client = client
 
@@ -251,9 +245,14 @@ def _run(self, **kwargs) -> Iterator[OutputAttr]:
                 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"]
+        sampling_params = kwargs.pop("sampling_params", None)
+        if sampling_params:
+            # only process kwargs that are supported by the OpenAI API
+            if "top_p" not in kwargs:
+                kwargs["top_p"] = sampling_params.get("top_p", None)
+
+            if sampling_params.get("top_k", None) is not None:
+                raise ValueError("OpenAI models do not support top_k sampling.")
 
         with self.client.streaming_chat_completions(
             model=self.model,
@@ -274,7 +273,8 @@ def _handle_stream(self, chunks: Iterator["ChatCompletionChunk"]) -> Iterator[Ou
             latency_ms = (t1 - t0) * 1000
             t0 = t1
 
-            if chunk.usage is not None:
+            # NOTE: use getattr here as litellm does not return usage
+            if getattr(chunk, "usage", None) is not None:
                 # Update token usage
                 usage.input_tokens += chunk.usage.prompt_tokens
                 # Estimate forward passes as number of completion tokens

guidance/models/_transformers.py

@@ -618,7 +618,7 @@ def __init__(
         enable_backtrack=True,
         enable_ff_tokens=True,
         enable_monitoring=True,
-        default_sampling_params: Optional[SamplingParams] = None,
+        sampling_params: Optional[SamplingParams] = None,
         **kwargs,
     ):
         """Build a new Transformers model object that represents a model in a given state."""
@@ -640,10 +640,10 @@ def __init__(
                 enable_ff_tokens=enable_ff_tokens,
                 enable_monitoring=enable_monitoring,
                 **kwargs,
-            ),
-            default_sampling_params=default_sampling_params,
+            )
         )
         super().__init__(
             interpreter=client,
+            sampling_params=SamplingParams() if sampling_params is None else sampling_params,
             echo=echo,
         )

guidance/models/experimental/_litellm.py

@@ -62,7 +62,7 @@ def streaming_chat_completions(
 class LiteLLMInterpreter(BaseOpenAIInterpreter):
     SUPPORTED_ENDPOINT_TYPES = ["openai", "azure_ai", "azure", "gemini", "anthropic", "xai", "hosted_vllm"]
 
-    def __init__(self, model_description: dict, default_sampling_params: Optional[SamplingParams], **kwargs):
+    def __init__(self, model_description: dict, **kwargs):
         try:
             import litellm
         except ImportError:
@@ -80,9 +80,7 @@ def __init__(self, model_description: dict, default_sampling_params: Optional[Sa
         # Otherwise, generation will fail for some endpoints.
         self.log_probs = False
 
-        super().__init__(
-            model=self.model, client=self.client, default_sampling_params=default_sampling_params, **kwargs
-        )
+        super().__init__(model=self.model, client=self.client, **kwargs)
 
     def _check_model(self, model_desc: dict) -> str:
         """Check if the model description is valid."""
@@ -232,24 +230,23 @@ def _grammar_vllm(self, node: GrammarNode, **kwargs) -> Iterator[OutputAttr]:
                     yield self.state.apply_capture(name=name, value=value, log_prob=log_probs, is_append=False)
 
     def _process_kwargs(self, **kwargs):
-        if "top_k" not in kwargs and "top_k" in self.default_sampling_params:
-            kwargs["top_k"] = self.default_sampling_params["top_k"]
+        sampling_params = kwargs.pop("sampling_params", None)
+        if sampling_params is None:
+            return kwargs
+
+        kwargs["top_p"] = sampling_params.pop("top_p", None)
+        kwargs["top_k"] = sampling_params.pop("top_k", None)
 
         return kwargs
 
 
 class LiteLLM(Model):
     def __init__(
-        self,
-        model_description: dict,
-        default_sampling_params: Optional[SamplingParams] = None,
-        echo: bool = True,
-        **kwargs,
+        self, model_description: dict, sampling_params: Optional[SamplingParams] = None, echo: bool = True, **kwargs
     ):
-        interpreter = LiteLLMInterpreter(
-            model_description=model_description, default_sampling_params=default_sampling_params, **kwargs
-        )
+        interpreter = LiteLLMInterpreter(model_description=model_description, **kwargs)
         super().__init__(
             interpreter=interpreter,
+            sampling_params=SamplingParams() if sampling_params is None else sampling_params,
             echo=echo,
         )