import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
import torch
import tempfile

# Load the model and tokenizer
model_name = "deepset/bert-base-cased-squad2"  # Swap for BioBERT or similar if needed
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)

# Sample context – replace with a richer medical dataset
context = """
Diabetes is a chronic 1health condition that affects how your body turns food into energy. 
Most food is broken down into glucose and released into the bloodstream. 
When blood sugar rises, the pancreas releases insulin. 
Insulin lets blood sugar enter the body’s cells for energy. 
There are three main types of diabetes: type 1, type 2, and gestational.
"""



def medical_chatbot(audio_input):
    # Transcribe audio to text
    transcription = audio_input["text"]
    if not transcription.strip():
        return "Please describe your symptoms or ask a medical question.", None

    result = qa_pipeline({
        "question": transcription,
        "context": context
    })

    answer = result["answer"]

    # Save TTS output to a temporary file
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
        engine.save_to_file(answer, fp.name)
        engine.runAndWait()
        audio_path = fp.name

    return answer, audio_path

# Gradio Interface
iface = gr.Interface(
    fn=medical_chatbot,
    outputs=[
        gr.Textbox(label="Answer"),
        
    ],
    title="AI Medical Chatbot (Voice)",
    description="Ask a medical question by speaking. The AI will respond with a text and voice answer. Not a substitute for professional medical advice."
)

if __name__ == "__main__":
    iface.launch()