import gradio as gr
import json

from transformers import pipeline
from pytesseract import pytesseract


# Load the model
nlp = pipeline(
    "document-question-answering",
    model="impira/layoutlm-document-qa",
)


# Function to perform the question answering
def perform_question_answering(image, question):
    answer = nlp(image, question)

    # Format the answer as JSON
    answer_json = json.dumps(answer, indent=4)

    return answer_json


# Create the Gradio interface
inputs = [
    gr.inputs.Image(type="pil", label="Upload Image"),
    gr.inputs.Textbox(label="Question")
]

outputs = gr.outputs.Textbox(label="Answer")

iface = gr.Interface(fn=perform_question_answering, inputs=inputs, outputs="text")


# Launch the Gradio interface
iface.launch()