Add parameters types

RasaHQ · Dec 10, 2019 · c6bd407 · c6bd407
1 parent 8c09c4d
commit c6bd407
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/rasa/nlu/training_data/training_data.py b/rasa/nlu/training_data/training_data.py
@@ -343,7 +343,7 @@ def validate(self) -> None:
                 )
 
     def train_test_split(
-        self, train_frac: float = 0.8, random_seed: int = None
+        self, train_frac: float = 0.8, random_seed: Optional[int] = None
     ) -> Tuple["TrainingData", "TrainingData"]:
         """Split into a training and test dataset,
         preserving the fraction of examples per intent."""
@@ -395,7 +395,9 @@ def build_nlg_stories_from_examples(examples) -> Dict[Text, list]:
                 ]
         return nlg_stories
 
-    def split_nlu_examples(self, train_frac, random_seed=None):
+    def split_nlu_examples(
+        self, train_frac: float, random_seed: Optional[int] = None
+    ) -> Tuple[list, list]:
         train, test = [], []
         for intent, count in self.examples_per_intent.items():
             ex = [e for e in self.intent_examples if e.data["intent"] == intent]