From d6ac011e7ef8c9293bdc7d02f1a0200087dd7fee Mon Sep 17 00:00:00 2001
From: Michelle Casbon <michellecasbon@google.com>
Date: Tue, 12 Dec 2023 23:15:49 -0800
Subject: [PATCH] Internal change.

PiperOrigin-RevId: 590465163
---
 seqio/experimental.py      | 2 ++
 seqio/experimental_test.py | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/seqio/experimental.py b/seqio/experimental.py
index 1963e1c1..77f3907e 100644
--- a/seqio/experimental.py
+++ b/seqio/experimental.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Experimental utilities for SeqIO."""
+
 import functools
 import inspect
 from typing import Callable, Iterable, Mapping, Optional, Sequence
@@ -86,6 +87,7 @@ def _no_op_mixture_registry_get(*args, **kwargs):
 
 
 def disable_registry():
+  """Disables the seqio TaskRegistry and MixtureRegistry."""
   _enfore_empty_registries()
   dataset_providers.TaskRegistry.add = _no_op_task_registry_add
   dataset_providers.TaskRegistry.add_provider = _no_op_task_registry_add
diff --git a/seqio/experimental_test.py b/seqio/experimental_test.py
index a17e029c..2f08521e 100644
--- a/seqio/experimental_test.py
+++ b/seqio/experimental_test.py
@@ -15,6 +15,8 @@
 """Tests for seqio.preprocessors."""
 
 import contextlib
+from unittest import mock
+
 from absl.testing import absltest
 from seqio import dataset_providers
 from seqio import experimental
@@ -984,5 +986,6 @@ def test_mixture_registry_get_error(self):
         MixtureRegistry.get('dummy_mixture')
 
 
+
 if __name__ == '__main__':
   absltest.main()