Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import pipeline | |
# 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): | |
model = load_model(model_name) | |
context = file.read().decode('utf-8') if file else "" | |
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." | |
return answer, f"{score:.2f}", score_explanation | |
# Define the Gradio interface | |
with gr.Blocks() as interface: | |
gr.Markdown( | |
""" | |
# Question Answering System | |
Upload a document and ask questions to get answers based on the context. | |
""") | |
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") | |
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") | |
# Define a status area for progress | |
status = gr.Markdown(value="") | |
def on_submit(model_name, file, question): | |
status.update(value="Loading model...") | |
answer, score, explanation = answer_question(model_name, file, question) | |
status.update(value="Model loaded") | |
return answer, score, explanation | |
submit_button.click( | |
on_submit, | |
inputs=[model_dropdown, file_input, question_input], | |
outputs=[answer_output, score_output, explanation_output] | |
) | |
if __name__ == "__main__": | |
interface.launch() | |