[Bugfix] Allow prefill of assistant response when using mistral_common (#9446)

2024-10-17 15:06:37 +00:00 · 2024-10-17 15:06:37 +00:00 · 5e443b594f
commit 5e443b594f
parent 9d30a056e7
1 changed files with 4 additions and 0 deletions
--- a/vllm/transformers_utils/tokenizers/mistral.py
+++ b/vllm/transformers_utils/tokenizers/mistral.py
@ -166,6 +166,10 @@ class MistralTokenizer:
                            tools: Optional[Dict[str, Any]] = None,
                            **kwargs) -> List[int]:

+        last_message = messages[-1]
+        if last_message["role"] == "assistant":
+            last_message["prefix"] = True
+
        request = ChatCompletionRequest(messages=messages,
                                        tools=tools)  # type: ignore[type-var]
        encoded = self.mistral.encode_chat_completion(request)