import gradio as gr
from transformers import pipeline

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

# Define the function to use the pipeline
def answer_question(image, question):
    result = pipe(image, question)
    return result[0]['answer']

# Create the Gradio interface
iface = gr.Interface(
    fn=answer_question,  # function to process input
    inputs=[
        gr.inputs.Image(type="filepath", label="Upload Document Image"),  # input for document image
        gr.inputs.Textbox(lines=2, label="Question"),  # input for question
    ],
    outputs="text",  # output will be text (the answer)
    title="Document Question Answering",  # Title of the interface
    description="Upload a document image and ask a question related to its content.",  # Description for users
)

# Launch the interface
iface.launch()