Adjust read_file_fn pytype.

zcharles8 · SeqIO · commit ff4cc65ce4e2 · 2024-01-22T19:45:05.000-08:00
PiperOrigin-RevId: 600556341
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,
