Add typing

Lightning-Universe · Feb 1, 2021 · 8d9775f · 8d9775f
1 parent 310f3ea
commit 8d9775f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/flash/text/seq2seq/core/data.py b/flash/text/seq2seq/core/data.py
@@ -144,13 +144,13 @@ def default_pipeline():
     @classmethod
     def from_files(
         cls,
-        train_file,
+        train_file: str,
         input: str = 'input',
         target: Optional[str] = None,
-        filetype="csv",
-        backbone="sshleifer/tiny-mbart",
-        valid_file=None,
-        test_file=None,
+        filetype: str = "csv",
+        backbone: str = "sshleifer/tiny-mbart",
+        valid_file: Optional[str] = None,
+        test_file: Optional[str] = None,
         max_source_length: int = 128,
         max_target_length: int = 128,
         padding: Union[str, bool] = 'max_length',
@@ -217,8 +217,8 @@ def from_file(
         predict_file: str,
         input: str = 'input',
         target: Optional[str] = None,
-        backbone="sshleifer/tiny-mbart",
-        filetype="csv",
+        backbone: str = "sshleifer/tiny-mbart",
+        filetype: str = "csv",
         max_source_length: int = 128,
         max_target_length: int = 128,
         padding: Union[str, bool] = 'max_length',

diff --git a/flash/text/seq2seq/summarization/data.py b/flash/text/seq2seq/summarization/data.py
@@ -29,13 +29,13 @@ def default_pipeline():
     @classmethod
     def from_files(
         cls,
-        train_file,
+        train_file: str,
         input: str = 'input',
         target: Optional[str] = None,
-        filetype="csv",
-        backbone="t5-small",
-        valid_file=None,
-        test_file=None,
+        filetype: str = "csv",
+        backbone: str = "t5-small",
+        valid_file: str = None,
+        test_file: str = None,
         max_source_length: int = 512,
         max_target_length: int = 128,
         padding: Union[str, bool] = 'max_length',
@@ -91,8 +91,8 @@ def from_file(
         predict_file: str,
         input: str = 'src_text',
         target: Optional[str] = None,
-        backbone="t5-small",
-        filetype="csv",
+        backbone: str = "t5-small",
+        filetype: str = "csv",
         max_source_length: int = 512,
         max_target_length: int = 128,
         padding: Union[str, bool] = 'longest',