Extractive_QA / app.py
Jitendra14355's picture
Update app.py
29c5efc verified
import gradio as gr
from transformers import pipeline
# 1. Load the RoBERTa model for Question Answering
# deepset/roberta-base-squad2 is a highly popular model for this task
model_name = "deepset/roberta-base-squad2"
qa_pipeline = pipeline("question-answering", model=model_name, tokenizer=model_name)
# 2. Define the prediction function
def extract_answer(context, question):
if not context.strip() or not question.strip():
return "Please provide both a context and a question."
try:
# The pipeline returns a dictionary with 'score', 'start', 'end', and 'answer'
result = qa_pipeline(question=question, context=context)
return result['answer']
except Exception as e:
return f"An error occurred: {str(e)}"
# 3. Create the Gradio Interface
iface = gr.Interface(
fn=extract_answer,
inputs=[
gr.Textbox(lines=10, label="Context Paragraph", placeholder="Paste the text you want the model to read here..."),
gr.Textbox(lines=2, label="Question", placeholder="What do you want to know from the text?")
],
outputs=gr.Textbox(label="Extracted Answer"),
title="Extractive QA with RoBERTa",
description="This application uses a pre-trained RoBERTa model to extract answers to questions based on a provided context.",
theme="default"
)
# 4. Launch the application
iface.launch()