From 443df38cc50ff4af99cbeb82b54460212ed7ca8c Mon Sep 17 00:00:00 2001
From: Zachary Charles <zachcharles@google.com>
Date: Mon, 22 Jan 2024 12:59:41 -0800
Subject: [PATCH] Adjust `read_file_fn` pytype.

PiperOrigin-RevId: 600545176
---
 seqio/dataset_providers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/seqio/dataset_providers.py b/seqio/dataset_providers.py
index 4ad6383a..a64bd382 100644
--- a/seqio/dataset_providers.py
+++ b/seqio/dataset_providers.py
@@ -583,7 +583,7 @@ class FileDataSource(DataSource):
 
   def __init__(
       self,
-      read_file_fn: Callable[[str], tf.data.Dataset],
+      read_file_fn: Callable[[tf.data.Dataset], tf.data.Dataset],
       split_to_filepattern: Mapping[str, Union[str, Iterable[str]]],
       num_input_examples: Optional[Mapping[str, int]] = None,
       caching_permitted: bool = True,