Update app.py

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 # import spaces
 import torch
 from transformers import AutoTokenizer
-from vllm import LLM, SamplingParams
+from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -19,7 +19,9 @@ 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"
-model = LLM(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.use_default_system_prompt = False
 
@@ -51,7 +53,7 @@ async def generate(
         repetition_penalty=repetition_penalty,
     )
 
-    stream = await model.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