keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl

keras_hub/src/samplers/sampler.py

@@ -0,0 +1,237 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+import keras
+from keras import ops
+from keras import random
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.utils.tensor_utils import any_equal
+
+
+@keras_hub_export("keras_hub.samplers.Sampler")
+class Sampler:
+    """Base sampler class.
+
+    Args:
+        temperature: float. optional. Used to control the
+            randomness of the sampling. The higher the temperature, the
+            more diverse the samples. Defaults to `1.0`.
+
+    Call arguments:
+        {{call_args}}
+
+    This base class can be extended to implement different auto-regressive
+    sampling methods. To do so, override the `get_next_token()` method, which
+    computes the next token based on a probability distribution over all
+    possible vocab entries.
+
+    Example:
+
+    ```python
+    causal_lm = keras_hub.models.GPT2CausalLM.from_preset("gpt2_base_en")
+
+    # Greedy search with some tokens forbidden.
+    class CustomSampler(keras_hub.samplers.Sampler):
+        def __init__(self, forbidden_tokens, **kwargs):
+            super().__init__(**kwargs)
+            self.forbidden_tokens = forbidden_tokens
+
+        def get_next_token(self, probs):
+            batch_size, vocab_size = keras.ops.shape(probs)
+            for id in self.forbidden_tokens:
+                update = keras.ops.zeros((batch_size, 1))
+                probs = keras.ops.slice_update(probs, (0, id), update)
+            return keras.ops.argmax(probs, axis=-1)
+
+    # 257 = "a" with a leading space, 262 = "the" with a leading space.
+    causal_lm.compile(sampler=CustomSampler(forbidden_tokens=[257, 262]))
+    causal_lm.summary()
+    causal_lm.generate(["That's strange"])
+    ```
+    """
+
+    def __init__(
+        self,
+        temperature=1.0,
+    ):
+        self.temperature = temperature
+        self._seed_generators = []
+
+    def __setattr__(self, name, value):
+        # We could update to the `Tracker` class from keras-core if our needs
+        # become more advanced (e.g. list assignment, nested trackables). For
+        # now, we only track `SeedGenerator` instances directly on the sampler.
+        if isinstance(value, random.SeedGenerator):
+            self._seed_generators.append(value)
+        return super().__setattr__(name, value)
+
+    @property
+    def variables(self):
+        variables = []
+        for sg in self._seed_generators:
+            variables.append(sg.state)
+        return variables
+
+    def __call__(
+        self,
+        next,
+        prompt,
+        cache=None,
+        index=0,
+        mask=None,
+        stop_token_ids=None,
+        hidden_states=None,
+        model=None,
+    ):
+        max_length = ops.shape(prompt)[-1]
+        # Make sure `max_length` and `index` are the same dtype.
+        index = ops.cast(index, "int32")
+        max_length = ops.cast(max_length, "int32")
+        if mask is None:
+            mask = ops.zeros_like(prompt, dtype="bool")
+        else:
+            mask = ops.cast(mask, dtype="bool")
+        # `ops.while_loop` will not accept `None` as a value for `loop_vars`.
+        cache = () if cache is None else cache
+
+        def cond(prompt, cache, index):
+            if stop_token_ids is None:
+                return True
+            # Stop if all sequences have produced a *new* id from stop_token_ids.
+            end_tokens = any_equal(prompt, stop_token_ids, ~mask)
+            prompt_done = ops.any(end_tokens, axis=-1)
+            return ops.logical_not(ops.all(prompt_done))
+
+        def body(prompt, cache, index):
+            # Compute the softmax distribution for the next token.
+            logits, _, cache = next(prompt, cache, index)
+            probabilities = self.compute_probabilities(logits)
+            # Compute the next token.
+            next_token = self.get_next_token(probabilities)
+            # Don't overwrite anywhere mask is True.
+            next_token = ops.cast(next_token, prompt.dtype)
+            next_token = ops.where(mask[:, index], prompt[:, index], next_token)
+            # Update the prompt with the next token.
+            next_token = next_token[:, None]
+            prompt = ops.slice_update(prompt, [0, index], next_token)
+
+            # Return the next prompt, cache and incremented index.
+            return (prompt, cache, index + 1)
+
+        prompt, _, _ = self.run_loop(
+            cond,
+            body,
+            loop_vars=(prompt, cache, index),
+            maximum_iterations=(max_length - index),
+            model=model,
+        )
+        return prompt
+
+    def compute_probabilities(self, logits):
+        """Compute token probabilities from logits.
+
+        This will always be done in full precision, regardless of dtype, and
+        scale by `temperature`.
+        """
+        logits = ops.cast(logits, "float32")
+        return keras.activations.softmax(logits / self.temperature)
+
+    def run_loop(
+        self, cond, body, model=None, loop_vars=None, maximum_iterations=None
+    ):
+        """Run ops.while_loops with a `StatelessScope` if necessary."""
+        if keras.config.backend() == "jax":
+            import itertools
+
+            if model:
+                model_trainable_variables = model.trainable_variables
+                model_non_trainable_variables = model.non_trainable_variables
+            else:
+                model_trainable_variables = []
+                model_non_trainable_variables = []
+
+            def stateless_cond(state, *loop_vars):
+                return cond(*loop_vars)
+
+            def stateless_body(state, *loop_vars):
+                (
+                    sampler_variables,
+                    trainable_variables,
+                    non_trainable_variables,
+                ) = state
+                mapping = itertools.chain(
+                    zip(self.variables, sampler_variables),
+                    zip(model_trainable_variables, trainable_variables),
+                    zip(model_non_trainable_variables, non_trainable_variables),
+                )
+                with keras.StatelessScope(state_mapping=mapping) as scope:
+                    loop_vars = body(*loop_vars)
+
+                sampler_variables = []
+                for v in self.variables:
+                    new_v = scope.get_current_value(v)
+                    sampler_variables.append(new_v if new_v is not None else v)
+                state = (
+                    sampler_variables,
+                    trainable_variables,
+                    non_trainable_variables,
+                )
+                return state, *loop_vars
+
+            variables = [ops.convert_to_tensor(v) for v in self.variables]
+            trainable_variables = [
+                ops.convert_to_tensor(v) for v in model_trainable_variables
+            ]
+            non_trainable_variables = [
+                ops.convert_to_tensor(v) for v in model_non_trainable_variables
+            ]
+            state = (
+                variables,
+                trainable_variables,
+                non_trainable_variables,
+            )
+            state, *loop_vars = ops.while_loop(
+                cond=stateless_cond,
+                body=stateless_body,
+                loop_vars=(state, *loop_vars),
+                maximum_iterations=maximum_iterations,
+            )
+            for ref_v, v in zip(self.variables, state[0]):
+                ref_v.assign(v)
+        else:
+            loop_vars = ops.while_loop(
+                cond=cond,
+                body=body,
+                loop_vars=(loop_vars),
+                maximum_iterations=maximum_iterations,
+            )
+        return loop_vars
+
+    def get_next_token(self, probabilities):
+        """Get the next token.
+        Args:
+            probabilities: a Tensor, the probability distribution for next
+                token over all vocab tokens.
+        Get the next token based on given probability distribution over tokens.
+        Subclasses must implement this method.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    def get_config(self):
+        return {"temperature": self.temperature}

keras_hub/src/samplers/serialization.py

@@ -0,0 +1,97 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.samplers.beam_sampler import BeamSampler
+from keras_hub.src.samplers.contrastive_sampler import ContrastiveSampler
+from keras_hub.src.samplers.greedy_sampler import GreedySampler
+from keras_hub.src.samplers.random_sampler import RandomSampler
+from keras_hub.src.samplers.top_k_sampler import TopKSampler
+from keras_hub.src.samplers.top_p_sampler import TopPSampler
+
+
+@keras_hub_export("keras_hub.samplers.serialize")
+def serialize(sampler):
+    return keras.saving.serialize_keras_object(sampler)
+
+
+@keras_hub_export("keras_hub.samplers.deserialize")
+def deserialize(config, custom_objects=None):
+    """Return a `Sampler` object from its config."""
+    all_classes = {
+        "beam": BeamSampler,
+        "contrastive": ContrastiveSampler,
+        "greedy": GreedySampler,
+        "random": RandomSampler,
+        "top_k": TopKSampler,
+        "top_p": TopPSampler,
+    }
+    return keras.saving.deserialize_keras_object(
+        config,
+        module_objects=all_classes,
+        custom_objects=custom_objects,
+        printable_module_name="samplers",
+    )
+
+
+@keras_hub_export("keras_hub.samplers.get")
+def get(identifier):
+    """Retrieve a KerasHub sampler by the identifier.
+
+    The `identifier` may be the string name of a sampler class or class.
+
+    >>> identifier = 'greedy'
+    >>> sampler = keras_hub.samplers.get(identifier)
+
+    You can also specify `config` of the sampler to this function by passing
+    dict containing `class_name` and `config` as an identifier. Also note that
+    the `class_name` must map to a `Sampler` class.
+
+    >>> cfg = {'class_name': 'keras_hub>GreedySampler', 'config': {}}
+    >>> sampler = keras_hub.samplers.get(cfg)
+
+    In the case that the `identifier` is a class, this method will return a new
+    instance of the class by its constructor.
+
+    Args:
+        identifier: String or dict that contains the sampler name or
+            configurations.
+
+    Returns:
+        Sampler instance base on the input identifier.
+
+    Raises:
+        ValueError: If the input identifier is not a supported type or in a bad
+            format.
+    """
+
+    if identifier is None:
+        return None
+    if isinstance(identifier, dict):
+        return deserialize(identifier)
+    elif isinstance(identifier, str):
+        if not identifier.islower():
+            raise KeyError(
+                "`keras_hub.samplers.get()` must take a lowercase string "
+                f"identifier, but received: {identifier}."
+            )
+        return deserialize(identifier)
+    elif callable(identifier):
+        return identifier
+    else:
+        raise ValueError(
+            "Could not interpret sampler identifier: " + str(identifier)
+        )

keras_hub/src/samplers/top_k_sampler.py

@@ -0,0 +1,92 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+from keras import ops
+from keras import random
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.samplers.sampler import Sampler
+
+
+@keras_hub_export("keras_hub.samplers.TopKSampler")
+class TopKSampler(Sampler):
+    """Top-K Sampler class.
+
+    This sampler implements top-k search algorithm. Briefly, top-k algorithm
+    randomly selects a token from the tokens of top K probability, with
+    selection chance determined by the probability.
+
+    Args:
+        k: int, the `k` value of top-k.
+        seed: int. The random seed. Defaults to `None`.
+
+    Call arguments:
+        {{call_args}}
+
+    Examples:
+    ```python
+    causal_lm = keras_hub.models.GPT2CausalLM.from_preset("gpt2_base_en")
+
+    # Pass by name to compile.
+    causal_lm.compile(sampler="top_k")
+    causal_lm.generate(["Keras is a"])
+
+    # Pass by object to compile.
+    sampler = keras_hub.samplers.TopKSampler(k=5, temperature=0.7)
+    causal_lm.compile(sampler=sampler)
+    causal_lm.generate(["Keras is a"])
+    ```
+    """
+
+    def __init__(
+        self,
+        k=5,
+        seed=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.k = k
+        self.seed = seed
+        self.seed_generator = random.SeedGenerator(seed)
+
+    def get_next_token(self, probabilities):
+        # Filter out top-k tokens.
+        top_k_pred, top_k_indices = ops.top_k(
+            probabilities,
+            k=self.k,
+            sorted=False,
+        )
+        # Sample the next token from the probability distribution.
+        sample_indices = random.categorical(
+            # tf does not support half precision multinomial sampling, so make
+            # sure we have full precision here.
+            ops.cast(ops.log(top_k_pred), "float32"),
+            1,
+            seed=self.seed_generator,
+            dtype="int32",
+        )
+
+        # Rearrange to get the next token idx from the original order.
+        output = ops.take_along_axis(top_k_indices, sample_indices, axis=-1)
+        return ops.squeeze(output, axis=-1)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "k": self.k,
+                "seed": self.seed,
+            }
+        )
+        return config

keras_hub/src/samplers/top_p_sampler.py

@@ -0,0 +1,113 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+from keras import ops
+from keras import random
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.samplers.sampler import Sampler
+
+
+@keras_hub_export("keras_hub.samplers.TopPSampler")
+class TopPSampler(Sampler):
+    """Top-P Sampler class.
+
+    This sampler implements top-p search algorithm. Top-p search selects tokens
+    from the smallest subset of output probabilities that sum to greater than
+    `p`. Put in another way, top-p will first order token predictions by
+    likelihood, and ignore all tokens after the cumulative probability of
+    selected tokens exceeds `p`, then select a token from the remaining tokens.
+
+    Args:
+        p: float, the `p` value of top-p.
+        k: int. If set, this argument defines a
+            heuristic "top-k" cutoff applied before the "top-p" sampling. All
+            logits not in the top `k` will be discarded, and the remaining
+            logits will be sorted to find a cutoff point for `p`. Setting this
+            arg can significantly speed sampling up by reducing the number
+            of tokens to sort. Defaults to `None`.
+        seed: int. The random seed. Defaults to `None`.
+
+    Call arguments:
+        {{call_args}}
+
+    Examples:
+    ```python
+    causal_lm = keras_hub.models.GPT2CausalLM.from_preset("gpt2_base_en")
+
+    # Pass by name to compile.
+    causal_lm.compile(sampler="top_p")
+    causal_lm.generate(["Keras is a"])
+
+    # Pass by object to compile.
+    sampler = keras_hub.samplers.TopPSampler(p=0.1, k=1_000)
+    causal_lm.compile(sampler=sampler)
+    causal_lm.generate(["Keras is a"])
+    ```
+    """
+
+    def __init__(
+        self,
+        p=0.1,
+        k=None,
+        seed=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.p = p
+        self.k = k
+        self.seed = seed
+        self.seed_generator = random.SeedGenerator(seed)
+
+    def get_next_token(self, probabilities):
+        cutoff = ops.shape(probabilities)[1]
+        if self.k is not None:
+            # If `k` is set, only sample from top `k` tokens.
+            cutoff = self.k
+        sorted_preds, sorted_indices = ops.top_k(
+            probabilities, k=cutoff, sorted=True
+        )
+        # Calculate cumulative probability distribution.
+        cumulative_probabilities = ops.cumsum(sorted_preds, axis=-1)
+        # Create a mask for the tokens to keep.
+        keep_mask = cumulative_probabilities <= self.p
+        # Shift to include the last token that exceed p.
+        shifted_keep_mask = ops.concatenate(
+            [ops.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1
+        )
+        # Filter out unmasked tokens and sample from filtered distribution.
+        probabilities = ops.where(
+            shifted_keep_mask,
+            sorted_preds,
+            ops.zeros(ops.shape(sorted_preds), dtype=sorted_preds.dtype),
+        )
+        sorted_next_token = random.categorical(
+            ops.log(probabilities),
+            1,
+            seed=self.seed_generator,
+            dtype="int32",
+        )
+        output = ops.take_along_axis(sorted_indices, sorted_next_token, axis=-1)
+        return ops.squeeze(output, axis=-1)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "p": self.p,
+                "k": self.k,
+                "seed": self.seed,
+            }
+        )
+        return config

keras_hub/src/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.