Update app.py

app.py

@@ -1,21 +1,21 @@
 
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from vllm import LLM, SamplingParams
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("braindeck/text2text", trust_remote_code=True, subfolder="checkpoints/model")
-model = AutoModelForCausalLM.from_pretrained("braindeck/text2text", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto", subfolder="checkpoints/model")
+# Load the model
+llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", trust_remote_code=True)
 
 def generate_response(prompt):
     """
     Generates a response from the model.
     """
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=512)
+    sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=512)
+    outputs = llm.generate(prompt, sampling_params)
     
-    # Decode the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract the generated text
+    generated_text = ""
+    for output in outputs:
+        generated_text += output.outputs[0].text
         
     return generated_text