Adjust read_file_fn pytype.

PiperOrigin-RevId: 600556341
google · Jan 23, 2024 · 674ec95 · 674ec95
1 parent 713c5f1
commit 674ec95
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/seqio/dataset_providers.py b/seqio/dataset_providers.py
@@ -583,7 +583,7 @@ class FileDataSource(DataSource):
 
   def __init__(
       self,
-      read_file_fn: Callable[[tf.data.Dataset], tf.data.Dataset],
+      read_file_fn: Callable[[str], tf.data.Dataset],
       split_to_filepattern: Mapping[str, Union[str, Iterable[str]]],
       num_input_examples: Optional[Mapping[str, int]] = None,
       caching_permitted: bool = True,