Update app.py

app.py

@@ -1,9 +1,6 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
-import gradio as gr
 
-# Define the device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model_id = "Narrativaai/BioGPT-Large-finetuned-chatdoctor"
 
@@ -11,31 +8,21 @@ 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,
+        top_p=0.75,
+        top_k=40,
         num_beams=2,
         **kwargs,
 ):
-    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:
-    """
     inputs = tokenizer(prompt, return_tensors="pt")
-    # 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"]
+    input_ids = inputs["input_ids"].to("cuda")
+    attention_mask = inputs["attention_mask"].to("cuda")
     generation_config = GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
         num_beams=num_beams,
         **kwargs,
     )
@@ -66,6 +53,9 @@ Hi i have sore lumps under the skin on my legs. they started on my left ankle an
 ### Response:
 """
 
+print(answer_question(example_prompt))
+
+
 def gui_interface(prompt):
     return answer_question(prompt)