Update app.py to use fine-tuned model

app.py

@@ -1,21 +1,20 @@
-
 import gradio as gr
-from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load the model
-llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", trust_remote_code=True)
+# 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")
 
 def generate_response(prompt):
     """
     Generates a response from the model.
     """
-    sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=512)
-    outputs = llm.generate(prompt, sampling_params)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=512)
     
-    # Extract the generated text
-    generated_text = ""
-    for output in outputs:
-        generated_text += output.outputs[0].text
+    # Decode the generated text
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
     return generated_text
 
@@ -40,4 +39,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()