quick fix middle token

mistralai · Dec 20, 2024 · 8fd8d2b · 8fd8d2b
1 parent 5cac5e1
commit 8fd8d2b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/mistral_common/tokens/tokenizers/sentencepiece.py b/src/mistral_common/tokens/tokenizers/sentencepiece.py
@@ -332,6 +332,7 @@ def __init__(self, tokenizer: Tokenizer, mm_encoder: Optional[MultiModalEncoder]
         self.BOS = self.tokenizer.get_control_token(SpecialTokens.bos.value)
         self.PREFIX = self.tokenizer.get_control_token(SpecialTokens.prefix.value)
         self.SUFFIX = self.tokenizer.get_control_token(SpecialTokens.suffix.value)
+        self.MIDDLE = self.tokenizer.get_control_token(SpecialTokens.middle.value)
 
     def encode_user_message(
         self,
@@ -453,6 +454,7 @@ def encode_fim(self, request: FIMRequest) -> Tokenized:
             *suffix_tokens,
             self.PREFIX,
             *prefix_tokens,
+            self.MIDDLE,
         ]
         return Tokenized(tokens=tokens, text=self.tokenizer.to_string(tokens))