Update app.py

app.py

@@ -1,39 +1,94 @@
-import random
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import torch
 
-tokenizer = AutoTokenizer.from_pretrained("docto/Docto-Bot")
-model = AutoModelForCausalLM.from_pretrained("docto/Docto-Bot")
-special_token = '<|endoftext|>'
-
-
-def get_reply(userinput):
-        prompt_text = f'Question: {userinput}\nAnswer:'
-        encoded_prompt = tokenizer.encode(prompt_text,
-                                  add_special_tokens = False,
-                                  return_tensors = 'pt')
-
-        output_sequences = model.generate(
-            input_ids = encoded_prompt,
-            max_length = 500,
-            temperature = 0.1,
-            top_k = 20,
-            top_p = 0.9,
-            repetition_penalty = 1,
-            do_sample = True,
-            num_return_sequences = 1
+MODEL_NAME = "docto/Docto-Bot"
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# Set pad token if missing
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+def get_reply(user_input):
+
+    if not user_input.strip():
+        return "Please enter a question."
+
+    try:
+        # Build prompt
+        prompt = f"Question: {user_input}\nAnswer:"
+
+        # Tokenize
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt"
+        ).to(device)
+
+        # Generate response
+        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,
+
+            eos_token_id=tokenizer.eos_token_id,
+
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+        # Decode
+        response = tokenizer.decode(
+            outputs[0],
+            skip_special_tokens=True
         )
 
-        # result = tokenizer.decode(random.choice(output_sequences))
-        # result = result[result.index("Answer: "):result.index(special_token)]
-        try:
-            result = tokenizer.decode(random.choice(output_sequences))
-            result = result[result.index("Answer: "):result.index(special_token)]
-            return (result[8:])
-            
-        except:
-            return "Sorry! I don\'t Know"
-            
-iface = gr.Interface(fn=get_reply, inputs=["text"], outputs=["textbox"]).launch()
-
- 
+        # Extract answer only
+        if "Answer:" in response:
+            response = response.split("Answer:", 1)[1]
+
+        return response.strip()
+
+    except Exception as e:
+        return f"Error: {e}"
+
+
+# Gradio UI
+iface = gr.Interface(
+    fn=get_reply,
+
+    inputs=gr.Textbox(
+        lines=2,
+        placeholder="Ask a question..."
+    ),
+
+    outputs=gr.Textbox(
+        label="Bot Response"
+    ),
+
+    title="Docto-Bot",
+
+    description="Medical Question Answering Bot",
+
+    allow_flagging="never"
+)
+
+iface.launch()