Merge pull request #199 from mistralai/improve_mistral_inference

Improve mistral inference

pyproject.toml

 [tool.poetry]
 name = "mistral_inference"
-version = "1.3.0"
+version = "1.3.1"
 description = ""
 authors = ["bam4d <bam4d@mistral.ai>"]
 readme = "README.md"

src/mistral_inference/__init__.py

-__version__ = "1.3.0"
+__version__ = "1.3.1"

src/mistral_inference/main.py

@@ -12,6 +12,7 @@ from mistral_common.protocol.instruct.messages import AssistantMessage, UserMess
 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