keras-hub-nightly 0.24.0.dev202511220420__py3-none-any.whl → 0.26.0.dev202601010440__py3-none-any.whl

keras_hub/src/models/gpt_oss/gpt_oss_decoder.py

@@ -0,0 +1,444 @@
+import keras
+from keras import ops
+
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_hub.src.models.gpt_oss.gpt_oss_attention import GptOssAttention
+from keras_hub.src.models.gpt_oss.gpt_oss_layer_norm import (
+    GptOssLayerNormalization,
+)
+from keras_hub.src.utils.keras_utils import clone_initializer
+
+
+class GptOssExperts(keras.layers.Layer):
+    """A layer containing the feed-forward expert networks for GPT-OSS.
+
+    This layer implements the expert networks as described in the GPT-OSS
+    paper. It uses a custom GLU activation.
+
+    Args:
+        num_experts: int. The total number of experts.
+        hidden_dim: int. The hidden size of the model.
+        intermediate_dim: int. The intermediate size of the feed-forward
+            network.
+        kernel_initializer: string. The initializer for the kernel
+            weights. Defaults to "glorot_uniform".
+        alpha: float. The alpha parameter for the custom GLU
+            activation. Defaults to `1.702`.
+        limit: float. The clamping limit for gate and up
+            projections. Defaults to `7.0`.
+    """
+
+    def __init__(
+        self,
+        num_experts,
+        hidden_dim,
+        intermediate_dim,
+        kernel_initializer="glorot_uniform",
+        alpha=1.702,
+        limit=7.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_experts = num_experts
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.alpha = alpha
+        self.limit = limit
+
+    def build(self, _):
+        self.gate_up_proj = self.add_weight(
+            shape=(
+                self.num_experts,
+                self.hidden_dim,
+                2 * self.intermediate_dim,
+            ),
+            initializer=self.kernel_initializer,
+            name="gate_up_proj",
+        )
+        self.gate_up_proj_bias = self.add_weight(
+            shape=(self.num_experts, 2 * self.intermediate_dim),
+            initializer="zeros",
+            name="gate_up_proj_bias",
+        )
+        self.down_proj = self.add_weight(
+            shape=(self.num_experts, self.intermediate_dim, self.hidden_dim),
+            initializer=self.kernel_initializer,
+            name="down_proj",
+        )
+        self.down_proj_bias = self.add_weight(
+            shape=(self.num_experts, self.hidden_dim),
+            initializer="zeros",
+            name="down_proj_bias",
+        )
+        self.built = True
+
+    def call(self, hidden_states):
+        # hidden_states shape: (num_tokens, hidden_dim)
+        # Einsum for batched matrix multiplication across experts.
+        # [num_experts, num_tokens, 2 * intermediate_dim]
+        gate_up = ops.einsum("th,ehm->etm", hidden_states, self.gate_up_proj)
+        gate_up = gate_up + self.gate_up_proj_bias[:, None, :]
+
+        gate = gate_up[..., ::2]
+        up = gate_up[..., 1::2]
+
+        gate = ops.clip(gate, -1e9, self.limit)
+        up = ops.clip(up, -self.limit, self.limit)
+
+        glu = gate * ops.sigmoid(gate * self.alpha)
+        gated_output = (up + 1) * glu
+
+        # [num_experts, num_tokens, hidden_dim]
+        out = ops.einsum("etm,emh->eth", gated_output, self.down_proj)
+        out = out + self.down_proj_bias[:, None, :]
+        return out
+
+
+class GptOssTopKRouter(keras.layers.Layer):
+    """A layer for routing tokens to the top-k experts.
+
+    Args:
+        num_experts: int. The total number of experts.
+        top_k: int. The number of experts to route each token to.
+        kernel_initializer: string. The initializer for the kernel
+            weights. Defaults to "glorot_uniform".
+    """
+
+    def __init__(
+        self,
+        num_experts,
+        top_k,
+        kernel_initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+
+    def build(self, hidden_states_shape):
+        self.router_dense = keras.layers.Dense(
+            self.num_experts,
+            kernel_initializer=self.kernel_initializer,
+            dtype=self.dtype_policy,
+            name="router_dense",
+        )
+        self.router_dense.build(hidden_states_shape)
+        self.built = True
+
+    def call(self, hidden_states):
+        # hidden_states shape: (num_tokens, hidden_dim)
+        router_logits = self.router_dense(hidden_states)
+
+        routing_weights, selected_experts = ops.top_k(
+            router_logits, k=self.top_k
+        )
+        routing_weights = ops.softmax(routing_weights, axis=-1)
+
+        expert_mask = ops.one_hot(selected_experts, self.num_experts)
+        expert_mask = ops.cast(expert_mask, dtype=routing_weights.dtype)
+
+        # Shape: (num_tokens, top_k, num_experts)
+        weighted_mask = expert_mask * ops.expand_dims(routing_weights, axis=-1)
+
+        # Shape: (num_tokens, num_experts)
+        router_scores = ops.sum(weighted_mask, axis=1)
+
+        return router_scores
+
+
+class GptOssSparseMoeBlock(keras.layers.Layer):
+    """GPT-OSS sparse Mixture of Experts (MoE) block.
+
+    This block combines a router and a set of expert networks to implement
+    the MoE layer.
+
+    Args:
+        hidden_dim: int. The hidden size of the model.
+        intermediate_dim: int. The intermediate size of the feed-forward
+            network.
+        num_experts: int. The total number of experts.
+        top_k: int. The number of experts to route each token to.
+            Defaults to 2.
+        kernel_initializer: string. The initializer for the kernel
+            weights. Defaults to "glorot_uniform".
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        intermediate_dim,
+        num_experts,
+        top_k=2,
+        kernel_initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.kernel_initializer = kernel_initializer
+
+    def build(self, decoder_sequence_shape):
+        self.router = GptOssTopKRouter(
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
+            name="router",
+        )
+        self.router.build(decoder_sequence_shape)
+
+        self.experts = GptOssExperts(
+            num_experts=self.num_experts,
+            hidden_dim=self.hidden_dim,
+            intermediate_dim=self.intermediate_dim,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
+            name="experts",
+        )
+        self.experts.build(decoder_sequence_shape)
+        self.built = True
+
+    def call(self, hidden_states):
+        batch_size, seq_len, _ = ops.shape(hidden_states)
+        hidden_states_flattened = ops.reshape(
+            hidden_states, (-1, self.hidden_dim)
+        )
+
+        router_scores = self.router(hidden_states_flattened)
+
+        expert_outputs = self.experts(hidden_states_flattened)
+
+        # Weight expert outputs by router scores and sum
+        # router_scores shape: (num_tokens, num_experts)
+        # expert_outputs shape: (num_experts, num_tokens, hidden_dim)
+        # Transpose scores for broadcasting: (num_experts, num_tokens)
+        router_scores_t = ops.transpose(router_scores)
+        # Expand for broadcasting: (num_experts, num_tokens, 1)
+        router_scores_expanded = ops.expand_dims(router_scores_t, axis=-1)
+
+        weighted_outputs = expert_outputs * router_scores_expanded
+        final_output = ops.sum(weighted_outputs, axis=0)
+
+        final_output = ops.reshape(
+            final_output, (batch_size, seq_len, self.hidden_dim)
+        )
+        return final_output, router_scores
+
+
+class GptOssTransformerDecoder(keras.layers.Layer):
+    """A GPT-OSS transformer decoder layer.
+
+    This layer implements the transformer decoder block from the GPT-OSS
+    model, which includes self-attention and a sparse MoE block.
+
+    Args:
+        intermediate_dim: int. The intermediate size of the feed-forward
+            network.
+        num_query_heads: int. The number of query attention heads.
+        num_key_value_heads: int. The number of key and value attention
+            heads.
+        num_experts: int. The total number of experts in the MoE layer.
+        top_k: int. The number of experts to route each token to.
+            Defaults to 2.
+        output_router_logits: bool. If True, the router logits will
+            be returned by the layer. Defaults to False.
+        rope_max_wavelength: int. The maximum wavelength for the
+            rotary position embedding. Defaults to 10000.
+        rope_scaling_factor: float. The scaling factor for the
+            rotary position embedding. Defaults to 1.0.
+        layer_norm_epsilon: float. The epsilon for layer
+            normalization. Defaults to 1e-6.
+        kernel_initializer: string. The initializer for the kernel
+            weights. Defaults to "glorot_uniform".
+        sliding_window: int. The size of the sliding window for
+            attention. Defaults to 4096.
+        dropout: float. The dropout rate. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        intermediate_dim,
+        num_query_heads,
+        num_key_value_heads,
+        num_experts,
+        top_k=2,
+        output_router_logits=False,
+        rope_max_wavelength=10000,
+        rope_scaling_factor=1.0,
+        layer_norm_epsilon=1e-6,
+        kernel_initializer="glorot_uniform",
+        sliding_window=4096,
+        dropout=0,
+        head_dim=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.intermediate_dim = intermediate_dim
+        self.num_query_heads = num_query_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.output_router_logits = output_router_logits
+        self.rope_max_wavelength = rope_max_wavelength
+        self.rope_scaling_factor = rope_scaling_factor
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.sliding_window = sliding_window
+        self.dropout = dropout
+        self.head_dim = head_dim
+        self.supports_masking = True
+
+    def build(self, decoder_sequence_shape):
+        self.hidden_dim = decoder_sequence_shape[-1]
+
+        self.self_attention_layer = GptOssAttention(
+            num_query_heads=self.num_query_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            rope_max_wavelength=self.rope_max_wavelength,
+            rope_scaling_factor=self.rope_scaling_factor,
+            sliding_window=self.sliding_window,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            dropout=self.dropout,
+            head_dim=self.head_dim,
+            dtype=self.dtype_policy,
+            name="self_attention",
+        )
+        self.self_attention_layer.build(decoder_sequence_shape)
+
+        self.input_layernorm = GptOssLayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="input_layernorm",
+        )
+        self.input_layernorm.build(decoder_sequence_shape)
+
+        self.post_attention_layernorm = GptOssLayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="post_attention_layernorm",
+        )
+        self.post_attention_layernorm.build(decoder_sequence_shape)
+
+        self.sparse_moe_block = GptOssSparseMoeBlock(
+            hidden_dim=self.hidden_dim,
+            intermediate_dim=self.intermediate_dim,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
+            name="sparse_moe_block",
+        )
+        self.sparse_moe_block.build(decoder_sequence_shape)
+
+        self.built = True
+
+    def call(
+        self,
+        decoder_sequence,
+        decoder_padding_mask=None,
+        decoder_attention_mask=None,
+        self_attention_cache=None,
+        self_attention_cache_update_index=None,
+        training=None,
+    ):
+        self_attention_mask = self._compute_self_attention_mask(
+            decoder_sequence=decoder_sequence,
+            decoder_padding_mask=decoder_padding_mask,
+            decoder_attention_mask=decoder_attention_mask,
+            self_attention_cache=self_attention_cache,
+            self_attention_cache_update_index=self_attention_cache_update_index,
+        )
+
+        residual = decoder_sequence
+        x = self.input_layernorm(decoder_sequence)
+
+        x = self.self_attention_layer(
+            hidden_states=x,
+            attention_mask=self_attention_mask,
+            cache=self_attention_cache,
+            cache_update_index=self_attention_cache_update_index,
+        )
+
+        if self_attention_cache is not None:
+            x, self_attention_cache = x
+
+        x = x + residual
+        residual = x
+
+        x = self.post_attention_layernorm(x)
+        x, router_logits = self.sparse_moe_block(x)
+
+        decoder_output = x + residual
+
+        output = (decoder_output,)
+        if self_attention_cache is not None:
+            output += (self_attention_cache,)
+        if self.output_router_logits:
+            output += (router_logits,)
+
+        return output[0] if len(output) == 1 else output
+
+    def _compute_self_attention_mask(
+        self,
+        decoder_sequence,
+        decoder_padding_mask,
+        decoder_attention_mask,
+        self_attention_cache,
+        self_attention_cache_update_index,
+    ):
+        decoder_mask = merge_padding_and_attention_mask(
+            decoder_sequence, decoder_padding_mask, decoder_attention_mask
+        )
+        batch_size = ops.shape(decoder_sequence)[0]
+        input_length = output_length = ops.shape(decoder_sequence)[1]
+
+        if self_attention_cache is not None:
+            input_length = ops.shape(self_attention_cache)[2]
+
+        cache_update_index = (
+            0
+            if self_attention_cache_update_index is None
+            else self_attention_cache_update_index
+        )
+
+        causal_mask = compute_causal_mask(
+            batch_size, input_length, output_length, cache_update_index
+        )
+
+        return (
+            ops.minimum(decoder_mask, causal_mask)
+            if decoder_mask is not None
+            else causal_mask
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "intermediate_dim": self.intermediate_dim,
+                "num_query_heads": self.num_query_heads,
+                "num_key_value_heads": self.num_key_value_heads,
+                "num_experts": self.num_experts,
+                "top_k": self.top_k,
+                "output_router_logits": self.output_router_logits,
+                "rope_max_wavelength": self.rope_max_wavelength,
+                "rope_scaling_factor": self.rope_scaling_factor,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "sliding_window": self.sliding_window,
+                "dropout": self.dropout,
+                "head_dim": self.head_dim,
+            }
+        )
+        return config

