Update adapt_tokenizer.py

adapt_tokenizer.py

@@ -1,9 +1,8 @@
-from typing import Union
-from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
-Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+from typing import Any
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
 NUM_SENTINEL_TOKENS: int = 100
 
-def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
+def adapt_tokenizer_for_denoising(tokenizer: PreTrainedTokenizerBase) -> None:
     """Adds sentinel tokens and padding token (if missing).
 
     Expands the tokenizer vocabulary to include sentinel tokens
@@ -34,7 +33,7 @@ class AutoTokenizerForMOD(AutoTokenizer):
     """
 
     @classmethod
-    def from_pretrained(cls, *args, **kwargs):
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedTokenizerBase:
         """See `AutoTokenizer.from_pretrained` docstring."""
         tokenizer = super().from_pretrained(*args, **kwargs)
         adapt_tokenizer_for_denoising(tokenizer)