A more generic packing op for seqio

PiperOrigin-RevId: 373368806
craffel · May 12, 2021 · 62a20f2 · 62a20f2
1 parent 2953766
commit 62a20f2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/seqio/feature_converters.py b/seqio/feature_converters.py
@@ -477,10 +477,11 @@ def __call__(self, ds: tf.data.Dataset,
     if self.pack:
       for k, v in expected_features.items():
         # Packing requires rank 1.
-        if v.rank != 1:
+        if v.rank != 1 and not self._use_custom_packing_ops:
           raise ValueError(
-              f"When packing is enabled, expected ranks must be 1. Got "
-              f"expected rank {v.rank} for feature {k}.")
+              f"When packing is enabled, expected ranks must be 1 or "
+              f"use_custom_packing_ops must be set. Got expected rank {v.rank} "
+              f"for feature {k}.")
       for k, v in self.PACKING_FEATURE_DTYPES.items():
         expected_features[k] = FeatureConverter.FeatureSpec(rank=1, dtype=v)
 

diff --git a/seqio/feature_converters_test.py b/seqio/feature_converters_test.py
@@ -331,8 +331,9 @@ def test_validate_dataset_rank_2_with_pack(self):
           "inputs": FeatureSpec(tf.int64, rank=2),
           "targets": FeatureSpec(tf.int64)
       }
-      expected_msg = ("When packing is enabled, expected ranks must be 1. Got "
-                      "expected rank 2 for feature inputs.")
+      expected_msg = ("When packing is enabled, expected ranks must be 1 or "
+                      "use_custom_packing_ops must be set. Got expected rank 2 "
+                      "for feature inputs.")
       with self.assertRaisesRegex(ValueError, expected_msg):
         converter(ds, task_feature_lengths)