From a238c9758583ccaeccdcbc217dddee75651cf26e Mon Sep 17 00:00:00 2001
From: bofeng huang <bofenghuang7@gmail.com>
Date: Sun, 5 May 2024 03:08:48 +0200
Subject: [PATCH] Fix vLLM template (#734)

* Update vllm_mixtral.py

* Fix template

* Fix template
---
 06_gpu_and_ml/llm-serving/vllm_gemma.py     | 2 +-
 06_gpu_and_ml/llm-serving/vllm_inference.py | 4 ++--
 06_gpu_and_ml/llm-serving/vllm_mixtral.py   | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/06_gpu_and_ml/llm-serving/vllm_gemma.py b/06_gpu_and_ml/llm-serving/vllm_gemma.py
index 2a3545961..634c6d47a 100644
--- a/06_gpu_and_ml/llm-serving/vllm_gemma.py
+++ b/06_gpu_and_ml/llm-serving/vllm_gemma.py
@@ -121,7 +121,7 @@ class Model:
     @modal.enter()
     def load(self):
         self.template = (
-            "start_of_turn>user\n{user}<end_of_turn>\n<start_of_turn>model"
+            "<start_of_turn>user\n{user}<end_of_turn>\n<start_of_turn>model\n"
         )
 
         # Load the model. Tip: Some models, like MPT, may require `trust_remote_code=true`.
diff --git a/06_gpu_and_ml/llm-serving/vllm_inference.py b/06_gpu_and_ml/llm-serving/vllm_inference.py
index 3f67aa908..c24e345db 100644
--- a/06_gpu_and_ml/llm-serving/vllm_inference.py
+++ b/06_gpu_and_ml/llm-serving/vllm_inference.py
@@ -109,11 +109,11 @@ class Model:
     def load_model(self):
         # Tip: models that are not fully implemented by Hugging Face may require `trust_remote_code=true`.
         self.llm = vllm.LLM(MODEL_DIR, tensor_parallel_size=GPU_CONFIG.count)
-        self.template = """<s>[INST] <<SYS>>
+        self.template = """[INST] <<SYS>>
 {system}
 <</SYS>>
 
-{user} [/INST] """
+{user} [/INST]"""
 
     @modal.method()
     def generate(self, user_questions):
diff --git a/06_gpu_and_ml/llm-serving/vllm_mixtral.py b/06_gpu_and_ml/llm-serving/vllm_mixtral.py
index 57618ae28..eb236b9cb 100644
--- a/06_gpu_and_ml/llm-serving/vllm_mixtral.py
+++ b/06_gpu_and_ml/llm-serving/vllm_mixtral.py
@@ -121,7 +121,7 @@ def start_engine(self):
             disable_log_stats=True,  # disable logging so we can stream tokens
             disable_log_requests=True,
         )
-        self.template = "<s> [INST] {user} [/INST] "
+        self.template = "[INST] {user} [/INST]"
 
         # this can take some time!
         self.engine = AsyncLLMEngine.from_engine_args(engine_args)