import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# Load model
model_name = "tmt3103/BioASQ-yesno-PudMedBERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def predict_yesno(context, question):
    inputs = tokenizer.encode_plus(
        question,
        context,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding="max_length"
    )

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = F.softmax(logits, dim=1).squeeze()
        pred_id = logits.argmax().item()

    label = "Yes" if pred_id == 1 else "No"
    return f"{label} (Confidence: {probs[pred_id]:.2f})"

# Gradio
with gr.Blocks() as demo:
    gr.Markdown("BioASQ Yes/No Question Answering")
    gr.Markdown("""
    This demo uses a fine-tuned BERT model to answer biomedical yes/no questions based on context.<br>
    **Instructions**:
    1. Type the context and your yes/no question.
    2. Click 'Predict' to get the answer.
    """)

    with gr.Row():
        with gr.Column():
            context_input = gr.Textbox(label="Context", lines=8, placeholder="Paste biomedical context here...")
            question_input = gr.Textbox(label="Question", lines=2, placeholder="Enter your question here...")
            predict_button = gr.Button("Predict")
        with gr.Column():
            output = gr.Textbox(label="Prediction")

    predict_button.click(
        fn=predict_yesno,
        inputs=[context_input, question_input],
        outputs=output
    )

demo.launch(share=True)