Spaces:
Paused
Paused
| # Copyright 2026 DeepMind Technologies Limited. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """Sampling methods.""" | |
| import abc | |
| import dataclasses | |
| from etils import enp | |
| import jax | |
| import jax.numpy as jnp | |
| from kauldron.ktyping import Float, Int, PRNGKey, typechecked # pylint: disable=g-multiple-import,g-importing-member | |
| class SamplingMethod(abc.ABC): | |
| """Base class for sampling methods.""" | |
| def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']: | |
| """Returns the next tokens to generate. | |
| Args: | |
| logits: Logits, as returned by the model (i.e. before softmax). | |
| rng: A random key. | |
| Returns: | |
| The next tokens to generate. | |
| """ | |
| raise NotImplementedError() | |
| class Greedy(SamplingMethod): | |
| """Greedy sampling.""" | |
| def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']: | |
| del rng | |
| return jnp.argmax(logits, axis=-1) | |
| class RandomSampling(SamplingMethod): | |
| """Simple random sampling.""" | |
| temperature: float = 1.0 | |
| def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']: | |
| return jax.random.categorical(rng, logits / self.temperature, axis=-1) | |
| class TopkSampling(SamplingMethod): | |
| """Top-k sampling.""" | |
| temperature: float = 1.0 | |
| k: int = 1 | |
| def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']: | |
| logits, batch_shape = enp.flatten(logits, '... V') | |
| batch_size = logits.shape[0] | |
| topk_values, topk_indices = jax.lax.top_k(logits, self.k) | |
| sampled_topk_indices = jax.random.categorical( | |
| rng, topk_values / self.temperature, axis=-1 | |
| ) | |
| batch_indices = jnp.arange(batch_size) | |
| topk_indices = topk_indices[batch_indices, sampled_topk_indices] | |
| return enp.unflatten(topk_indices, batch_shape, '...') | |
| class TopPSampling(SamplingMethod): | |
| """Top-p (Nucleus) Sampling.""" | |
| p: float = 0.9 | |
| temperature: float = 1.0 | |
| def get_next_tokens(self, logits: Float['... V'], rng: PRNGKey) -> Int['...']: # pyrefly: ignore[bad-index] | |
| # temperature scaling | |
| logits = logits / self.temperature | |
| if self.p < 1.0: | |
| sorted_logits = jnp.sort(logits, axis=-1, descending=True) | |
| cumulative_probs = jnp.cumsum( | |
| jax.nn.softmax(sorted_logits, axis=-1), axis=-1 | |
| ) | |
| # get the index of the first token with cumulative probability >= p. | |
| cutoff_index = jnp.sum(cumulative_probs < self.p, axis=-1, keepdims=True) | |
| # get the logit value where the cutoff is. | |
| cutoff_logit = jnp.take_along_axis(sorted_logits, cutoff_index, axis=-1) | |
| # select logit values that are smaller than the cutoff logit. | |
| logits = jnp.where( | |
| logits < cutoff_logit, | |
| jnp.finfo(logits.dtype).min, | |
| logits, | |
| ) | |
| return jax.random.categorical(rng, logits, axis=-1) | |