QandA / app.py
rajeshuriti's picture
Created QandA AI app
2f2ce5c
import gradio as gr
from transformers import pipeline
# Part 1: Load the model ONCE
print("Loading the MobileBERT model...")
info_extractor = pipeline("question-answering", model="csarron/mobilebert-uncased-squad-v2")
print("Model loaded successfully!")
# Part 2: Create the function that the UI will call
# This function takes the document and question from the UI,
# gets the answer from the model, and returns it.
def extract_information(context, question):
print(f"Extracting answer for question: '{question}'")
result = info_extractor(question=question, context=context)
return result['answer']
# Part 3: Build and launch the Gradio Interface
print("Launching Gradio interface...")
iface = gr.Interface(
fn=extract_information,
inputs=[
gr.Textbox(lines=7, label="Document", placeholder="Paste the document or text you want to ask questions about..."),
gr.Textbox(label="Question", placeholder="What specific detail are you looking for?")
],
outputs=gr.Textbox(label="Answer"),
title="πŸ’‘ Efficient Information Extractor",
description="Ask a question about the document below to pull out specific details using a MobileBERT model."
)
iface.launch()