keras-hub 0.25.1__py3-none-any.whl → 0.26.0.dev0__py3-none-any.whl

keras_hub/layers/__init__.py

@@ -138,6 +138,27 @@ from keras_hub.src.models.sam.sam_mask_decoder import (
 from keras_hub.src.models.sam.sam_prompt_encoder import (
     SAMPromptEncoder as SAMPromptEncoder,
 )
+from keras_hub.src.models.sam3.sam3_detr_decoder import (
+    SAM3DetrDecoder as SAM3DetrDecoder,
+)
+from keras_hub.src.models.sam3.sam3_detr_encoder import (
+    SAM3DetrEncoder as SAM3DetrEncoder,
+)
+from keras_hub.src.models.sam3.sam3_geometry_encoder import (
+    SAM3GeometryEncoder as SAM3GeometryEncoder,
+)
+from keras_hub.src.models.sam3.sam3_image_converter import (
+    SAM3ImageConverter as SAM3ImageConverter,
+)
+from keras_hub.src.models.sam3.sam3_mask_decoder import (
+    SAM3MaskDecoder as SAM3MaskDecoder,
+)
+from keras_hub.src.models.sam3.sam3_text_encoder import (
+    SAM3TextEncoder as SAM3TextEncoder,
+)
+from keras_hub.src.models.sam3.sam3_vision_encoder import (
+    SAM3VisionEncoder as SAM3VisionEncoder,
+)
 from keras_hub.src.models.segformer.segformer_image_converter import (
     SegFormerImageConverter as SegFormerImageConverter,
 )

keras_hub/models/__init__.py

@@ -211,6 +211,12 @@ from keras_hub.src.models.distil_bert.distil_bert_text_classifier_preprocessor i
 from keras_hub.src.models.distil_bert.distil_bert_tokenizer import (
     DistilBertTokenizer as DistilBertTokenizer,
 )
+from keras_hub.src.models.edrec.edrec_backbone import (
+    EdRecBackbone as EdRecBackbone,
+)
+from keras_hub.src.models.edrec.edrec_seq2seq_lm import (
+    EdRecSeq2SeqLM as EdRecSeq2SeqLM,
+)
 from keras_hub.src.models.efficientnet.efficientnet_backbone import (
     EfficientNetBackbone as EfficientNetBackbone,
 )
@@ -629,6 +635,15 @@ from keras_hub.src.models.roformer_v2.roformer_v2_text_classifier_preprocessor i
 from keras_hub.src.models.roformer_v2.roformer_v2_tokenizer import (
     RoformerV2Tokenizer as RoformerV2Tokenizer,
 )
+from keras_hub.src.models.rwkv7.rwkv7_backbone import (
+    RWKV7Backbone as RWKV7Backbone,
+)
+from keras_hub.src.models.rwkv7.rwkv7_causal_lm import (
+    RWKV7CausalLM as RWKV7CausalLM,
+)
+from keras_hub.src.models.rwkv7.rwkv7_causal_lm_preprocessor import (
+    RWKV7CausalLMPreprocessor as RWKV7CausalLMPreprocessor,
+)
 from keras_hub.src.models.sam.sam_backbone import SAMBackbone as SAMBackbone
 from keras_hub.src.models.sam.sam_image_segmenter import (
     SAMImageSegmenter as SAMImageSegmenter,
@@ -636,6 +651,18 @@ from keras_hub.src.models.sam.sam_image_segmenter import (
 from keras_hub.src.models.sam.sam_image_segmenter_preprocessor import (
     SAMImageSegmenterPreprocessor as SAMImageSegmenterPreprocessor,
 )
+from keras_hub.src.models.sam3.sam3_pc_backbone import (
+    SAM3PromptableConceptBackbone as SAM3PromptableConceptBackbone,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter import (
+    SAM3PromptableConceptImageSegmenter as SAM3PromptableConceptImageSegmenter,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter_preprocessor import (
+    SAM3PromptableConceptImageSegmenterPreprocessor as SAM3PromptableConceptImageSegmenterPreprocessor,
+)
+from keras_hub.src.models.sam3.sam3_tokenizer import (
+    SAM3Tokenizer as SAM3Tokenizer,
+)
 from keras_hub.src.models.segformer.segformer_backbone import (
     SegFormerBackbone as SegFormerBackbone,
 )

keras_hub/src/layers/modeling/non_max_supression.py

@@ -290,16 +290,19 @@ def non_max_suppression(
         "int32",
     )
     idx = ops.minimum(idx, num_boxes - 1)
+    idx = ops.cast(idx, "int32")
 
     index_offsets = ops.cast(ops.arange(batch_size) * num_boxes, "int32")
     take_along_axis_idx = ops.reshape(
         idx + ops.expand_dims(index_offsets, 1), [-1]
     )
+    take_along_axis_idx = ops.cast(take_along_axis_idx, "int32")
 
     if keras.backend.backend() != "tensorflow":
-        idx = ops.take_along_axis(
-            ops.reshape(sorted_indices, [-1]), take_along_axis_idx
+        sorted_indices_int = ops.cast(
+            ops.reshape(sorted_indices, [-1]), "int32"
         )
+        idx = ops.take_along_axis(sorted_indices_int, take_along_axis_idx)
     else:
         import tensorflow as tf
 

keras_hub/src/layers/modeling/reversible_embedding.py

@@ -1,281 +1,8 @@
-import inspect
-
 import keras
-from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 
 
 @keras_hub_export("keras_hub.layers.ReversibleEmbedding")
-class ReversibleEmbedding(keras.layers.Embedding):
-    """An embedding layer which can project backwards to the input dim.
-
-    This layer is an extension of `keras.layers.Embedding` for language models.
-    This layer can be called "in reverse" with `reverse=True`, in which case the
-    layer will linearly project from `output_dim` back to `input_dim`.
-
-    By default, the reverse projection will use the transpose of the
-    `embeddings` weights to project to `input_dim` (weights are "tied"). If
-    `tie_weights=False`, the model will use a separate, trainable variable for
-    reverse projection.
-
-    This layer has no bias terms.
-
-    Args:
-        input_dim: Integer. Size of the vocabulary,
-            i.e. maximum integer index + 1.
-        output_dim: Integer. Dimension of the dense embedding.
-        tie_weights: Boolean, whether or not the matrix for embedding and
-            the matrix for the `reverse` projection should share the same
-            weights.
-        embeddings_initializer: Initializer for the `embeddings`
-            matrix (see `keras.initializers`).
-        embeddings_regularizer: Regularizer function applied to
-            the `embeddings` matrix (see `keras.regularizers`).
-        embeddings_constraint: Constraint function applied to
-            the `embeddings` matrix (see `keras.constraints`).
-        mask_zero: Boolean, whether or not the input value 0 is a special
-            "padding" value that should be masked out.
-        reverse_dtype: The dtype for the reverse projection computation.
-            Defaults to the `compute_dtype` of the layer.
-        logit_soft_cap: If `logit_soft_cap` is set and `reverse=True`, the
-            output logits will be scaled by
-            `tanh(logits / logit_soft_cap) * logit_soft_cap`. This narrows the
-            range of output logits and can improve training.
-        **kwargs: other keyword arguments passed to `keras.layers.Embedding`,
-            including `name`, `trainable`, `dtype` etc.
-
-    Call arguments:
-        inputs: The tensor inputs to the layer.
-        reverse: Boolean. If `True` the layer will perform a linear projection
-            from `output_dim` to `input_dim`, instead of a normal embedding
-            call. Default to `False`.
-
-    Example:
-    ```python
-    batch_size = 16
-    vocab_size = 100
-    hidden_dim = 32
-    seq_length = 50
-
-    # Generate random inputs.
-    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
-
-    embedding = keras_hub.layers.ReversibleEmbedding(vocab_size, hidden_dim)
-    # Embed tokens to shape `(batch_size, seq_length, hidden_dim)`.
-    hidden_states = embedding(token_ids)
-    # Project hidden states to shape `(batch_size, seq_length, vocab_size)`.
-    logits = embedding(hidden_states, reverse=True)
-    ```
-
-    References:
-    - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
-    - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
-    """
-
-    def __init__(
-        self,
-        input_dim,
-        output_dim,
-        tie_weights=True,
-        embeddings_initializer="uniform",
-        embeddings_regularizer=None,
-        embeddings_constraint=None,
-        mask_zero=False,
-        reverse_dtype=None,
-        logit_soft_cap=None,
-        **kwargs,
-    ):
-        super().__init__(
-            input_dim,
-            output_dim,
-            embeddings_initializer=embeddings_initializer,
-            embeddings_regularizer=embeddings_regularizer,
-            embeddings_constraint=embeddings_constraint,
-            mask_zero=mask_zero,
-            **kwargs,
-        )
-        self.tie_weights = tie_weights
-        self.reverse_dtype = reverse_dtype
-        self.logit_soft_cap = logit_soft_cap
-
-    def build(self, inputs_shape=None):
-        super().build(inputs_shape)
-        if (
-            not self.tie_weights
-            and getattr(self, "quantization_mode", None) != "int8"
-        ):
-            self.reverse_embeddings = self.add_weight(
-                name="reverse_embeddings",
-                shape=(self.output_dim, self.input_dim),
-                initializer=self.embeddings_initializer,
-                dtype=self.dtype,
-            )
-
-    def call(self, inputs, reverse=False):
-        if reverse:
-            if self.tie_weights:
-                kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
-            else:
-                kernel = self.reverse_embeddings
-            if self.reverse_dtype is not None:
-                inputs = ops.cast(inputs, self.reverse_dtype)
-                kernel = ops.cast(kernel, self.reverse_dtype)
-            logits = ops.matmul(inputs, kernel)
-            # Optionally soft-cap logits.
-            if self.logit_soft_cap is not None:
-                soft_cap = self.logit_soft_cap
-                logits = ops.tanh(logits / soft_cap) * soft_cap
-            return logits
-
-        return super().call(inputs)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "tie_weights": self.tie_weights,
-                "reverse_dtype": self.reverse_dtype,
-                "logit_soft_cap": self.logit_soft_cap,
-            }
-        )
-        return config
-
-    def save_own_variables(self, store):
-        if not self.built:
-            return
-        super().save_own_variables(store)
-        target_variables = []
-        if not self.tie_weights:
-            # Store the reverse embedding weights as the last weights.
-            target_variables.append(self.reverse_embeddings)
-            if getattr(self, "quantization_mode", None) == "int8":
-                target_variables.append(self.reverse_embeddings_scale)
-            for i, variable in enumerate(target_variables, start=len(store)):
-                store[str(i)] = variable
-
-    def load_own_variables(self, store):
-        if not self.built:
-            self.build()
-        super().load_own_variables(store)
-        if not self.tie_weights:
-            # Last weights in the stores are the reverse embedding weights.
-            target_variables = [self.reverse_embeddings]
-            if getattr(self, "quantization_mode", None) == "int8":
-                target_variables.append(self.reverse_embeddings_scale)
-            for i, variable in enumerate(
-                target_variables, start=len(store) - len(target_variables)
-            ):
-                variable.assign(store[str(i)])
-
-    def compute_output_spec(self, inputs, reverse=False):
-        output_shape = list(inputs.shape)
-        if reverse:
-            output_shape[-1] = self.input_dim
-        else:
-            output_shape += [self.output_dim]
-        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)
-
-    # Quantization-related (int8) methods
-
-    def quantized_call(self, inputs, reverse=False):
-        # TODO (hongyu): This function could be removed once we add `*args` and
-        # `**kwargs` for `Embedding.quantized_call`
-        if self.quantization_mode == "int8":
-            return self._int8_call(inputs, reverse=reverse)
-        else:
-            self._quantization_mode_error(self.quantization_mode)
-
-    def _int8_build(self, embeddings_shape=None):
-        if (
-            "embeddings_shape"
-            in inspect.signature(super()._int8_build).parameters
-        ):
-            if embeddings_shape is None:
-                embeddings_shape = (self.input_dim, self.output_dim)
-            super()._int8_build(embeddings_shape=embeddings_shape)
-        else:
-            # Backward compatibility for older versions of Keras.
-            super()._int8_build()
-        self.inputs_quantizer = keras.quantizers.AbsMaxQuantizer(axis=-1)
-        if not self.tie_weights:
-            self.reverse_embeddings = self.add_weight(
-                name="reverse_embeddings",
-                shape=(self.output_dim, self.input_dim),
-                initializer="zeros",
-                dtype="int8",
-                trainable=False,
-            )
-            self.reverse_embeddings_scale = self.add_weight(
-                name="reverse_embeddings_scale",
-                shape=(self.input_dim,),
-                initializer="ones",
-                trainable=False,
-            )
-        self._is_quantized = True
-
-    def _int8_call(self, inputs, reverse=False):
-        if reverse:
-            if self.tie_weights:
-                kernel = ops.transpose(self._embeddings)
-                scale = ops.transpose(self.embeddings_scale)
-            else:
-                kernel = self.reverse_embeddings
-                scale = self.reverse_embeddings_scale
-            inputs, inputs_scale = self.inputs_quantizer(inputs)
-            logits = ops.matmul(inputs, kernel)
-            # De-scale outputs
-            logits = ops.cast(logits, self.compute_dtype)
-            logits = ops.divide(logits, ops.multiply(inputs_scale, scale))
-            # Optionally soft-cap logits.
-            if self.logit_soft_cap is not None:
-                soft_cap = self.logit_soft_cap
-                logits = ops.tanh(logits / soft_cap) * soft_cap
-            return logits
-
-        return super()._int8_call(inputs)
-
-    def quantize(self, mode, type_check=True, config=None):
-        del config
-        if type_check and type(self) is not ReversibleEmbedding:
-            raise self._not_implemented_error(self.quantize)
-
-        def abs_max_quantize(inputs, axis):
-            return keras.quantizers.abs_max_quantize(
-                inputs, axis=axis, to_numpy=True
-            )
-
-        if mode != "int8":
-            raise NotImplementedError(
-                "Invalid quantization mode. Expected 'int8'. "
-                f"Received: quantization_mode={mode}"
-            )
-
-        embeddings_shape = (self.input_dim, self.output_dim)
-        if mode == "int8":
-            embeddings, embeddings_scale = abs_max_quantize(
-                self._embeddings, axis=-1
-            )
-            embeddings_scale = ops.squeeze(embeddings_scale, axis=-1)
-            del self._embeddings
-            if not self.tie_weights:
-                reverse_embeddings, reverse_embeddings_scale = abs_max_quantize(
-                    self.reverse_embeddings, axis=0
-                )
-                reverse_embeddings_scale = ops.squeeze(
-                    reverse_embeddings_scale, axis=0
-                )
-                del self.reverse_embeddings
-        self.quantized_build(embeddings_shape, mode)
-        if mode == "int8":
-            self._embeddings.assign(embeddings)
-            self.embeddings_scale.assign(embeddings_scale)
-            if not self.tie_weights:
-                self.reverse_embeddings.assign(reverse_embeddings)
-                self.reverse_embeddings_scale.assign(reverse_embeddings_scale)
-
-        if self.dtype_policy.quantization_mode is None:
-            policy = keras.dtype_policies.get(
-                f"{mode}_from_{self.dtype_policy.name}"
-            )
-            self.dtype_policy = policy
+class ReversibleEmbedding(keras.layers.ReversibleEmbedding):
+    pass

keras_hub/src/layers/modeling/token_and_position_embedding.py

@@ -1,10 +1,10 @@
 import keras
+from keras.layers import ReversibleEmbedding
+from keras.src.backend import get_keras_mask
+from keras.src.backend import set_keras_mask
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.utils.keras_utils import clone_initializer
 
 
@@ -128,10 +128,10 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             positions=positions,
         )
         outputs = embedded_tokens + embedded_positions
+        mask = get_keras_mask(embedded_tokens)
+        if mask is not None:
+            set_keras_mask(outputs, mask)
         return outputs
 
-    def compute_mask(self, inputs, mask=None):
-        return self.token_embedding.compute_mask(inputs, mask=mask)
-
     def compute_output_shape(self, input_shape):
         return tuple(input_shape) + (self.embedding_dim,)

keras_hub/src/layers/modeling/transformer_layer_utils.py

@@ -1,11 +1,12 @@
 from absl import logging
 from keras import ops
+from keras.src.backend import get_keras_mask
 
 
 def _check_masks_shapes(inputs, padding_mask, attention_mask):
     mask = padding_mask
-    if hasattr(inputs, "_keras_mask") and mask is None:
-        mask = inputs._keras_mask
+    if mask is None:
+        mask = get_keras_mask(inputs)
     if mask is not None:
         if len(mask.shape) != 2:
             raise ValueError(
@@ -68,17 +69,16 @@ def merge_padding_and_attention_mask(
         returned mask is padding_mask with one additional axis.
     """
     _check_masks_shapes(inputs, padding_mask, attention_mask)
-    mask = padding_mask
-    if hasattr(inputs, "_keras_mask"):
-        if mask is None:
-            # If no padding mask is explicitly provided, we look for padding
-            # mask from the input data.
-            mask = inputs._keras_mask
-        else:
+    # We look for a padding mask from the input data.
+    mask = get_keras_mask(inputs)
+    # But if padding mask is explicitly provided, we use it.
+    if padding_mask is not None:
+        if mask is not None:
             logging.warning(
                 "You are explicitly setting `padding_mask` while the `inputs` "
                 "have built-in mask, so the built-in mask is ignored."
             )
+        mask = padding_mask
     if mask is not None:
         # Add an axis for broadcasting, the attention mask should be 2D
         # (not including the batch axis).

keras_hub/src/layers/preprocessing/masked_lm_mask_generator.py

@@ -7,9 +7,11 @@ from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
-    import tensorflow_text as tf_text
 except ImportError:
     tf = None
+try:
+    import tensorflow_text as tf_text
+except ImportError:
     tf_text = None
 
 

keras_hub/src/layers/preprocessing/multi_segment_packer.py

@@ -8,9 +8,11 @@ from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
-    import tensorflow_text as tf_text
 except ImportError:
     tf = None
+try:
+    import tensorflow_text as tf_text
+except ImportError:
     tf_text = None
 
 

keras_hub/src/models/albert/albert_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate

keras_hub/src/models/backbone.py

@@ -107,6 +107,9 @@ class Backbone(keras.Model):
     def from_config(cls, config):
         # The default `from_config()` for functional models will return a
         # vanilla `keras.Model`. We override it to get a subclass instance back.
+        config = config.copy()
+        if "dtype" in config and isinstance(config["dtype"], dict):
+            config["dtype"] = keras.dtype_policies.get(config["dtype"])
         return cls(**config)
 
     @classproperty

keras_hub/src/models/bart/bart_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone

keras_hub/src/models/bert/bert_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate
@@ -35,7 +33,7 @@ class BertBackbone(Backbone):
         vocabulary_size: int. The size of the token vocabulary.
         num_layers: int. The number of transformer layers.
         num_heads: int. The number of attention heads for each transformer.
-            The hidden size must be divisible by the number of attention heads.
+            The hidden_dim must be divisible by the number of attention heads.
         hidden_dim: int. The size of the transformer encoding and pooler layers.
         intermediate_dim: int. The output dimension of the first Dense layer in
             a two-layer feedforward network for each transformer.

keras_hub/src/models/bloom/bloom_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.bloom.bloom_decoder import BloomDecoder
 

keras_hub/src/models/causal_lm.py

@@ -196,7 +196,7 @@ class CausalLM(Task):
 
                 # Create an explicit tuple of all variable state.
                 state = (
-                    self.sampler.variables,
+                    [v.value for v in self.sampler.variables],
                     # Use the explicit variable.value to preserve the
                     # sharding spec of distribution.
                     [v.value for v in self.trainable_variables],
@@ -431,7 +431,7 @@ class CausalLM(Task):
         self.generate_function = None
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         backbone = self.backbone

keras_hub/src/models/deberta_v3/deberta_v3_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.deberta_v3.disentangled_attention_encoder import (
     DisentangledAttentionEncoder,