Update app.py

app.py

@@ -1,71 +1,32 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the PastelMed model and tokenizer
+# Load the fine-tuned model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("harishussain12/PastelMed")
 model = AutoModelForCausalLM.from_pretrained("harishussain12/PastelMed")
 
-# Patient database (example, can be expanded)
-patients_db = []
-
-# Doctor's assistant function
-def doctor_assistant(question):
-    # Encode the input question
-    inputs = tokenizer(question, return_tensors="pt")
-    
-    # Generate a response from the model
-    outputs = model.generate(inputs["input_ids"], max_length=200, num_return_sequences=1)
+# Function to generate response from the model
+def doctor_consultant(query):
+    # Encode the input query and generate the model's response
+    inputs = tokenizer(query, return_tensors="pt")
+    outputs = model.generate(inputs['input_ids'], max_length=200, num_return_sequences=1, no_repeat_ngram_size=2, top_p=0.95, temperature=0.7)
     
-    # Decode the generated response
+    # Decode the output and return the response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
     return response
 
-# Function to register a patient (simplified version)
-def register_patient(name, age, gender, password):
-    patient_id = len(patients_db) + 1
-    patients_db.append({
-        "ID": patient_id,
-        "Name": name,
-        "Age": age,
-        "Gender": gender,
-        "Password": password,
-        "Diagnosis": "",
-        "Medications": "",
-        "Precautions": "",
-        "Doctor": ""
-    })
-    return f"✅ Patient {name} registered successfully. Patient ID: {patient_id}"
-
-# Gradio Interface for Doctor Assistance
-doctor_assistant_interface = gr.Interface(
-    fn=doctor_assistant,
-    inputs=gr.Textbox(label="Ask a Question to the Doctor Assistant"),
-    outputs="text",
-    title="Doctor Assistant",
-    description="Ask the assistant for medical advice and it will generate a response based on the PastelMed model."
-)
-
-# Gradio Interface for Patient Registration (for testing)
-registration_interface = gr.Interface(
-    fn=register_patient,
-    inputs=[
-        gr.Textbox(label="Patient Name"),
-        gr.Number(label="Age"),
-        gr.Radio(label="Gender", choices=["Male", "Female", "Other"]),
-        gr.Textbox(label="Set Password", type="password"),
-    ],
-    outputs="text",
-)
-
-# Gradio App Layout
+# Gradio Interface
 with gr.Blocks() as app:
-    gr.Markdown("# Medico GPT")
+    gr.Markdown("# Doctor Consultant Assistant")
     
-    with gr.Tab("Patient Registration"):
-        registration_interface.render()
+    with gr.Row():
+        gr.Textbox(label="Ask the Doctor", placeholder="Enter your symptoms or question", lines=3, elem_id="input_text")
     
-    with gr.Tab("Doctor Assistant"):
-        doctor_assistant_interface.render()
-
-app.launch(share=True)
+    with gr.Row():
+        gr.Button("Get Response", elem_id="response_button")
+    
+    with gr.Row():
+        gr.Textbox(label="Doctor's Response", elem_id="response_output", interactive=False)
+    
+    # Connect the function to the interface
+    gr.Interface(fn=doctor_consultant, inputs="text", outputs="text").launch(share=True)