keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/causal_lm.py

@@ -274,6 +274,7 @@ class CausalLM(Task):
         inputs,
         max_length=None,
         stop_token_ids="auto",
+        strip_prompt=False,
     ):
         """Generate text given prompt `inputs`.
 
@@ -302,13 +303,18 @@ class CausalLM(Task):
                 `preprocessor`. If `preprocessor` is `None`, `inputs` should be
                 should be padded to the desired maximum length and this argument
                 will be ignored.
-            stop_token_ids: Optional. `None`, "auto", or tuple of token ids. Defaults
-                to "auto" which uses the `preprocessor.tokenizer.end_token_id`.
-                Not specifying a processor will produce an error. None stops
-                generation after generating `max_length` tokens. You may also
-                specify a list of token id's the model should stop on. Note that
-                sequences of tokens will each be interpreted as a stop token,
-                multi-token stop sequences are not supported.
+            stop_token_ids: Optional. `None`, "auto", or tuple of token ids.
+                Defaults to "auto" which uses the
+                `preprocessor.tokenizer.end_token_id`. Not specifying a
+                processor will produce an error. None stops generation after
+                generating `max_length` tokens. You may also specify a list of
+                token id's the model should stop on. Note that sequences of
+                tokens will each be interpreted as a stop token, multi-token
+                stop sequences are not supported.
+            strip_prompt: Optional. By default, generate() returns the full
+                prompt followed by its completion generated by the model. If
+                this option is set to True, only the newly generated text is
+                returned.
         """
         # Setup our three main passes.
         # 1. Optionally preprocessing strings to dense integer tensors.
@@ -318,14 +324,19 @@ class CausalLM(Task):
 
         if self.preprocessor is None and stop_token_ids == "auto":
             raise ValueError(
-                'A `preprocessor` must be attached to the model if `stop_token_ids="auto"`. '
-                "Currently `preprocessor=None`. To call `generate()` with preprocessing "
-                "detached, either pass `stop_token_ids=None` to always generate until "
-                "`max_length` or pass a tuple of token ids that should terminate generation "
+                "A `preprocessor` must be attached to the model if "
+                '`stop_token_ids="auto"`. Currently `preprocessor=None`. To '
+                "call `generate()` with preprocessing detached, either pass "
+                "`stop_token_ids=None` to always generate until `max_length` "
+                "or pass a tuple of token ids that should terminate generation "
                 "as `stop_token_ids`."
             )
         elif stop_token_ids == "auto":
             stop_token_ids = [self.preprocessor.tokenizer.end_token_id]
+            # Some models like Llama3 use two end tokens: <|eot_id|> in
+            # "instruct" versions and <|end_of_text|> in others.
+            if hasattr(self.preprocessor.tokenizer, "end_token2_id"):
+                stop_token_ids.append(self.preprocessor.tokenizer.end_token2_id)
 
         def preprocess(x):
             return self.preprocessor.generate_preprocess(
@@ -335,6 +346,34 @@ class CausalLM(Task):
         def generate(x):
             return generate_function(x, stop_token_ids=stop_token_ids)
 
+        def strip_prompt_function(x, prompt):
+            # This function removes the prompt from the generated
+            # response, in a batch-friendly fashion.
+            y = {}
+            prompt_mask = prompt["padding_mask"]
+            seq_len = prompt_mask.shape[1]
+
+            # We need to shift every output sequence by the size of the prompt.
+            shifts = -ops.sum(ops.cast(prompt_mask, "int"), axis=1) % seq_len
+            ix = ops.arange(seq_len, dtype="int")
+            ix = ops.expand_dims(ix, axis=0) - ops.expand_dims(shifts, axis=1)
+
+            # This produces the desired shift (in fact a rollover).
+            def roll_sequence(seq):
+                return ops.take_along_axis(seq, ix, axis=1)
+
+            # The shifting rolls the content over so the prompt is at the end of
+            # the sequence and the generated text is at the beginning. We mask
+            # it to retain the generated text only.
+            y["padding_mask"] = ops.logical_xor(
+                roll_sequence(prompt_mask), roll_sequence(x["padding_mask"])
+            )
+            # we assume the mask is enough and there is no need to zero-out the
+            # values
+            y["token_ids"] = roll_sequence(x["token_ids"])
+
+            return y
+
         def postprocess(x):
             return self.preprocessor.generate_postprocess(x)
 
@@ -343,7 +382,12 @@ class CausalLM(Task):
 
         if self.preprocessor is not None:
             inputs = [preprocess(x) for x in inputs]
-        outputs = [generate(x) for x in inputs]
+
+        if strip_prompt:
+            outputs = [strip_prompt_function(generate(x), x) for x in inputs]
+        else:
+            outputs = [generate(x) for x in inputs]
+
         if self.preprocessor is not None:
             outputs = [postprocess(x) for x in outputs]
 

keras_hub/src/models/clip/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.clip.clip_backbone import CLIPBackbone
+from keras_hub.src.models.clip.clip_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, CLIPBackbone)

