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

keras_hub/src/models/sam3/sam3_text_encoder.py

@@ -0,0 +1,212 @@
+import keras
+from keras import layers
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.token_and_position_embedding import (
+    TokenAndPositionEmbedding,
+)
+from keras_hub.src.models.sam3.sam3_utils import create_bidirectional_mask
+
+
+class CLIPEncoderLayer(layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        num_heads,
+        intermediate_dim,
+        intermediate_activation="gelu",
+        layer_norm_epsilon=1e-5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.num_heads = int(num_heads)
+        self.intermediate_dim = int(intermediate_dim)
+        self.intermediate_activation = intermediate_activation
+        self.layer_norm_epsilon = float(layer_norm_epsilon)
+
+        self.layer_norm_1 = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm_1",
+        )
+        self.attention = layers.MultiHeadAttention(
+            num_heads,
+            hidden_dim // num_heads,
+            dtype=self.dtype_policy,
+            name="attention",
+        )
+        self.layer_norm_2 = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm_2",
+        )
+        self.dense_1 = layers.Dense(
+            self.intermediate_dim, dtype=self.dtype_policy, name="dense_1"
+        )
+        self.activation = layers.Activation(
+            intermediate_activation, dtype=self.dtype_policy, name="activation"
+        )
+        self.dense_2 = layers.Dense(
+            self.hidden_dim, dtype=self.dtype_policy, name="dense_2"
+        )
+
+    def build(self, inputs_shape, attention_mask_shape):
+        self.layer_norm_1.build(inputs_shape)
+        self.attention.build(inputs_shape, inputs_shape, inputs_shape)
+        self.layer_norm_2.build(inputs_shape)
+        self.dense_1.build(inputs_shape)
+        input_shape = self.dense_1.compute_output_shape(inputs_shape)
+        self.dense_2.build(input_shape)
+
+    def compute_output_shape(self, inputs_shape, attention_mask_shape):
+        outputs_shape = list(inputs_shape)
+        outputs_shape[-1] = self.hidden_dim
+        return outputs_shape
+
+    def call(self, inputs, attention_mask, training=None):
+        residual = inputs
+        x = self.layer_norm_1(inputs)
+        x = self.attention(
+            x,
+            x,
+            x,
+            attention_mask=ops.cast(attention_mask, dtype="bool"),
+            training=training,
+        )
+        x = ops.add(residual, x)
+
+        residual = x
+        x = self.dense_1(self.layer_norm_2(residual))
+        x = self.activation(x)
+        x = self.dense_2(x)
+        x = ops.add(residual, x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "num_heads": self.num_heads,
+                "intermediate_dim": self.intermediate_dim,
+                "intermediate_activation": self.intermediate_activation,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+
+@keras_hub_export("keras_hub.layers.SAM3TextEncoder")
+class SAM3TextEncoder(layers.Layer):
+    """A text encoder for the Segment Anything Model 3 (SAM3).
+
+    This layer implements a CLIP-style text encoder. It processes token IDs and
+    padding masks to produce text embeddings that are used as prompts for
+    segmentation.
+
+    Args:
+        vocabulary_size: int. The size of the vocabulary.
+        embedding_dim: int. The dimension of the token embeddings.
+        hidden_dim: int. The hidden dimension of the transformer layers.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads.
+        intermediate_dim: int. The dimension of the intermediate layer in the
+            transformer's MLP.
+        intermediate_activation: str. The activation function for the
+            transformer layers. Defaults to `"gelu"`.
+        max_sequence_length: int. The maximum sequence length. Defaults to
+            `32`.
+        layer_norm_epsilon: float. The epsilon value for layer normalization.
+            Defaults to `1e-6`.
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        embedding_dim,
+        hidden_dim,
+        num_layers,
+        num_heads,
+        intermediate_dim,
+        intermediate_activation="gelu",
+        max_sequence_length=32,
+        layer_norm_epsilon=1e-5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocabulary_size = int(vocabulary_size)
+        self.embedding_dim = int(embedding_dim)
+        self.hidden_dim = int(hidden_dim)
+        self.num_layers = int(num_layers)
+        self.num_heads = int(num_heads)
+        self.intermediate_dim = int(intermediate_dim)
+        self.intermediate_activation = intermediate_activation
+        self.max_sequence_length = int(max_sequence_length)
+        self.layer_norm_epsilon = float(layer_norm_epsilon)
+
+        self.embedding = TokenAndPositionEmbedding(
+            vocabulary_size=self.vocabulary_size,
+            sequence_length=self.max_sequence_length,
+            embedding_dim=self.embedding_dim,
+            dtype=self.dtype_policy,
+            name="embedding",
+        )
+        self.encoder_layers = [
+            CLIPEncoderLayer(
+                self.hidden_dim,
+                self.num_heads,
+                self.intermediate_dim,
+                self.intermediate_activation,
+                dtype=self.dtype_policy,
+                name=f"encoder_layer_{i}",
+            )
+            for i in range(self.num_layers)
+        ]
+        self.layer_norm = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm",
+        )
+
+    def build(self, token_ids_shape, padding_masks_shape):
+        self.embedding.build(token_ids_shape)
+        x_shape = self.embedding.compute_output_shape(token_ids_shape)
+        for layer in self.encoder_layers:
+            layer.build(x_shape, padding_masks_shape)
+        self.layer_norm.build(x_shape)
+
+    def call(self, token_ids, padding_masks, training=None):
+        x = self.embedding(token_ids, training=training)
+        padding_masks = create_bidirectional_mask(x, padding_masks)
+        for layer in self.encoder_layers:
+            x = layer(x, padding_masks, training=training)
+        x = self.layer_norm(x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "embedding_dim": self.embedding_dim,
+                "hidden_dim": self.hidden_dim,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "intermediate_dim": self.intermediate_dim,
+                "intermediate_activation": self.intermediate_activation,
+                "max_sequence_length": self.max_sequence_length,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, token_ids_shape, padding_masks_shape):
+        return self.embedding.compute_output_shape(token_ids_shape)
+
+    def compute_output_spec(self, token_ids, padding_masks):
+        output_shape = self.compute_output_shape(
+            token_ids.shape, padding_masks.shape
+        )
+        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)

keras_hub/src/models/sam3/sam3_tokenizer.py

@@ -0,0 +1,65 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.clip.clip_tokenizer import CLIPTokenizer
+from keras_hub.src.models.sam3.sam3_pc_backbone import (
+    SAM3PromptableConceptBackbone,
+)
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.SAM3Tokenizer",
+        "keras_hub.models.SAM3Tokenizer",
+    ]
+)
+class SAM3Tokenizer(CLIPTokenizer):
+    """A SAM3 tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_hub.tokenizers.BytePairTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by SAM3
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a SAM3 preset.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    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. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space.
+
+    Examples:
+
+    ```python
+    # Unbatched input.
+    tokenizer = keras_hub.models.SAM3Tokenizer.from_preset("sam3_pcs")
+    tokenizer("The quick brown fox jumped.")
+
+    # Batched input.
+    tokenizer(["The quick brown fox jumped.", "The fox slept."])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+    ```
+    """
+
+    backbone_cls = SAM3PromptableConceptBackbone
+
+    def __init__(self, vocabulary=None, merges=None, **kwargs):
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            pad_with_end_token=True,
+            **kwargs,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        del config["pad_with_end_token"]  # Always True for SAM3Tokenizer.
+        return config

keras_hub/src/models/sam3/sam3_utils.py

@@ -0,0 +1,134 @@
+from keras import ops
+
+
+def window_partition(x, height, width, window_size, hidden_dim):
+    x = ops.reshape(
+        x,
+        (
+            -1,
+            height // window_size,
+            window_size,
+            width // window_size,
+            window_size,
+            hidden_dim,
+        ),
+    )
+    x = ops.transpose(x, axes=(0, 1, 3, 2, 4, 5))
+    x = ops.reshape(x, (-1, window_size, window_size, hidden_dim))
+    return x
+
+
+def window_unpartition(x, height, width, window_size, hidden_dim):
+    x = ops.reshape(
+        x,
+        (
+            -1,
+            height // window_size,
+            width // window_size,
+            window_size,
+            window_size,
+            hidden_dim,
+        ),
+    )
+    x = ops.transpose(x, axes=(0, 1, 3, 2, 4, 5))
+    x = ops.reshape(x, (-1, height, width, hidden_dim))
+    return x
+
+
+def box_cxcywh_to_xyxy(boxes):
+    x_c, y_c, w, h = ops.unstack(boxes, num=4, axis=-1)
+    return ops.stack(
+        [
+            ops.subtract(x_c, ops.multiply(0.5, w)),
+            ops.subtract(y_c, ops.multiply(0.5, h)),
+            ops.add(x_c, ops.multiply(0.5, w)),
+            ops.add(y_c, ops.multiply(0.5, h)),
+        ],
+        axis=-1,
+    )
+
+
+def concatenate_padded_sequences(
+    sequence1, mask1, sequence_len1, sequence2, mask2, sequence_len2, hidden_dim
+):
+    """Concatenate two sequences with padding masks.
+
+    Args:
+        sequence1: A tensor of shape (batch_size, sequence_length1, hidden_dim).
+        mask1: A boolean tensor of shape (batch_size, sequence_length1).
+        sequence2: A tensor of shape (batch_size, sequence_length2, hidden_dim).
+        mask2: A boolean tensor of shape (batch_size, sequence_length2).
+        hidden_dim: An integer representing the hidden dimension.
+
+    Returns:
+        concatenated_sequence: A tensor of shape
+            (batch_size, sequence_length1 + sequence_length2, hidden_dim).
+        concatenated_mask: A boolean tensor of shape
+            (batch_size, sequence_length1 + sequence_length2).
+    """
+    batch_size = ops.shape(sequence1)[0]
+    max_length = sequence_len1 + sequence_len2
+
+    actual_sequence_1_lengths = ops.sum(ops.cast(mask1, dtype="int32"), axis=1)
+    actual_sequence_2_lengths = ops.sum(ops.cast(mask2, dtype="int32"), axis=1)
+    final_lengths = ops.add(
+        actual_sequence_1_lengths, actual_sequence_2_lengths
+    )
+
+    concatenated_mask = ops.less(
+        ops.tile(
+            ops.expand_dims(ops.arange(max_length, dtype="int32"), axis=0),
+            [batch_size, 1],
+        ),
+        ops.expand_dims(final_lengths, axis=1),
+    )
+    concatenated_sequence = ops.concatenate(
+        [
+            sequence1,
+            ops.zeros(
+                (batch_size, sequence_len2, hidden_dim), dtype=sequence1.dtype
+            ),
+        ],
+        axis=1,
+    )
+
+    # Create the indices.
+    indices = ops.tile(
+        ops.expand_dims(ops.arange(sequence_len2, dtype="int32"), axis=0),
+        [batch_size, 1],
+    )
+    indices = ops.add(
+        indices, ops.expand_dims(actual_sequence_1_lengths, axis=-1)
+    )
+    # Adjust the indices to account for batch dimension.
+    to_add = ops.multiply(ops.arange(batch_size, dtype="int32"), max_length)
+    indices = ops.add(
+        indices, ops.cast(ops.expand_dims(to_add, axis=-1), "int32")
+    )
+    # `ops.scatter_update` requires 2D indices. We flatten the inputs before
+    # scattering and reshape back after.
+    flat_concatenated_sequence = ops.scatter_update(
+        ops.reshape(concatenated_sequence, (-1, hidden_dim)),
+        ops.reshape(indices, (-1, 1)),
+        ops.reshape(
+            ops.cast(sequence2, concatenated_sequence.dtype), (-1, hidden_dim)
+        ),
+    )
+
+    concatenated_sequence = ops.reshape(
+        flat_concatenated_sequence, (batch_size, -1, hidden_dim)
+    )
+    return concatenated_sequence, concatenated_mask
+
+
+def create_bidirectional_mask(input_embeds, attention_mask):
+    seq_len = ops.shape(input_embeds)[1]
+    attention_mask = attention_mask[:, None, None, :]
+    return ops.tile(attention_mask, (1, 1, seq_len, 1))
+
+
+def inverse_sigmoid(x, eps=1e-3):
+    x = ops.clip(x, 0.0, 1.0)
+    x1 = ops.maximum(x, eps)
+    x2 = ops.maximum(ops.subtract(1.0, x), eps)
+    return ops.log(ops.divide(x1, x2))