Update app.py

app.py

@@ -1,6 +1,8 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
+# Define the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model_id = "Narrativaai/BioGPT-Large-finetuned-chatdoctor"
 
@@ -8,6 +10,9 @@ tokenizer = AutoTokenizer.from_pretrained("microsoft/BioGPT-Large")
 
 model = AutoModelForCausalLM.from_pretrained(model_id)
 
+# Move the model to the device
+model = model.to(device)
+
 def answer_question(
         prompt,
         temperature=0.1,
@@ -17,8 +22,10 @@ def answer_question(
         **kwargs,
 ):
     inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"].to("cuda")
-    attention_mask = inputs["attention_mask"].to("cuda")
+    # Move the inputs to the device
+    inputs = {key: val.to(device) for key, val in inputs.items()}
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
     generation_config = GenerationConfig(
         temperature=temperature,
         top_p=top_p,
@@ -57,6 +64,19 @@ print(answer_question(example_prompt))
 
 
 def gui_interface(prompt):
+    
+    prompt="""
+    Below is an instruction that describes a task, paired with an input that provides further context.Write a response that appropriately completes the request.
+    
+    ### Instruction:
+    If you are a doctor, please answer the medical questions based on the patient's description.
+    
+    ### Input:
+    """+prompt+"""
+    ### Response:
+    """
+
+    print(p)
     return answer_question(prompt)
 
 iface = gr.Interface(fn=gui_interface, inputs="text", outputs="text")