import gradio as gr
from transformers import pipeline
import PyPDF2
import markdown

# Preload models
models = {
    "distilbert-base-uncased-distilled-squad": "distilbert-base-uncased-distilled-squad",
    "roberta-base-squad2": "deepset/roberta-base-squad2",
    "bert-large-uncased-whole-word-masking-finetuned-squad": "bert-large-uncased-whole-word-masking-finetuned-squad",
    "albert-base-v2": "twmkn9/albert-base-v2-squad2",
    "xlm-roberta-large-squad2": "deepset/xlm-roberta-large-squad2"
}

loaded_models = {}

def load_model(model_name):
    if model_name not in loaded_models:
        loaded_models[model_name] = pipeline("question-answering", model=models[model_name])
    return loaded_models[model_name]

def answer_question(model_name, file, question, status):
    status = "Loading model..."
    model = load_model(model_name)
    
    if file is not None:
        if file.type == "application/pdf":
            pdf_reader = PyPDF2.PdfFileReader(file)
            context = ""
            for page_num in range(pdf_reader.numPages):
                context += pdf_reader.getPage(page_num).extract_text()
        elif file.name.endswith(".md"):
            context = file.read().decode('utf-8')
            context = markdown.markdown(context)
        else:
            context = file.read().decode('utf-8')
    else:
        context = ""
    
    result = model(question=question, context=context)
    answer = result['answer']
    score = result['score']
    
    # Explain score
    score_explanation = f"The confidence score ranges from 0 to 1, where a higher score indicates higher confidence in the answer's correctness. In this case, the score is {score:.2f}. A score closer to 1 implies the model is very confident about the answer."
    
    status = "Model loaded"
    return answer, f"{score:.2f}", score_explanation, status

# Define the Gradio interface
with gr.Blocks() as interface:
    gr.Markdown(
        """
        # Question Answering System
        Upload a document (text, PDF, or Markdown) and ask questions to get answers based on the context.
        
        **Supported File Types**: `.txt`, `.pdf`, `.md`
        """)
    
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=list(models.keys()),
            label="Select Model",
            value="distilbert-base-uncased-distilled-squad"
        )
    
    with gr.Row():
        file_input = gr.File(label="Upload Document", file_types=["text", "pdf", "markdown"])
        question_input = gr.Textbox(lines=2, placeholder="Enter your question here...", label="Question")
    
    with gr.Row():
        answer_output = gr.Textbox(label="Answer")
        score_output = gr.Textbox(label="Confidence Score")
        explanation_output = gr.Textbox(label="Score Explanation")
    
    with gr.Row():
        submit_button = gr.Button("Submit")
    
    status_output = gr.Markdown(value="")

    def on_submit(model_name, file, question):
        return answer_question(model_name, file, question, status="Loading model...")

    submit_button.click(
        on_submit,
        inputs=[model_dropdown, file_input, question_input],
        outputs=[answer_output, score_output, explanation_output, status_output]
    )

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