tf-models-nightly 2.20.0.dev20250806__py2.py3-none-any.whl → 2.20.0.dev20250808__py2.py3-none-any.whl

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -274,6 +274,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
       input_tensor_shape = input_shape
     elif isinstance(input_shape, (list, tuple)):
       input_tensor_shape = tf.TensorShape(input_shape[0])
+    elif isinstance(input_shape, dict):
+      input_tensor_shape = tf.TensorShape(input_shape["input_tensor"])
     else:
       raise ValueError(
           "The type of input shape argument is not supported, got: %s" %
@@ -546,11 +548,13 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     """Transformer self-attention encoder block call.
 
     Args:
-      inputs: a single tensor or a list of tensors. `input tensor` as the single
-        sequence of embeddings. [`input tensor`, `attention mask`] to have the
-        additional attention mask. [`query tensor`, `key value tensor`,
-        `attention mask`] to have separate input streams for the query, and
-        key/value to the multi-head attention.
+      inputs: a single tensor or a list of tensors, or a dictionary. `input
+        tensor` as the single sequence of embeddings. [`input tensor`,
+        `attention mask`] to have the additional attention mask. [`query
+        tensor`, `key value tensor`, `attention mask`] to have separate input
+        streams for the query, and key/value to the multi-head attention. If
+        dictionary is provided, it must contain the following keys:
+        `input_tensor`, `attention_mask`, `key_value_tensor`.
       output_range: the sequence output range, [0, output_range) for slicing the
         target sequence. `None` means the target sequence is not sliced. If you
         would like to have no change to the model training, it is better to only
@@ -568,6 +572,21 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
       else:
         raise ValueError("Unexpected inputs to %s with length at %d" %
                          (self.__class__, len(inputs)))
+    elif isinstance(inputs, dict):
+      if not set(inputs.keys()).issubset(
+          set(["input_tensor", "key_value_tensor", "attention_mask"])
+      ):
+        raise ValueError(
+            f"Unexpected keys in input dictionary to: {inputs.keys()}"
+        )
+      try:
+        input_tensor = inputs["input_tensor"]
+      except KeyError as e:
+        raise ValueError(
+            "Missing required key `input_tensor` in input dictionary."
+        ) from e
+      key_value = inputs.get("key_value_tensor", None)
+      attention_mask = inputs.get("attention_mask", None)
     else:
       input_tensor, key_value, attention_mask = (inputs, None, None)
 

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -55,6 +55,25 @@ class TransformerEncoderBlockLayerTest(
     # The default output of a transformer layer should be the same as the input.
     self.assertEqual(data_tensor.shape.as_list(), output_tensor.shape.as_list())
 
+  def test_layer_creation_with_dict_inputs(self, transformer_cls):
+    test_layer = transformer_cls(
+        num_attention_heads=10, inner_dim=2048, inner_activation='relu'
+    )
+    sequence_length = 21
+    width = 80
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    # Create a 2-dimensional input (the first dimension is implicit).
+    mask_tensor = tf_keras.Input(shape=(sequence_length, sequence_length))
+    inputs = {
+        'input_tensor': data_tensor,
+        'key_value_tensor': data_tensor,
+        'attention_mask': mask_tensor,
+    }
+    output_tensor = test_layer(inputs)
+    # The default output of a transformer layer should be the same as the input.
+    self.assertEqual(data_tensor.shape.as_list(), output_tensor.shape.as_list())
+
   def test_layer_invocation(self, transformer_cls):
     test_layer = transformer_cls(
         num_attention_heads=10, inner_dim=2048, inner_activation='relu')
@@ -88,6 +107,40 @@ class TransformerEncoderBlockLayerTest(
     # Create a model from the test layer.
     model = tf_keras.Model([data_tensor, mask_tensor], output_tensor)
 
+    # Invoke the model on test data. We can't validate the output data itself
+    # (the NN is too complex) but this will rule out structural runtime errors.
+    batch_size = 6
+    input_data = 10 * np.random.random_sample(
+        (batch_size, sequence_length, width)
+    )
+    # The attention mask should be of shape (batch, from_seq_len, to_seq_len),
+    # which here is (batch, sequence_length, sequence_length)
+    mask_data = np.random.randint(
+        2, size=(batch_size, sequence_length, sequence_length)
+    )
+    _ = model.predict([input_data, mask_data])
+
+  def test_layer_invocation_with_dict_inputs(self, transformer_cls):
+    test_layer = transformer_cls(
+        num_attention_heads=10, inner_dim=2048, inner_activation='relu'
+    )
+    sequence_length = 21
+    width = 80
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    # Create a 2-dimensional input (the first dimension is implicit).
+    mask_tensor = tf_keras.Input(shape=(sequence_length, sequence_length))
+    inputs = {
+        'input_tensor': data_tensor,
+        'key_value_tensor': data_tensor,
+        'attention_mask': mask_tensor,
+    }
+
+    output_tensor = test_layer(inputs)
+
+    # Create a model from the test layer.
+    model = tf_keras.Model([data_tensor, mask_tensor], output_tensor)
+
     # Invoke the model on test data. We can't validate the output data itself
     # (the NN is too complex) but this will rule out structural runtime errors.
     batch_size = 6

{tf_models_nightly-2.20.0.dev20250806.dist-info → tf_models_nightly-2.20.0.dev20250808.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.20.0.dev20250806
+Version: 2.20.0.dev20250808
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.20.0.dev20250806.dist-info → tf_models_nightly-2.20.0.dev20250808.dist-info}/RECORD

@@ -363,8 +363,8 @@ official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=QWq1dJqQUPe5n69K3
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=omzTkCBEk2TOkHEYDEBwve6WsOitX7IIJHzeKXdqDq0,11012
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=pSCONEZRI4J9_6QLTJ3g_ynUYLrRXsJ1c2YMSiOV_60,8893
 official/nlp/modeling/layers/transformer.py,sha256=VjUO-gVj_PnavbT_vSrg5NDKMr0SRSiqSg5ktd42m5M,20087
-official/nlp/modeling/layers/transformer_encoder_block.py,sha256=E-WeoxsjByL-lkPAyEjDbvt1_3ghcIpCXnoLSeCDKFQ,27953
-official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=Y1Byz7RgF6puMIU3WRyoWghp5a3rYemgTRFRDohfy2Q,35402
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=5GJgtK1mdTxMDYVWfUoBAI_GvjDL0zO9AWtKCovSZiU,28789
+official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=7yBgv1UNmfOFre6txF_Rq93RLc1TJwnJ7-Dz4p55sy4,37602
 official/nlp/modeling/layers/transformer_scaffold.py,sha256=qmzhCJvbbFVF9zDqnfO4Zs2JDXwKhK7iEBOhsU6-KpQ,15704
 official/nlp/modeling/layers/transformer_scaffold_test.py,sha256=dRJwesTBKm-mF5mDHrHfVpVNnxa-Wx-fj_4ZHDPTpE0,19920
 official/nlp/modeling/layers/transformer_test.py,sha256=-pk9cdz9UlMpCIkGRkCKsMmjdRGi0seySaaB_2dwmXw,5522
@@ -1248,9 +1248,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=yiAneltAW3NHSj3fUSvHNBjfq0MGZ
 tensorflow_models/nlp/__init__.py,sha256=8uQd4wI6Zc4IJMPjtQifMeWVbPFkTxqYh66wfivCOL4,807
 tensorflow_models/uplift/__init__.py,sha256=NzaweFf4ZmhRb2l_fuV6bP-2N8oSO3xu6xJqVb1UmpY,999
 tensorflow_models/vision/__init__.py,sha256=ks420Ooqzi0hU7HnQpM5rylLaE-YcJdJkBx_umVaXlE,833
-tf_models_nightly-2.20.0.dev20250806.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.20.0.dev20250806.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.20.0.dev20250806.dist-info/METADATA,sha256=exENzp9z9a-9fvO8_TricogELR_dG3RhNBpsjWys0zM,1432
-tf_models_nightly-2.20.0.dev20250806.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.20.0.dev20250806.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.20.0.dev20250806.dist-info/RECORD,,
+tf_models_nightly-2.20.0.dev20250808.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.20.0.dev20250808.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.20.0.dev20250808.dist-info/METADATA,sha256=w4VBtg2nGIUWJ7FrIO7pKf_mUesG5m7vu1ZCyARsVes,1432
+tf_models_nightly-2.20.0.dev20250808.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.20.0.dev20250808.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.20.0.dev20250808.dist-info/RECORD,,