Update app.py

app.py

@@ -1,11 +1,66 @@
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL = "docto/Docto-Bot"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModelForCausalLM.from_pretrained(MODEL)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+def get_reply(user_input):
+
+    prompt = f"Question: {user_input}\nAnswer:"
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt"
+    ).to(device)
+
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.9,
+        repetition_penalty=1.15,
+        no_repeat_ngram_size=3,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
+    response = tokenizer.decode(
+        outputs[0],
+        skip_special_tokens=True
+    )
+
+    if "Answer:" in response:
+        response = response.split("Answer:", 1)[1]
+
+    return response.strip()
+
+
 iface = gr.Interface(
     fn=get_reply,
+
     inputs=gr.Textbox(
         lines=2,
-        placeholder="Ask a question..."
+        placeholder="Ask a medical question..."
     ),
-    outputs=gr.Textbox(label="Bot Response"),
+
+    outputs=gr.Textbox(
+        label="Response"
+    ),
+
     title="Docto-Bot",
+
     description="Medical Question Answering Bot"
 )