Spaces:
Sleeping
Sleeping
| 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() | |