update type annotation per PR feedback

axolotl-ai-cloud · Sep 19, 2023 · 7571db3 · 7571db3
1 parent e136a7b
commit 7571db3
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py
@@ -66,11 +66,13 @@ def _get_assistant_token(self):
             pass
         return False
 
-    def _tokenize(self, prompt: str, add_eos_token=True, strip_bos_token=False):
-        result: Union[Dict[str, List[Union[bool, int]]], BatchEncoding]
+    def _tokenize(
+        self, prompt: str, add_eos_token: bool = True, strip_bos_token: bool = False
+    ) -> BatchEncoding:
+        result: BatchEncoding
         if not prompt.strip():
             LOG.warning("Empty text requested for tokenization.")
-            result = {"input_ids": [], "attention_mask": []}
+            result = BatchEncoding(data={"input_ids": [], "attention_mask": []})
         else:
             result = self.tokenizer(
                 prompt,