Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import uuid
 
 import gradio as gr
-# import spaces
+
 import torch
 from transformers import AutoTokenizer
 from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
@@ -10,22 +10,27 @@ from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+MODEL_ID = "neuralmagic/OpenHermes-2.5-Mistral-7B-pruned50"
 
-DESCRIPTION = """\
-# NM vLLM Hermes Mistral Chat
+DESCRIPTION = f"""\
+# NM vLLM Chat
+Model: {MODEL_ID}
 """
 
 if not torch.cuda.is_available():
     raise ValueError("Running on CPU 🥶 This demo does not work on CPU.")
 
-model_id = "neuralmagic/OpenHermes-2.5-Mistral-7B-pruned50"
-engine_args = AsyncEngineArgs(model=model_id, sparsity="sparse_w16a16", max_model_len=MAX_INPUT_TOKEN_LENGTH)
+engine_args = AsyncEngineArgs(
+    model=MODEL_ID, 
+    sparsity="sparse_w16a16", 
+    max_model_len=MAX_INPUT_TOKEN_LENGTH
+)
 engine = AsyncLLMEngine.from_engine_args(engine_args)
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 tokenizer.use_default_system_prompt = False
 
-# @spaces.GPU
+
 async def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -37,13 +42,22 @@ async def generate(
     repetition_penalty: float = 1.2,
 ):
     conversation = []
+
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
+
     for user, assistant in chat_history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+        conversation.extend(
+            [
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant},
+            ]
+        )
     conversation.append({"role": "user", "content": message})
 
-    formatted_conversation = tokenizer.apply_chat_template(conversation, tokenize=False)
+    formatted_conversation = tokenizer.apply_chat_template(
+        conversation, tokenize=False, add_generation_prompt=True
+    )
 
     sampling_params = SamplingParams(
         max_tokens=max_new_tokens,
@@ -53,8 +67,10 @@ async def generate(
         repetition_penalty=repetition_penalty,
     )
 
-    stream = await engine.add_request(uuid.uuid4().hex, formatted_conversation, sampling_params)
-    
+    stream = await engine.add_request(
+        uuid.uuid4().hex, formatted_conversation, sampling_params
+    )
+
     async for request_output in stream:
         text = request_output.outputs[0].text
         yield text
@@ -112,8 +128,10 @@ chat_interface = gr.ChatInterface(
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+    gr.DuplicateButton(
+        value="Duplicate Space for private use", elem_id="duplicate-button"
+    )
     chat_interface.render()
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()