Stop asserting key order in bart preprocessor (#1221)

keras-team · Aug 21, 2023 · 22e21af · 22e21af
1 parent bab767a
commit 22e21af
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/keras_nlp/models/bart/bart_preprocessor.py b/keras_nlp/models/bart/bart_preprocessor.py
@@ -179,7 +179,7 @@ def get_config(self):
     def call(self, x, y=None, sample_weight=None):
         if not (
             isinstance(x, dict)
-            and ["encoder_text", "decoder_text"] == list(x.keys())
+            and all(k in x for k in ("encoder_text", "decoder_text"))
         ):
             raise ValueError(
                 '`x` must be a dictionary, containing the keys `"encoder_text"`'

diff --git a/keras_nlp/models/bart/bart_preprocessor_test.py b/keras_nlp/models/bart/bart_preprocessor_test.py
@@ -74,6 +74,22 @@ def test_tokenize_strings(self):
             output["decoder_padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 0]
         )
 
+    def test_key_order(self):
+        self.assertAllClose(
+            self.preprocessor(
+                {
+                    "encoder_text": " airplane at airport",
+                    "decoder_text": " kohli is the best",
+                }
+            ),
+            self.preprocessor(
+                {
+                    "decoder_text": " kohli is the best",
+                    "encoder_text": " airplane at airport",
+                }
+            ),
+        )
+
     def test_tokenize_list_of_strings(self):
         input_data = {
             "encoder_text": [" airplane at airport"] * 4,