Fix sampler API: validate TopP args, register, drop gumbel

zeromodels/samplers/__init__.py

@@ -1,5 +1,5 @@
 from zeromodels.samplers.greedy_sampler import GreedySampler
-from zeromodels.samplers.sampler import Sampler, categorical, gumbel
+from zeromodels.samplers.sampler import Sampler, categorical
 from zeromodels.samplers.top_k_sampler import TopKSampler
 from zeromodels.samplers.top_p_sampler import TopPSampler
 
@@ -9,5 +9,4 @@
     "TopKSampler",
     "TopPSampler",
     "categorical",
-    "gumbel",
 ]

zeromodels/samplers/greedy_sampler.py

@@ -1,8 +1,10 @@
+import keras
 from keras import ops
 
 from zeromodels.samplers.sampler import Sampler
 
 
+@keras.saving.register_keras_serializable(package="zeromodels")
 class GreedySampler(Sampler):
     """Deterministic argmax: the default decoding strategy."""
 

zeromodels/samplers/sampler.py

@@ -1,3 +1,4 @@
+import keras
 from keras import ops
 
 # Rejected tokens are pushed here rather than to -inf: finite, so their softmax
@@ -6,6 +7,7 @@
 NEG_INF = -1e9
 
 
+@keras.saving.register_keras_serializable(package="zeromodels")
 class Sampler:
     """Maps logits ``(batch, vocab)`` + per-step uniform ``noise`` to next ids.
 
@@ -24,6 +26,11 @@ class Sampler:
     keeps everything (greedy); ``TopKSampler`` / ``TopPSampler`` override it. It is
     split out from ``sample`` so the kept set can be compared against the reference
     warpers without drawing a token.
+
+    ``get_config`` returns exactly the constructor kwargs, so :meth:`from_config`
+    round-trips every subclass without an override. Besides Keras serialization, the
+    decode engine hashes ``get_config()`` into its compiled-function cache key, so two
+    samplers that differ only in a setting never share a traced function.
     """
 
     stochastic = False
@@ -37,11 +44,9 @@ def filter_logits(self, logits):
     def get_config(self):
         return {}
 
-
-def gumbel(noise):
-    # uniform(0, 1) -> Gumbel(0, 1)
-    u = ops.clip(noise, 1e-9, 1.0)
-    return -ops.log(-ops.log(u))
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def categorical(masked_logits, noise):

zeromodels/samplers/top_k_sampler.py

@@ -1,3 +1,4 @@
+import keras
 from keras import ops
 
 from zeromodels.samplers.sampler import (
@@ -8,6 +9,7 @@
 )
 
 
+@keras.saving.register_keras_serializable(package="zeromodels")
 class TopKSampler(Sampler):
     """Sample from the ``k`` highest-logit tokens (temperature-scaled).
 

zeromodels/samplers/top_p_sampler.py

@@ -1,3 +1,4 @@
+import keras
 from keras import ops
 
 from zeromodels.samplers.sampler import (
@@ -8,22 +9,35 @@
 )
 
 
+@keras.saving.register_keras_serializable(package="zeromodels")
 class TopPSampler(Sampler):
     """Nucleus sampling: the smallest set of top tokens with cumulative prob >= ``p``.
 
     Matches Hugging Face's ``TopPLogitsWarper``: sort by probability, keep tokens
     while the prefix *before* them holds less than ``p`` of the mass (so the token
     that crosses ``p`` is kept too), push the rest to ``NEG_INF``, then draw with an
     inverse-CDF categorical draw on the pre-supplied per-row noise. At least
-    ``min_tokens_to_keep`` (default 1) top tokens are always kept, so ``p <= 0``
-    falls back to greedy rather than masking everything. ``temperature`` must be
-    strictly positive.
+    ``min_tokens_to_keep`` top tokens are always kept.
+
+    Args are validated up front, because an out-of-range value would not raise on its
+    own -- it would silently decode with a different strategy than requested. ``p``
+    must be in ``(0, 1]``, ``min_tokens_to_keep`` must be ``>= 1``, and
+    ``temperature`` must be strictly positive.
     """
 
     stochastic = True
 
     def __init__(self, p=0.9, temperature=1.0, min_tokens_to_keep=1):
         validate_temperature(temperature)
+        if not (0.0 < float(p) <= 1.0):
+            raise ValueError(
+                f"p must be in the range (0, 1], got {p!r}. Use p=1.0 to sample from "
+                "the full distribution, or GreedySampler() for greedy decoding."
+            )
+        if int(min_tokens_to_keep) < 1:
+            raise ValueError(
+                f"min_tokens_to_keep must be >= 1, got {min_tokens_to_keep!r}."
+            )
         self.p = float(p)
         self.temperature = float(temperature)
         self.min_tokens_to_keep = int(min_tokens_to_keep)
@@ -33,9 +47,7 @@ def filter_logits(self, logits):
         sorted_logits = ops.take_along_axis(logits, order, axis=-1)
         probs = ops.softmax(sorted_logits, axis=-1)
         cumulative = ops.cumsum(probs, axis=-1) - probs  # exclusive prefix
-        keep_sorted = cumulative < self.p  # nucleus (sorted order)
-        # Always keep the top ``min_tokens_to_keep``: without this floor ``p <= 0``
-        # keeps nothing and the draw over an all-NEG_INF row is uniform-random.
+        keep_sorted = cumulative < self.p
         floor = ops.arange(logits.shape[-1]) < self.min_tokens_to_keep
         keep_sorted = ops.logical_or(keep_sorted, floor[None, :])
         inverse = ops.argsort(order, axis=-1)  # scatter back to vocab order