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()