test_space / test_app.py
Raptor1038's picture
Initial commit
b67003e
raw
history blame
1.57 kB
import gradio as gr
import tensorflow as tf
import pandas as pd
from transformers import AutoTokenizer, TFAutoModelForQuestionAnswering
# Load the pre-trained model and tokenizer
# Import required libraries
# Load the tokenizer and model
model_name = "distilbert-base-uncased-distilled-squad"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = TFAutoModelForQuestionAnswering.from_pretrained(model_name)
def read_text_file(file):
with open(file.name) as f:
content = f.read()
return content
# Define a function to perform the question answering
def answer_question(doc, question):
# context = doc.read().decode('utf-8')
context = read_text_file(doc)
# Tokenize the context and question
inputs = tokenizer(question, context, return_tensors="tf")
# Get the answer span
start_scores = model(inputs)[0]
end_scores = model(inputs)[1]
answer_start = tf.argmax(start_scores, axis=1).numpy()[0]
answer_end = tf.argmax(end_scores, axis=1).numpy()[0] + 1
answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
return answer
# Create a Gradio interface
interface = gr.Interface(
fn=answer_question,
inputs=[
gr.inputs.File(label="doc"),
gr.inputs.Textbox(label="question")
],
outputs=gr.outputs.Textbox(label="answer"),
title="Document Question Answering",
description="Upload a document and ask a question about its contents.",
theme="default"
)
# Launch the interface
interface.launch()