keras-hub-nightly 0.16.1.dev202410150342__tar.gz → 0.16.1.dev202410170342__tar.gz

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-hub-nightly
-Version: 0.16.1.dev202410150342
+Version: 0.16.1.dev202410170342
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/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`.
 
@@ -309,6 +310,9 @@ class CausalLM(Task):
                 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.
@@ -339,6 +343,33 @@ 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)
 
@@ -347,7 +378,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_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/deeplab_v3/deeplab_v3_presets.py

@@ -13,6 +13,6 @@ backbone_presets = {
             "path": "deeplabv3",
             "model_card": "https://arxiv.org/abs/1802.02611",
         },
-        "kaggle_handle": "kaggle://keras/deeplabv3/keras/deeplabv3_plus_resnet50_pascalvoc/3",
+        "kaggle_handle": "kaggle://keras/deeplabv3/keras/deeplab_v3_plus_resnet50_pascalvoc/3",
     },
 }

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/gemma/gemma_backbone.py

@@ -224,7 +224,7 @@ class GemmaBackbone(Backbone):
 
         Example:
         ```
-        # Feel free to change the mesh shape to balance data and model parallel
+        # Feel free to change the mesh shape to balance data and model parallelism
         mesh = keras.distribution.DeviceMesh(
             shape=(1, 8), axis_names=('batch', 'model'),
             devices=keras.distribution.list_devices())