keras_hub/src/models/gpt_oss/gpt_oss_layer_norm.py

@@ -0,0 +1,34 @@
+import keras
+from keras import ops
+
+
+# NOTE: `keras.layers.LayerNormalization(rms_scaling=True)`
+# does not produce the same results.
+class GptOssLayerNormalization(keras.layers.Layer):
+    """A normalization layer for Gpt-Oss that implements RMS normalization."""
+
+    def __init__(self, epsilon=1e-6, **kwargs):
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        dim = input_shape[-1]
+        self.scale = self.add_weight(
+            name="scale",
+            trainable=True,
+            shape=(dim,),
+            initializer="ones",
+            dtype=self.variable_dtype,
+        )
+        self.built = True
+
+    def call(self, x):
+        x = ops.cast(x, "float32")
+        var = ops.mean(ops.power(x, 2), axis=-1, keepdims=True)
+        x = x * ops.rsqrt(var + self.epsilon)
+        return ops.cast(x * self.scale, self.compute_dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"epsilon": self.epsilon})
+        return config

keras_hub/src/models/gpt_oss/gpt_oss_presets.py

@@ -0,0 +1,51 @@
+backbone_presets = {
+    "gpt_oss_20b_en": {
+        "metadata": {
+            "description": (
+                "This preset has 21 billion total parameters, "
+                "with 3.6 billion active parameters, a 128k context "
+                "length, and is de-quantized from MXFP4."
+            ),
+            "params": 20_914_757_184,
+            "path": "gpt_oss",
+        },
+        "kaggle_handle": "kaggle://keras/gpt-oss/keras/gpt_oss_20b_en/1",
+    },
+    "gpt_oss_120b_en": {
+        "metadata": {
+            "description": (
+                "This preset has 117 billion total parameters, "
+                "with 5.1 billion active parameters, a 128k context "
+                "length, and is de-quantized from MXFP4."
+            ),
+            "params": 116_829_156_672,
+            "path": "gpt_oss",
+        },
+        "kaggle_handle": "kaggle://keras/gpt-oss/keras/gpt_oss_120b_en/1",
+    },
+    "gpt_oss_safeguard_20b_en": {
+        "metadata": {
+            "description": (
+                "Open-weight safety reasoning model with 21 billion "
+                "total parameters,with 3.6 billion active "
+                "parameters, a context length of over 128k, "
+                "and is de-quantized from MXFP4."
+            ),
+            "params": 20_914_757_184,
+            "path": "gpt_oss",
+        },
+        "kaggle_handle": "kaggle://keras/gpt-oss-safeguard/keras/gpt_oss_safeguard_20b_en/1",
+    },
+    "gpt_oss_safeguard_120b_en": {
+        "metadata": {
+            "description": (
+                "Open-weight safety reasoning model with 117 billion "
+                "total parameters,with 5.1 billion active parameters, "
+                "a 128k context length, and is de-quantized from MXFP4."
+            ),
+            "params": 116_829_156_672,
+            "path": "gpt_oss",
+        },
+        "kaggle_handle": "kaggle://keras/gpt-oss-safeguard/keras/gpt_oss_safeguard_120b_en/1",
+    },
+}

