From f4bc5f8c45487b31aa0f53419714f32229ca8854 Mon Sep 17 00:00:00 2001 From: Daniil Gavrilov Date: Thu, 28 Jan 2021 22:04:09 +0300 Subject: [PATCH] type annotation for BPE.decode fixed (#83) Co-authored-by: d.gavrilov --- youtokentome/youtokentome.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/youtokentome/youtokentome.py b/youtokentome/youtokentome.py index b728141..593febf 100644 --- 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):