Update app.py

app.py

@@ -1,23 +1,28 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from vllm import LLM
+from vllm.sampling_params import SamplingParams
+import spaces
+# Define the model and sampling parameters
+model_name = "mistralai/Ministral-8B-Instruct-2410"
+sampling_params = SamplingParams(max_tokens=8192)
 
-# Load the model and tokenizer
-model_name = "mrcuddle/Ministral-Instruct-2410-8B-DPO-RP"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Initialize the LLM model
+llm = LLM(model=model_name, tokenizer_mode="mistral", config_format="mistral", load_format="mistral")
 
+@spaces.GPU
 # Define the chatbot function
 def chatbot(message, history):
-    # Tokenize the input
-    inputs = tokenizer(message, return_tensors="pt")
+    # Prepare the messages for the model
+    messages = [
+        {
+            "role": "user",
+            "content": message
+        },
+    ]
 
     # Generate the response
-    with torch.no_grad():
-        outputs = model.generate(**inputs)
-
-    # Decode the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    outputs = llm.chat(messages, sampling_params=sampling_params)
+    response = outputs[0].outputs[0].text
 
     # Append the response to the history
     history.append((message, response))
@@ -25,7 +30,7 @@ def chatbot(message, history):
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Chatbot using mrcuddle/Ministral-Instruct-2410-8B-DPO-RP")
+    gr.Markdown("## Chatbot using mistralai/Ministral-8B-Instruct-2410")
 
     # Create a Chatbot component
     chatbot = gr.Chatbot([], elem_id="chatbot")