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