tosin2013's picture
Update app.py
dc15a5c verified
raw
history blame
2.75 kB
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()