Allow features to be immutable

PiperOrigin-RevId: 502624945

setup.py

@@ -46,6 +46,7 @@
         'absl-py',
         'babel',
         'editdistance',
+        'immutabledict',
         'gin-config',
         'mesh-tensorflow[transformer]>=0.1.13',
         'nltk',

t5/data/dataset_providers.py

@@ -22,6 +22,7 @@
 between.
 """
 
+from collections.abc import Mapping
 import re
 
 import seqio
@@ -88,7 +89,7 @@ def __init__(self,
       raise ValueError("output_features must be non-empty.")
     if output_features is None:
       output_features = seqio.Feature(utils.get_default_vocabulary())
-    if isinstance(output_features, dict):
+    if isinstance(output_features, Mapping):
       pass
     elif isinstance(output_features, seqio.Feature):
       output_features = {k: output_features for k in _DEFAULT_FEATURE_KEYS}

t5/data/dataset_providers_test.py

@@ -16,6 +16,7 @@
 import os
 
 from absl.testing import absltest
+import immutabledict
 import seqio
 from seqio import test_utils
 from t5.data import dataset_providers
@@ -53,6 +54,15 @@ def test_tfds_task(self):
         "t5_tfds_task", dataset_providers.TfdsTask, tfds_name="fake:0.0.0")
     self.verify_task_matches_fake_datasets("t5_tfds_task", use_cached=False)
 
+  def test_immutabledict_features(self):
+    _add_t5_task(
+        "t5_tfds_task", dataset_providers.TfdsTask, tfds_name="fake:0.0.0",
+        output_features=immutabledict.immutabledict({
+            "inputs": seqio.Feature(test_utils.sentencepiece_vocab()),
+            "targets": seqio.Feature(test_utils.sentencepiece_vocab())
+        }))
+    self.verify_task_matches_fake_datasets("t5_tfds_task", use_cached=False)
+
   def test_function_task(self):
     _add_t5_task(
         "t5_fn_task",