keras_hub/src/models/gpt_oss/gpt_oss_tokenizer.py

@@ -0,0 +1,39 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.gpt_oss.gpt_oss_backbone import GptOssBackbone
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.GptOssTokenizer",
+        "keras_hub.models.GptOssTokenizer",
+    ]
+)
+class GptOssTokenizer(BytePairTokenizer):
+    """A GptOss tokenizer using BytePair encoding.
+
+    Tokenizer is a subclass of `keras_hub.tokenizers.BytePairTokenizer`.
+    It uses a BytePair encoding model to tokenize strings. It also adds special
+    tokens for the start and end of a sequence.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules.
+    """
+
+    backbone_cls = GptOssBackbone
+
+    def __init__(self, vocabulary=None, merges=None, **kwargs):
+        """Initializes the GptOssTokenizer.
+
+        Args:
+            vocabulary: string or dict, maps token to integer ids.
+            merges: string or list, contains the merge rule.
+            **kwargs: Additional keyword arguments.
+        """
+        self.start_token_id = None
+        self._add_special_token("<|endoftext|>", "end_token")
+        self.pad_token_id = 0
+        super().__init__(vocabulary=vocabulary, merges=merges, **kwargs)

keras_hub/src/models/llama/llama_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.llama.llama_decoder import LlamaTransformerDecoder
 from keras_hub.src.models.llama.llama_layernorm import LlamaLayerNorm

