|
|
import gradio as gr |
|
|
from transformers import pipeline |
|
|
import PyPDF2 |
|
|
import markdown |
|
|
|
|
|
|
|
|
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'] |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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() |
|
|
|