@@ -232,11 +232,19 @@ class GemmaBackbone(Backbone):
             mesh, model_parallel_dim_name="model")
 
         distribution = keras.distribution.ModelParallel(
-            mesh, layout_map, batch_dim_name='batch')
+            layout_map=layout_map, batch_dim_name='batch')
         with distribution.scope():
            gemma_model = keras_hub.models.GemmaCausalLM.from_preset()
         ```
 
+        To see how the layout map was applied, load the model then run (for one decoder block):
+        ```
+        embedding_layer = gemma_model.backbone.get_layer("token_embedding")
+        decoder_block_1 = gemma_model.backbone.get_layer('decoder_block_1')
+        for variable in embedding_layer.weights + decoder_block_1.weights:
+            print(f'{variable.path:<58}  {str(variable.shape):<16}  {str(variable.value.sharding.spec)}')
+        ```
+
         Args:
             device_mesh: The `keras.distribution.DeviceMesh` instance for
                 distribution.
@@ -246,7 +254,7 @@ class GemmaBackbone(Backbone):
                 the data should be partition on.
         Return:
             `keras.distribution.LayoutMap` that contains the sharding spec
-            of all the model weights.
+            for all the model weights.
         """
         # The weight path and shape of the Gemma backbone is like below (for 2G)
         # token_embedding/embeddings,  (256128, 2048), 524550144

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/llama/llama_backbone.py

@@ -175,3 +175,121 @@ class LlamaBackbone(Backbone):
             }
         )
         return config
+
+    @staticmethod
+    def get_layout_map(
+        device_mesh,
+        model_parallel_dim_name="model",
+        data_parallel_dim_name="batch",
+    ):
+        """Get a `keras.distribution.LayoutMap` for model parallel distribution.
+
+        The returned `LayoutMap` contains the sharding spec for the Llama
+        backbone weights, so that you can use it to distribute weights across
+        the accelerators.
+
+        Example:
+        ```
+        # Feel free to change the mesh shape to balance data and model parallelism
+        mesh = keras.distribution.DeviceMesh(
+            shape=(1, 8),
+            axis_names=('batch', 'model'),
+            devices=keras.distribution.list_devices(),
+        )
+        layout_map = LlamaBackbone.get_layout_map(
+            mesh,
+            model_parallel_dim_name="model",
+        )
+
+        distribution = keras.distribution.ModelParallel(
+            layout_map=layout_map,
+            batch_dim_name='batch',
+        )
+
+        with distribution.scope():
+           llama_model = keras_hub.models.LlamaCausalLM.from_preset()
+        ```
+
+        To see how the layout map was applied, load the model then run (for one decoder block):
+        ```
+        embedding_layer = llama_model.backbone.get_layer("token_embedding")
+        decoder_block_1 = llama_model.backbone.get_layer('transformer_layer_0')
+        for variable in embedding_layer.weights + decoder_block_1.weights:
+            print(f'{variable.path:<58}  {str(variable.shape):<16}  {str(variable.value.sharding.spec)}')
+        ```
+
+        Args:
+            device_mesh: The `keras.distribution.DeviceMesh` instance for
+                distribution.
+            model_parallel_dim_name: The axis name of the device mesh, where
+                the weights should be partition on.
+            data_parallel_dim_name: The axis name of the device mesh, where
+                the data should be partition on.
+        Return:
+            `keras.distribution.LayoutMap` that contains the sharding spec
+            for all the model weights.
+        """
+        # The weight path and shape of the Llama backbone is like below
+        # token_embedding/embeddings                                  (128256, 2048)
+        # repeat block for decoder
+        # transformer_layer_0/self_attention/query/kernel             (2048, 32, 64)
+        # transformer_layer_0/self_attention/key/kernel               (2048, 8, 64)
+        # transformer_layer_0/self_attention/value/kernel             (2048, 8, 64)
+        # transformer_layer_0/self_attention/attention_output/kernel  (32, 64, 2048)
+        # transformer_layer_0/self_attention_layernorm/scale          (2048,)
+        # transformer_layer_0/feedforward_intermediate_dense/kernel   (2048, 8192)
+        # transformer_layer_0/feedforward_gate_dense/kernel           (2048, 8192)
+        # transformer_layer_0/feedforward_output_dense/kernel         (8192, 2048)
+        # transformer_layer_0/feedforward_layernorm/scale             (2048,)
+
+        if not isinstance(device_mesh, keras.distribution.DeviceMesh):
+            raise ValueError(
+                "Invalid device_mesh type. Expected `keras.distribution.Device`,"
+                f" got {type(device_mesh)}"
+            )
+        if model_parallel_dim_name not in device_mesh.axis_names:
+            raise ValueError(
+                f"{model_parallel_dim_name} is not found in the "
+                f"device_mesh.axis_names. {device_mesh.axis_name=}"
+            )
+        if data_parallel_dim_name not in device_mesh.axis_names:
+            raise ValueError(
+                f"{data_parallel_dim_name} is not found in the "
+                f"device_mesh.axis_names. {device_mesh.axis_name=}"
+            )
+        # Note that it is possible to further config the mesh to be 3D, eg
+        # (data, seq, model). We leave it as 2D for now for simplicity.
+        data_dim = data_parallel_dim_name
+        model_dim = model_parallel_dim_name
+        # The sharding config is based on the Gemma team training config.
+        # See https://arxiv.org/abs/2403.08295
+        layout_map = keras.distribution.LayoutMap(device_mesh)
+        layout_map["token_embedding/embeddings"] = (model_dim, data_dim)
+        layout_map[
+            "transformer_layer.*self_attention.*(query|key|value).kernel"
+        ] = (
+            model_dim,
+            data_dim,
+            None,
+        )
+        layout_map["transformer_layer.*attention_output.kernel"] = (
+            model_dim,
+            None,
+            data_dim,
+        )
+        layout_map[
+            "transformer_layer.*feedforward_intermediate_dense.kernel"
+        ] = (
+            data_dim,
+            model_dim,
+        )
+        layout_map["transformer_layer.*feedforward_gate_dense.kernel"] = (
+            data_dim,
+            model_dim,
+        )
+        layout_map["transformer_layer.*feedforward_output_dense.kernel"] = (
+            model_dim,
+            data_dim,
+        )
+
+        return layout_map

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/llama/llama_causal_lm.py

@@ -42,7 +42,9 @@ class LlamaCausalLM(CausalLM):
         self.preprocessor = preprocessor
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
         super().__init__(

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/mistral/mistral_causal_lm.py

@@ -42,7 +42,9 @@ class MistralCausalLM(CausalLM):
         self.preprocessor = preprocessor
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
         super().__init__(

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/pali_gemma/pali_gemma_causal_lm.py

@@ -110,7 +110,9 @@ class PaliGemmaCausalLM(CausalLM):
         self.backbone = backbone
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_state = backbone(inputs=inputs)
         outputs = backbone.token_embedding(hidden_state, reverse=True)
         outputs = outputs[:, backbone.image_sequence_length :, :]

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/models/phi3/phi3_causal_lm.py

@@ -41,7 +41,9 @@ class Phi3CausalLM(CausalLM):
         self.preprocessor = preprocessor
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
         super().__init__(

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/tests/test_case.py

@@ -569,6 +569,15 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         ds = tf.data.Dataset.from_tensor_slices(train_data).batch(batch_size)
         x, y, sw = keras.utils.unpack_x_y_sample_weight(train_data)
 
+        # Test: the tree struct output by the
+        # preprocessor must match what model expects.
+        preprocessed_data = preprocessor(*train_data)[0]
+        tree.assert_same_structure(
+            preprocessed_data,
+            task._inputs_struct,
+            check_types=False,
+        )
+
         # Test predict.
         output = task.predict(x)
         if expected_output_shape is not None:

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.16.1.dev202410150342"
+__version__ = "0.16.1.dev202410170342"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/keras_hub_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-hub-nightly
-Version: 0.16.1.dev202410150342
+Version: 0.16.1.dev202410170342
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.16.1.dev202410150342 → keras_hub_nightly-0.16.1.dev202410170342}/setup.py

@@ -23,7 +23,7 @@ def get_version(rel_path):
 
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
-VERSION = "0.16.1.dev202410150342"  # get_version("keras_hub/src/version_utils.py")
+VERSION = "0.16.1.dev202410170342"  # get_version("keras_hub/src/version_utils.py")
 
 setup(
     name="keras-hub-nightly",  # "keras-hub",