keras_hub/src/models/clip/clip_backbone.py

@@ -0,0 +1,286 @@
+import math
+
+from keras import layers
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+
+
+class CLIPVisionPooler(layers.Layer):
+    """The vision pooler layer of CLIP.
+
+    `CLIPVisionPooler` will extracts the first token (index `0`) from the
+    sequence of the vision embeddings as the pooled outputs.
+
+    Call arguments:
+        vision_embeddings: A tensor of shape
+            `(batch_size, sequence_length, hidden_dim)`.
+    """
+
+    def call(self, vision_embeddings):
+        return vision_embeddings[:, 0, :]
+
+    def compute_output_shape(self, input_shape):
+        return (input_shape[0], input_shape[-1])
+
+
+class CLIPTextPooler(layers.Layer):
+    """The text pooler layer of CLIP.
+
+    `CLIPTextPooler` extracts the text embeddings at the positions of EOS tokens
+    as the pooled outputs.
+
+    Call arguments:
+        text_embeddings: A tensor of shape
+            `(batch_size, sequence_length, hidden_dim)`.
+        token_ids: A tensor of shape `(batch_size, max_tokens)`, used to
+            identify the positions of EOS tokens.
+    """
+
+    def call(self, text_embeddings, token_ids):
+        # `keepdims` is not supported in `keras<=3.1`.
+        eos_index = ops.argmax(token_ids, axis=-1)
+        eos_index = ops.expand_dims(eos_index, axis=-1)
+        eos_index = ops.expand_dims(eos_index, axis=-1)
+        pooled_outputs = ops.take_along_axis(text_embeddings, eos_index, axis=1)
+        return ops.squeeze(pooled_outputs, axis=1)
+
+    def compute_output_shape(self, input_shape):
+        return (input_shape[0], input_shape[-1])
+
+
+class CLIPHead(layers.Layer):
+    """The head layer of CLIP.
+
+    `CLIPHead` takes `vision_embedding` and `text_embedding` as inputs to
+    compute the corresponding logits. Both embeddings are L2 normalized and used
+    to compute pairwise cosine similarity. The resulting logits are then scaled
+    by a learnable `logit_scale` parameter.
+
+    Call arguments:
+        vision_embedding: A tensor of shape `(batch_size, hidden_dim)`.
+        text_embedding: A tensor of shape `(batch_size, hidden_dim)`.
+    """
+
+    def build(self, input_shape):
+        self.logit_scale = self.add_weight(
+            shape=(),
+            initializer=lambda *a, **kw: math.log(1 / 0.07),
+            trainable=True,
+            dtype=self.variable_dtype,
+            name="logit_scale",
+        )
+
+    def call(self, vision_embedding, text_embedding):
+        normalized_vision_embedding = ops.sqrt(
+            ops.sum(ops.power(vision_embedding, 2), axis=-1, keepdims=True)
+        )
+        normalized_text_embedding = ops.sqrt(
+            ops.sum(ops.power(text_embedding, 2), axis=-1, keepdims=True)
+        )
+        vision_embedding = vision_embedding / normalized_vision_embedding
+        text_embedding = text_embedding / normalized_text_embedding
+        logit_scale = ops.exp(self.logit_scale)
+        text_logits = (
+            ops.matmul(
+                text_embedding,
+                ops.transpose(vision_embedding),
+            )
+            * logit_scale
+        )
+        vision_logits = ops.transpose(text_logits)
+        return vision_logits, text_logits
+
+    def compute_output_shape(
+        self, vision_embedding_shape, text_embedding_shape
+    ):
+        vision_logits_shape = (
+            vision_embedding_shape[0],
+            text_embedding_shape[0],
+        )
+        text_logits_shape = (
+            text_embedding_shape[0],
+            vision_embedding_shape[0],
+        )
+        return vision_logits_shape, text_logits_shape
+
+
+@keras_hub_export("keras_hub.models.CLIPBackbone")
+class CLIPBackbone(Backbone):
+    """CLIP core network with hyperparameters.
+
+    This backbone implements the base architecture for Contrastive
+    Language-Image Pretraining (CLIP) model. It includes a vision and text
+    encoders and the corresponding projection layers. This backbone will output
+    the final logit scores corresponding to each image and token input. These
+    values are cosine similarities between the corresponding image and text
+    features.
+
+    The default constructor gives a fully customizable, randomly initialized
+    CLIP model with any number of layers, heads, and embedding dimensions. To
+    load preset architectures and weights, use the `from_preset` constructor.
+
+    Args:
+        vision_encoder: The CLIP vision encoder for encoding the input images.
+        text_encoder: The CLIP text encoder for encoding the input tokens.
+        projection_dim: int. The size of the projection layer.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done a float32 precision regardless of dtype.
+
+    Example:
+    ```python
+    input_data = {
+        "images": np.ones(shape=(1, 224, 224, 3), dtype="float32"),
+        "token_ids": np.ones(shape=(1, 12), dtype="int32"),
+    }
+
+    # Pretrained CLIP model.
+    model = keras_hub.models.CLIPBackbone.from_preset("clip_vit_base_patch32")
+    model(input_data)
+
+    # Randomly initialized CLIP model with custom config.
+    vision_encoder = keras_hub.models.CLIPVisionEncoder(
+        patch_size=32,
+        hidden_dim=768,
+        num_layers=8,
+        num_heads=8,
+        intermediate_dim=2048,
+        image_shape=(384, 384, 3),
+    )
+    text_encoder = keras_hub.models.CLIPTextEncoder(
+        vocabulary_size=49408,
+        embedding_dim=768,
+        hidden_dim=768,
+        num_layers=8,
+        num_heads=8,
+        intermediate_dim=2048,
+    )
+    model = keras_hub.models.CLIPBackbone(
+        vision_encoder=vision_encoder,
+        text_encoder=text_encoder,
+        projection_dim=256,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vision_encoder,
+        text_encoder,
+        projection_dim,
+        dtype=None,
+        name=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.vision_encoder = vision_encoder
+        self.text_encoder = text_encoder
+        self.vision_pooler = CLIPVisionPooler(dtype=dtype, name="vision_pooler")
+        self.text_pooler = CLIPTextPooler(dtype=dtype, name="text_pooler")
+        self.vision_projection = layers.Dense(
+            projection_dim,
+            use_bias=False,
+            dtype=dtype,
+            name="vision_projection",
+        )
+        self.text_projection = layers.Dense(
+            projection_dim,
+            use_bias=False,
+            dtype=dtype,
+            name="text_projection",
+        )
+        self.clip_head = CLIPHead(dtype=dtype, name="clip_head")
+
+        # === Functional Model ===
+        image_input = layers.Input(
+            shape=self.vision_encoder.image_shape, name="images"
+        )
+        token_id_input = layers.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        vision_embeddings = self.get_vision_embeddings(image_input)
+        text_embeddings = self.get_text_embeddings(token_id_input)
+        vision_logits, text_logits = self.clip_head(
+            vision_embeddings, text_embeddings
+        )
+
+        super().__init__(
+            inputs={
+                "images": image_input,
+                "token_ids": token_id_input,
+            },
+            outputs={
+                "vision_logits": vision_logits,
+                "text_logits": text_logits,
+            },
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.projection_dim = projection_dim
+
+    def get_vision_embeddings(self, images):
+        """Get the embeddings from the vision encoder.
+
+        Args:
+            images: The input tensor for the vision encoder.
+
+        Returns:
+            The output embeddings obtained by applying projection layer to the
+            pooled output of the vision encoder.
+        """
+        vision_outputs = self.vision_encoder({"images": images})
+        vision_outputs = self.vision_pooler(vision_outputs)
+        return self.vision_projection(vision_outputs)
+
+    def get_text_embeddings(self, token_ids):
+        """Get the embeddings from the text encoder.
+
+        Args:
+            token_ids: The input int tensor for the text encoder.
+
+        Returns:
+            The output embeddings obtained by applying projection layer to the
+            pooled output of the text encoder.
+        """
+        text_outputs = self.text_encoder({"token_ids": token_ids})
+        text_outputs = self.text_pooler(text_outputs, token_ids)
+        return self.text_projection(text_outputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vision_encoder": layers.serialize(self.vision_encoder),
+                "text_encoder": layers.serialize(self.text_encoder),
+                "projection_dim": self.projection_dim,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        config = config.copy()
+
+        # Propagate `dtype` to submodels if needed.
+        if "dtype" in config and config["dtype"] is not None:
+            dtype_config = config["dtype"]
+            if "dtype" not in config["vision_encoder"]["config"]:
+                config["vision_encoder"]["config"]["dtype"] = dtype_config
+            if "dtype" not in config["text_encoder"]["config"]:
+                config["text_encoder"]["config"]["dtype"] = dtype_config
+
+        # We expect submodels to be instantiated.
+        config["vision_encoder"] = layers.deserialize(
+            config["vision_encoder"], custom_objects=custom_objects
+        )
+        config["text_encoder"] = layers.deserialize(
+            config["text_encoder"], custom_objects=custom_objects
+        )
+        return cls(**config)

keras_hub/src/models/clip/clip_encoder_block.py

@@ -7,6 +7,16 @@ def quick_gelu(x):
     return x * ops.sigmoid(1.702 * x)
 
 
+# TODO: Deprecate this in favor of `keras.layers.MultiHeadAttention` once the
+# dtype compatibility issue is resolved.
+class CLIPMultiHeadAttention(layers.MultiHeadAttention):
+    def _masked_softmax(self, attention_scores, attention_mask=None):
+        attention_scores = super()._masked_softmax(
+            attention_scores, attention_mask
+        )
+        return ops.cast(attention_scores, self._value_dense.compute_dtype)
+
+
 class CLIPEncoderBlock(layers.Layer):
     def __init__(
         self,
@@ -14,6 +24,7 @@ class CLIPEncoderBlock(layers.Layer):
         num_heads,
         intermediate_dim,
         intermediate_activation="quick_gelu",
+        use_causal_mask=True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -26,21 +37,22 @@ class CLIPEncoderBlock(layers.Layer):
         self.num_heads = num_heads
         self.intermediate_dim = intermediate_dim
         self.intermediate_activation = intermediate_activation
+        self.use_causal_mask = use_causal_mask
 
         if intermediate_activation == "quick_gelu":
             intermediate_activation = quick_gelu
 
         self.layer_norm_1 = layers.LayerNormalization(
-            epsilon=1e-5, dtype="float32", name="layer_norm_1"
+            epsilon=1e-5, dtype=self.dtype_policy, name="layer_norm_1"
         )
-        self.attention = layers.MultiHeadAttention(
+        self.attention = CLIPMultiHeadAttention(
             num_heads,
             hidden_dim // num_heads,
             dtype=self.dtype_policy,
             name="attention",
         )
         self.layer_norm_2 = layers.LayerNormalization(
-            epsilon=1e-5, dtype="float32", name="layer_norm_2"
+            epsilon=1e-5, dtype=self.dtype_policy, name="layer_norm_2"
         )
         self.dense_1 = layers.Dense(
             self.intermediate_dim, dtype=self.dtype_policy, name="dense_1"
@@ -73,7 +85,9 @@ class CLIPEncoderBlock(layers.Layer):
     def call(self, x, training=None):
         residual = x
         x = self.layer_norm_1(x)
-        x = self.attention(x, x, x, training=training, use_causal_mask=True)
+        x = self.attention(
+            x, x, x, training=training, use_causal_mask=self.use_causal_mask
+        )
         x = ops.add(residual, x)
 
         residual = x
@@ -91,6 +105,7 @@ class CLIPEncoderBlock(layers.Layer):
                 "num_heads": self.num_heads,
                 "intermediate_dim": self.intermediate_dim,
                 "intermediate_activation": self.intermediate_activation,
+                "use_causal_mask": self.use_causal_mask,
             }
         )
         return config

keras_hub/src/models/clip/clip_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.clip.clip_backbone import CLIPBackbone
+
+
+@keras_hub_export("keras_hub.layers.CLIPImageConverter")
+class CLIPImageConverter(ImageConverter):
+    backbone_cls = CLIPBackbone

keras_hub/src/models/clip/clip_presets.py

@@ -0,0 +1,93 @@
+"""CLIP model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "clip_vit_base_patch16": {
+        "metadata": {
+            "description": (
+                "150 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 16, CLIP model."
+            ),
+            "params": 149620934,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_base_patch16/2",
+    },
+    "clip_vit_base_patch32": {
+        "metadata": {
+            "description": (
+                "151 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 32, CLIP model."
+            ),
+            "params": 151277363,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_base_patch32/2",
+    },
+    "clip_vit_large_patch14": {
+        "metadata": {
+            "description": (
+                "428 million parameter, 24-layer for vision and 12-layer for "
+                "text, patch size of 14, CLIP model."
+            ),
+            "params": 427616770,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_large_patch14/2",
+    },
+    "clip_vit_large_patch14_336": {
+        "metadata": {
+            "description": (
+                "428 million parameter, 24-layer for vision and 12-layer for "
+                "text, patch size of 14, image size of 336, CLIP model."
+            ),
+            "params": 427944770,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_large_patch14_336/2",
+    },
+    "clip_vit_b_32_laion2b_s34b_b79k": {
+        "metadata": {
+            "description": (
+                "151 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 32, Open CLIP model."
+            ),
+            "params": 151277363,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_b_32_laion2b_s34b_b79k/2",
+    },
+    "clip_vit_h_14_laion2b_s32b_b79k": {
+        "metadata": {
+            "description": (
+                "986 million parameter, 32-layer for vision and 24-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 986109698,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_h_14_laion2b_s32b_b79k/2",
+    },
+    "clip_vit_g_14_laion2b_s12b_b42k": {
+        "metadata": {
+            "description": (
+                "1.4 billion parameter, 40-layer for vision and 24-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 1366678530,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_g_14_laion2b_s12b_b42k/2",
+    },
+    "clip_vit_bigg_14_laion2b_39b_b160k": {
+        "metadata": {
+            "description": (
+                "2.5 billion parameter, 48-layer for vision and 32-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 2539567362,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_bigg_14_laion2b_39b_b160k/2",
+    },
+}

keras_hub/src/models/clip/clip_text_encoder.py

@@ -1,5 +1,6 @@
 from keras import layers
 
+from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
@@ -7,6 +8,7 @@ from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.clip.clip_encoder_block import CLIPEncoderBlock
 
 
+@keras_hub_export("keras_hub.models.CLIPTextEncoder")
 class CLIPTextEncoder(Backbone):
     """CLIP text core network with hyperparameters.
 
@@ -80,7 +82,7 @@ class CLIPTextEncoder(Backbone):
             for i in range(num_layers)
         ]
         self.layer_norm = layers.LayerNormalization(
-            epsilon=1e-6, dtype="float32", name=f"{prefix}layer_norm"
+            epsilon=1e-6, dtype=dtype, name=f"{prefix}layer_norm"
         )
 
         # === Functional Model ===
@@ -106,6 +108,7 @@ class CLIPTextEncoder(Backbone):
         super().__init__(
             inputs={"token_ids": token_id_input},
             outputs=outputs,
+            dtype=dtype,
             name=name,
             **kwargs,
         )

keras_hub/src/models/clip/clip_tokenizer.py

@@ -1,4 +1,5 @@
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.clip.clip_backbone import CLIPBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_hub.src.tokenizers.byte_pair_tokenizer import convert_to_ragged_batch
 from keras_hub.src.tokenizers.byte_pair_tokenizer import split_strings_for_bpe
@@ -39,11 +40,25 @@ class CLIPTokenizer(BytePairTokenizer):
             should have one merge rule per line. Every merge rule contains
             merge entities separated by a space.
         pad_with_end_token: bool. Whether to pad the output with `end_token`.
-    """
 
-    # TODO: Add example and `backbone_cls` once we have a CLIP model.
+    Examples:
+
+    ```python
+    # Unbatched input.
+    tokenizer = keras_hub.models.CLIPTokenizer.from_preset(
+        "clip_vit_base_patch32"
+    )
+    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 = None
+    backbone_cls = CLIPBackbone
 
     def __init__(
         self,