update type hint

wemysschen · Nov 9, 2023 · f75a2c0 · f75a2c0
1 parent 8936898
commit f75a2c0
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/qianfan/resources/tools/tokenizer.py b/src/qianfan/resources/tools/tokenizer.py
@@ -17,7 +17,7 @@
 """
 
 import unicodedata
-from typing import Any
+from typing import Any, Literal
 
 from qianfan import get_config
 from qianfan.consts import Consts
@@ -33,7 +33,11 @@ class Tokenizer(object):
 
     @classmethod
     def count_tokens(
-        cls, text: str, mode: str = "local", model: str = "ERNIE-Bot", **kwargs: Any
+        cls,
+        text: str,
+        mode: Literal["local", "remote"] = "local",
+        model: str = "ERNIE-Bot",
+        **kwargs: Any,
     ) -> int:
         """
         Count the number of tokens in a given text.