type annotation for BPE.decode fixed (VKCOM#83)

Co-authored-by: d.gavrilov <[email protected]>
vitalymegabyte · Jan 28, 2021 · f4bc5f8 · f4bc5f8
1 parent 9c758c1
commit f4bc5f8
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/youtokentome/youtokentome.py b/youtokentome/youtokentome.py
@@ -1,8 +1,7 @@
+import _youtokentome_cython
 from enum import Enum
 from typing import List, Union, Optional, Collection
 
-import _youtokentome_cython
-
 
 class OutputType(Enum):
     ID = 1
@@ -82,8 +81,10 @@ def id_to_subword(self, id: int) -> str:
         return self.bpe_cython.id_to_subword(id)
 
     def decode(
-        self, ids: List[int], ignore_ids: Optional[Collection[int]] = None
-    ) -> str:
+        self,
+        ids: Union[List[int], List[List[int]]],
+        ignore_ids: Optional[Collection[int]] = None,
+    ) -> List[str]:
         return self.bpe_cython.decode(ids, ignore_ids)
 
     def __getstate__(self):