tosin2013's picture
Update app.py
55e40d9 verified
raw
history blame
2.75 kB
import gradio as gr
from transformers import pipeline
from transformers.pipelines import PipelineException
# 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:
try:
loaded_models[model_name] = pipeline("question-answering", model=models[model_name])
except PipelineException as e:
return str(e)
return loaded_models[model_name]
def answer_question(model_name, file, question):
model = load_model(model_name)
if isinstance(model, str):
return model, "", ""
context = file.read() 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")
gr.Markdown("### Progress")
with gr.Row():
progress_bar = gr.Progress(label="Loading Model...")
submit_button.click(
answer_question,
inputs=[model_dropdown, file_input, question_input],
outputs=[answer_output, score_output, explanation_output],
show_progress=progress_bar,
)
if __name__ == "__main__":
interface.launch()