keras-hub 0.25.0.dev0__py3-none-any.whl → 0.26.0.dev0__py3-none-any.whl

keras_hub/src/models/rwkv7/rwkv7_backbone.py

@@ -0,0 +1,180 @@
+import keras
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.rwkv7.rwkv7_layer import RWKV7_Block
+
+
+def rwkv7_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras_hub_export("keras_hub.models.RWKV7Backbone")
+class RWKV7Backbone(Backbone):
+    """The RWKV7 Transformer core architecture with hyperparameters.
+
+    This network implements a RNN-based decoder network,
+    Goose, as described in
+    [RWKV-7](https://arxiv.org/abs/2503.14456).
+
+    This network implements a Modern RNN architecture based on linear
+    attention mechanisms with recurrent processing, as described in the
+    RWKV papers. It includes the embedding lookups and RWKV-7 blocks.
+
+    The default constructor gives a fully customizable, randomly initialized
+    RWKV-7 model with any number of layers, heads, and embedding dimensions.
+    To load preset architectures and weights, use the `from_preset`
+    constructor.
+
+    Args:
+        hidden_size: int. The size of the transformer encoding and pooling
+            layers.
+        head_size: int. The size of each attention head.
+        num_layers: int. The number of transformer layers.
+        vocabulary_size: int. The size of the token vocabulary.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        gate_lora: int. LoRA dimension for gating. Defaults to `128` .
+        mv_lora: int. LoRA dimension for value mixing. Defaults to `32` .
+        aaa_lora: int. LoRA dimension for alpha parameters.Defaults to `64` .
+        decay_lora: int. LoRA dimension for decay parameters.Defaults to `64` .
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
+        dropout_rate: float. Dropout rate for the dropout layer.
+
+    Examples:
+
+    ```python
+    input_data = np.ones(shape=(1, 12), dtype="int32")
+
+
+    # Randomly initialized RWKV-7 decoder with custom config.
+    model = keras_hub.models.RWKV7Backbone(
+        vocabulary_size=10,
+        hidden_size=512,
+        num_layers=2,
+        head_size=64,
+        intermediate_dim=1024,
+        dtype="float32"
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        head_size,
+        num_layers,
+        vocabulary_size,
+        intermediate_dim,
+        gate_lora=128,
+        mv_lora=32,
+        aaa_lora=64,
+        decay_lora=64,
+        dtype=None,
+        dropout_rate=0,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.token_embedding = keras.layers.Embedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_size,
+            embeddings_initializer=rwkv7_kernel_initializer(),
+            dtype=dtype,
+            name="token_embedding",
+        )
+
+        self.output_layer_norm = keras.layers.LayerNormalization(
+            epsilon=1e-5,
+            name="output_norm",
+            dtype=dtype,
+        )
+        self.dropout = keras.layers.Dropout(
+            dropout_rate,
+            dtype=dtype,
+            name="dropout",
+        )
+        self.rwkv_layers = []
+        for i in range(num_layers):
+            layer = RWKV7_Block(
+                hidden_size,
+                head_size,
+                intermediate_dim,
+                gate_lora,
+                mv_lora,
+                aaa_lora,
+                decay_lora,
+                use_initial_norm=i == 0,
+                kernel_initializer=rwkv7_kernel_initializer(),
+                dtype=dtype,
+                name=f"rwkv_layer_{i}",
+            )
+
+            self.rwkv_layers.append(layer)
+        self.head = keras.layers.Dense(
+            units=vocabulary_size,
+            kernel_initializer=rwkv7_kernel_initializer(),
+            use_bias=False,
+            name="head",
+            dtype=dtype,
+        )
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+
+        x = self.token_embedding(token_id_input)
+        padding_mask = ops.cast(padding_mask_input, dtype=x.dtype)
+        v_first = None
+        for rwkv_layer in self.rwkv_layers:
+            x, v_first = rwkv_layer(x, v_first, padding_mask)
+            x = self.dropout(x)
+        sequence_output = self.output_layer_norm(x)
+        sequence_output = self.head(sequence_output)
+
+        super().__init__(
+            inputs={
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
+            },
+            outputs=sequence_output,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        self.num_layers = num_layers
+        self.head_size = head_size
+        self.hidden_size = hidden_size
+        self.gate_lora = gate_lora
+        self.mv_lora = mv_lora
+        self.aaa_lora = aaa_lora
+        self.decay_lora = decay_lora
+        self.vocabulary_size = vocabulary_size
+        self.dropout_rate = dropout_rate
+        self.intermediate_dim = intermediate_dim
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_size": self.hidden_size,
+                "head_size": self.head_size,
+                "gate_lora": self.gate_lora,
+                "mv_lora": self.mv_lora,
+                "aaa_lora": self.aaa_lora,
+                "decay_lora": self.decay_lora,
+                "vocabulary_size": self.vocabulary_size,
+                "dropout_rate": self.dropout_rate,
+                "intermediate_dim": self.intermediate_dim,
+                "num_layers": self.num_layers,
+            }
+        )
+        return config

keras_hub/src/models/rwkv7/rwkv7_causal_lm.py

@@ -0,0 +1,259 @@
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.causal_lm import CausalLM
+from keras_hub.src.models.rwkv7.rwkv7_backbone import RWKV7Backbone
+from keras_hub.src.models.rwkv7.rwkv7_causal_lm_preprocessor import (
+    RWKV7CausalLMPreprocessor,
+)
+from keras_hub.src.utils.tensor_utils import any_equal
+
+
+@keras_hub_export("keras_hub.models.RWKV7CausalLM")
+class RWKV7CausalLM(CausalLM):
+    """An end-to-end RWKV-7 model for causal language modeling.
+
+    A causal language model (LM) predicts the next token based on previous
+    tokens. This task setup can be used to train the model unsupervised on
+    plain text input, or to autoregressively generate plain text similar to
+    the data used for training. This task can be used for pre-training or
+    fine-tuning a RWKV-7 model, simply by calling `fit()`.
+
+    This model has a generate() method, which generates text based on a
+    prompt. The generation strategy used is controlled by an additional
+    sampler argument on `compile()`. You can recompile the model with
+    different `keras_hub.samplers` objects to control the generation. By
+    default, `"greedy"` sampling will be used.
+
+    Args:
+        backbone: A `keras_hub.models.RWKV7Backbone` instance.
+        preprocessor: A `keras_hub.models.RWKV7CausalLMPreprocessor` or `None`.
+            If `None`, this model will not apply preprocessing, and inputs
+            should be preprocessed before calling the model.
+
+    Examples:
+    ```python
+    # Initialize the tokenizer and load assets from a local path.
+    tokenizer = RWKVTokenizer()
+    tokenizer.load_assets(rwkv_path)
+
+    # Create a preprocessor with a sequence length of 8.
+    preprocessor = RWKV7CausalLMPreprocessor(tokenizer, sequence_length=8)
+
+    # Initialize the model with a backbone and preprocessor.
+    causal_lm = RWKV7CausalLM(backbone, preprocessor)
+
+    # you also can load model by from_preset
+    rwkv_path = "RWKV7_G1a_0.1B"
+    tokenizer = RWKVTokenizer.from_preset(rwkv_path)
+    causal_lm = RWKV7CausalLM.from_preset(rwkv_path)
+
+    prompts = ["Bubble sort\n```python", "Hello World\n```python\n"]
+
+    causal_lm.compile(sampler="greedy")
+
+    outputs = causal_lm.generate(prompts, max_length=128)
+    for out in outputs:
+        print(out)
+        print("-" * 100)
+    ```
+    """
+
+    backbone_cls = RWKV7Backbone
+    preprocessor_cls = RWKV7CausalLMPreprocessor
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        """Initialize the RWKV-7 causal language model.
+
+        Args:
+            backbone: The backbone model.
+            preprocessor: The preprocessor for tokenization.
+            **kwargs: Additional keyword arguments.
+        """
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        super().__init__(
+            inputs=backbone.input,
+            outputs=backbone.output,
+            **kwargs,
+        )
+
+    def call_with_cache(
+        self,
+        token_ids,
+        cache,
+        compute_head=True,
+        padding_mask=None,
+        rnn_mode=True,
+    ):
+        """Forward pass of `RWKV7CausalLM` with cache.
+
+        `call_with_cache` adds an additional forward pass for the model for
+        autoregressive inference. Unlike calling the model directly, this method
+        allows caching previous state Tensors in RWKV layers, and avoids
+        recomputing the outputs of seen tokens.
+
+        Args:
+            token_ids: a dense int Tensor with shape `(batch_size, max_length)`.
+            cache: a dense float Tensor, the cache of state and token values.
+            compute_head: bool, whether to compute the output head.
+            padding_mask: a dense bool Tensor, the padding mask.
+            rnn_mode: bool, whether to use RNN mode.
+
+        Returns:
+            A (logits, hidden_states, cache) tuple. Where `logits` is the
+            language model logits for the input token_ids, `hidden_states` is
+            the final hidden representation of the input tokens, and `cache` is
+            the decoding cache.
+        """
+        state_cache, last_token_cache = cache
+        x = self.backbone.token_embedding(token_ids)
+        if padding_mask is None:
+            padding_mask = ops.not_equal(token_ids, 0)
+        padding_mask = ops.cast(padding_mask, x.dtype)
+        v_first = None
+        updated_state_cache = []
+        updated_last_token_cache = []
+
+        for i in range(self.backbone.num_layers):
+            current_state_cache = state_cache[:, i, ...]
+            current_token_cache = last_token_cache[:, i, ...]
+            x, v_first, new_cache_state, cache_tmix_x, cache_cmix_x = (
+                self.backbone.rwkv_layers[i].generate_call(
+                    x,
+                    v_first=v_first,
+                    padding_mask=padding_mask,
+                    cache_state=current_state_cache,
+                    cache_tmix_x=current_token_cache[:, 0],
+                    cache_cmix_x=current_token_cache[:, 1],
+                    rnn_mode=rnn_mode,
+                )
+            )
+            new_token_cache = ops.stack([cache_tmix_x, cache_cmix_x], axis=1)
+            updated_state_cache.append(new_cache_state)
+            updated_last_token_cache.append(new_token_cache)
+        cache = [
+            ops.stack(updated_state_cache, axis=1),
+            ops.stack(updated_last_token_cache, axis=1),
+        ]
+        hidden_states = x = self.backbone.output_layer_norm(x)
+        if compute_head:
+            logits = self.backbone.head(x)
+        else:
+            logits = None
+        return logits, hidden_states, cache
+
+    def _build_cache(self, token_ids, padding_mask):
+        """Build an empty cache for use with `call_with_cache()`."""
+        batch_size = ops.shape(token_ids)[0]
+        num_layers = self.backbone.num_layers
+        head_dim = self.backbone.head_size
+        hidden_size = self.backbone.hidden_size
+        num_heads = hidden_size // head_dim
+
+        state_cache = ops.zeros(
+            [batch_size, num_layers, num_heads, head_dim, head_dim],
+            dtype=self.compute_dtype,
+        )
+        last_token_cache = ops.zeros(
+            [batch_size, num_layers, 2, 1, hidden_size],
+            dtype=self.compute_dtype,
+        )
+        cache = [state_cache, last_token_cache]
+
+        # Seed the cache.
+        # Prefill stage can use kernel for better performance
+        _, hidden_states, cache = self.call_with_cache(
+            token_ids,
+            cache,
+            rnn_mode=False,
+            compute_head=False,
+            padding_mask=padding_mask,
+        )
+
+        return hidden_states, cache
+
+    def generate_step(
+        self,
+        inputs,
+        stop_token_ids=None,
+    ):
+        """A compilable generation function for a single batch of inputs.
+
+        This function represents the inner, XLA-compilable, generation function
+        for a single batch of inputs. Inputs should have the same structure as
+        model inputs, a dictionary with keys `"token_ids"` and `"padding_mask"`.
+
+        Args:
+            inputs: A dictionary with keys `"token_ids"`, `"padding_mask"`, and
+                `"predict_token_ids"` with batched tensor values.
+            stop_token_ids: Tuple of id's of the end token to stop on. If all
+                sequences have produced a new stop token, generation
+                will stop.
+        """
+        token_ids, padding_mask, predict_token_ids = (
+            inputs["token_ids"],
+            inputs["padding_mask"],
+            inputs["predict_token_ids"],
+        )
+        # Create and seed cache with a single forward pass.
+
+        hidden_states, cache = self._build_cache(
+            token_ids, inputs["input_padding_mask"]
+        )
+
+        def next(prompt, cache, index):
+            # The cache index is the index of our previous token.
+            cache_update_index = index - 1
+            batch_size = ops.shape(prompt)[0]
+            prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1])
+            logits, hidden_states, cache = self.call_with_cache(
+                prompt,
+                cache,
+            )
+            return (
+                ops.squeeze(logits, axis=1),
+                ops.squeeze(hidden_states, axis=1),
+                cache,
+            )
+
+        output_ids = self.sampler(
+            next=next,
+            prompt=predict_token_ids,
+            cache=cache,
+            index=1,
+            mask=padding_mask,
+            stop_token_ids=stop_token_ids,
+            hidden_states=hidden_states,
+            model=self,
+        )
+        padding_mask = ops.concatenate(
+            [
+                ops.cast(ops.not_equal(token_ids, 0), padding_mask.dtype),
+                padding_mask,
+            ],
+            axis=1,
+        )
+        token_ids = ops.concatenate([token_ids, output_ids], axis=1)
+
+        # Compute an output padding mask with the token ids we updated.
+        if stop_token_ids is not None:
+            # Build a mask of stop token locations not in the original
+            # prompt (not in locations where `padding_mask` is True).
+            end_locations = any_equal(
+                token_ids, stop_token_ids, ops.logical_not(padding_mask)
+            )
+            end_locations = ops.cast(end_locations, "int32")
+            # Use cumsum to get ones in all locations after end_locations.
+            cumsum = ops.cast(ops.cumsum(end_locations, axis=-1), "int32")
+            overflow = cumsum - end_locations
+            # Our padding mask is the inverse of these overflow locations.
+            padding_mask = ops.logical_not(ops.cast(overflow, "bool"))
+        else:
+            # Without early stopping, all locations will have been updated.
+            padding_mask = ops.ones_like(token_ids, dtype="bool")
+        return {
+            "token_ids": token_ids,
+            "padding_mask": ops.cast(padding_mask, token_ids.dtype),
+        }

keras_hub/src/models/rwkv7/rwkv7_causal_lm_preprocessor.py

@@ -0,0 +1,214 @@
+import keras
+import numpy as np
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
+from keras_hub.src.models.rwkv7.rwkv7_backbone import RWKV7Backbone
+from keras_hub.src.models.rwkv7.rwkv7_tokenizer import RWKVTokenizer
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+from keras_hub.src.utils.tensor_utils import tf
+
+
+@keras_hub_export("keras_hub.models.RWKV7CausalLMPreprocessor")
+class RWKV7CausalLMPreprocessor(CausalLMPreprocessor):
+    """RWKV-7 Causal LM preprocessor.
+
+    This preprocessing layer is meant for use with
+    `keras_hub.models.RWKV7CausalLM`. By default, it will take in batches of
+    strings, and return outputs in a `(x, y, sample_weight)` format, where the
+    `y` label is the next token id in the `x` sequence.
+
+    For use with generation, the layer also exposes two methods
+    `generate_preprocess()` and `generate_postprocess()`. When this preprocessor
+    is attached to a `keras_hub.models.RWKV7CausalLM` instance, these methods
+    will be called implicitly in generate(). They can also be called
+    standalone (e.g. to precompute preprocessing inputs for generation in a
+    separate process).
+
+    Args:
+        tokenizer: A `keras_hub.models.RWKVTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence. Default is `False`.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Label data. Should always be `None` as the layer generates labels.
+        sample_weight: Label weights. Should always be `None` as the layer
+            generates label weights.
+        sequence_length: Pass to override the configured sequence_length of
+            the layer.
+
+
+    Examples:
+    ```python
+    # Initialize the tokenizer and load assets from a local path.
+    tokenizer = RWKVTokenizer()
+    tokenizer.load_assets(rwkv_path)
+
+    # Create a preprocessor with a sequence length of 8.
+    preprocessor = RWKV7CausalLMPreprocessor(tokenizer, sequence_length=8)
+
+    # Tokenize and pack a batch of sentences.
+    preprocessor(["Bubble sort\n```python", "Hello World\n```python\n"])
+
+    # Preprocess inputs for generation with a maximum generation length of 16.
+    preprocessor.generate_preprocess(
+        ["Bubble sort\n```python", "Hello World\n```python\n"], 16
+    )
+    ```
+    """
+
+    backbone_cls = RWKV7Backbone
+    tokenizer_cls = RWKVTokenizer
+
+    def __init__(
+        self,
+        tokenizer,
+        add_start_token=False,
+        **kwargs,
+    ):
+        """Initialize the preprocessor.
+
+        Args:
+            tokenizer: The tokenizer to use.
+            add_start_token: Whether to add start token.
+            **kwargs: Additional arguments.
+        """
+        super().__init__(
+            tokenizer=tokenizer, add_start_token=add_start_token, **kwargs
+        )
+
+    @preprocessing_function
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        """Preprocess the input for training.
+
+        Args:
+            x: Input text data.
+            y: Target data (optional).
+            sample_weight: Sample weights (optional).
+            sequence_length: Desired sequence length.
+
+        Returns:
+            Preprocessed data tuple (x, y, sample_weight).
+        """
+        if isinstance(x, str):
+            x = [x]
+        sequence_length = sequence_length or self.sequence_length
+        # Pad length to multiples of 16 to meet kernel requirements
+        if sequence_length is None:
+            raise ValueError("sequence_length must be specified.")
+        if keras.config.backend() in ["torch", "jax"]:
+            # When using rwkv_ops, ensure sequence_length is divisible by 16.
+            try:
+                import rwkv_ops  # noqa: F401
+
+                if sequence_length % 16 != 0:
+                    sequence_length += (16 - sequence_length % 16) % 16
+            except ImportError:
+                pass
+        x = self.tokenizer(x)
+
+        token_ids, padding_mask = self.packer(
+            x, sequence_length=sequence_length + 1, add_end_value=False
+        )
+
+        # The last token does not have a next token, so we truncate it out.
+        x = {
+            "token_ids": token_ids[..., :-1],
+            "padding_mask": padding_mask[..., :-1],
+        }
+        # Target `y` will be the next token.
+        y, sample_weight = token_ids[..., 1:], padding_mask[..., 1:]
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    def build(self, input_shape):
+        self.packer = StartEndPacker(
+            start_value=None,
+            end_value=None,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+            padding_side="left",  # RWKV uses left-padding exclusively
+        )
+        self.built = True
+
+    @preprocessing_function
+    def generate_preprocess(
+        self,
+        x,
+        sequence_length=None,
+    ):
+        """Preprocess input for generation.
+
+        Args:
+            x: Input text data.
+            sequence_length: Maximum generation length.
+
+        Returns:
+            Dictionary with preprocessed inputs for generation.
+        """
+        if isinstance(x, str):
+            x = [x]
+
+        if not self.built:
+            self.build(None)
+        # Align with Keras API
+        # Input sequence_length is the maximum generation length
+        # While self.sequence_length corresponds to the prefill max length
+        generate_length = sequence_length
+        if sequence_length is None:
+            raise ValueError("`sequence_length` must be specified.")
+        sequence_length = self.sequence_length
+
+        x = [t[-sequence_length:] for t in self.tokenizer(x)]
+        y = tf.zeros((len(x), generate_length), "int32")
+        # Utilize RNN characteristics where prefill and decode are two sequences
+        # But the first token of decode should be the last token of prefill
+        start_token = [[t[-1]] for t in x]
+        x = [np.array(t[:-1]) if len(t) > 1 else [0] for t in x]
+        x = tf.ragged.constant(x)
+        token_ids, input_padding_mask = self.packer(
+            x, sequence_length=sequence_length, add_end_value=False
+        )
+        start_token = tf.convert_to_tensor(start_token, "int32")
+
+        y = tf.concat([start_token, y], axis=1)
+        padding_mask = tf.not_equal(y, 0)
+
+        return {
+            "token_ids": token_ids,
+            "input_padding_mask": input_padding_mask,
+            "padding_mask": padding_mask,
+            "predict_token_ids": y,
+        }
+
+    @preprocessing_function
+    def generate_postprocess(
+        self,
+        x,
+    ):
+        """Convert integer token output to strings for generation.
+
+        This method reverses `generate_preprocess()`, by first removing all
+        padding and start/end tokens, and then converting the integer sequence
+        back to a string.
+
+        Args:
+            x: Dictionary containing token_ids and padding_mask.
+
+        Returns:
+            Detokenized string output.
+        """
+        if not self.built:
+            self.build(None)
+
+        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
+        return self.tokenizer.detokenize(token_ids * padding_mask)