Update app.py

app.py

@@ -1,23 +1,15 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+\import gradio as gr
+from transformers import pipeline
 
-# Load the model and tokenizer with trust_remote_code=True
-model_name = "Flmc/DISC-MedLLM"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
+# Set up the pipeline
+pipe = pipeline("text-generation", model="Flmc/DISC-MedLLM", trust_remote_code=True)
 
-# Function to generate responses
+# Define the function for generating responses
 def generate_response(input_text):
     if not input_text.strip():
         return "Please enter some text to generate a response."
     
-    inputs = tokenizer(input_text, return_tensors="pt")
-    if torch.cuda.is_available():
-        inputs = inputs.to("cuda")
-        model.to("cuda")
-    outputs = model.generate(**inputs, max_new_tokens=150)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = pipe(input_text, max_new_tokens=150, do_sample=True)[0]["generated_text"]
     return response
 
 # Gradio interface