Improve

mistralai · Jul 20, 2024 · 64d0d70 · 64d0d70
1 parent 58230b4
commit 64d0d70
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/mistral_inference/main.py b/src/mistral_inference/main.py
@@ -12,6 +12,7 @@
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 from mistral_common.tokens.tokenizers.base import Tokenizer
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.tokens.tokenizers.tekken import Tekkenizer, SpecialTokenPolicy
 from mistral_common.tokens.tokenizers.sentencepiece import is_sentencepiece
 from mistral_common.tokens.tokenizers.tekken import is_tekken
 
@@ -36,6 +37,9 @@ def load_tokenizer(model_path: Path) -> MistralTokenizer:
 
     mistral_tokenizer = MistralTokenizer.from_file(str(model_path / tokenizer[0]))
 
+    if isinstance(mistral_tokenizer.instruct_tokenizer.tokenizer, Tekkenizer):
+        mistral_tokenizer.instruct_tokenizer.tokenizer.special_token_policy = SpecialTokenPolicy.KEEP
+
     logging.info(f"Loaded tokenizer of type {mistral_tokenizer.instruct_tokenizer.__class__}")
 
     return mistral_tokenizer