keras_hub/src/models/llama3/llama3_presets.py

@@ -63,7 +63,7 @@ backbone_presets = {
             "params": 8030261248,
             "path": "llama3",
         },
-        "kaggle_handle": ("kaggle://keras/llama3/keras/lama3.1_instruct_8b/2"),
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.1_instruct_8b/2"),
     },
     "llama3.1_guard_8b": {
         "metadata": {

keras_hub/src/models/masked_lm.py

@@ -84,3 +84,25 @@ class MaskedLM(Task):
             weighted_metrics=weighted_metrics,
             **kwargs,
         )
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        # Check for standard backbone structure.
+        if not hasattr(backbone, "transformer_layers"):
+            return None
+
+        # Check for embedding.
+        embedding = getattr(backbone, "token_embedding", None)
+        if embedding is None:
+            embedding = getattr(backbone, "embedding", None)
+
+        if embedding is None:
+            return None
+
+        return {
+            "pre_block_layers": [embedding],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/mistral/mistral_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.mistral.mistral_layer_norm import (
     MistralLayerNormalization,

keras_hub/src/models/mixtral/mixtral_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.mixtral.mixtral_decoder import (
     MixtralTransformerDecoder,

keras_hub/src/models/moonshine/moonshine_backbone.py

@@ -1,9 +1,7 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.moonshine.moonshine_decoder import (
     MoonshineDecoderBlock,

keras_hub/src/models/pali_gemma/pali_gemma_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.models.pali_gemma.pali_gemma_decoder_block import (