import gradio as gr
from gradio_pdf import PDF
from pdf2image import convert_from_path
from transformers import pipeline
from pathlib import Path
from pypdf import PdfReader
import numpy as np

dir_ = Path(__file__).parent

p = pipeline(
    "document-question-answering",
    model="impira/layoutlm-document-qa",
)

def qa(question: str, doc: str) -> str:
    reader  = PdfReader(doc)

    text = []
    for p in np.arange(0, len(reader.pages), 1):
      page = reader.pages[int(p)]
    
      # extracting text from page
      text.append(page.extract_text())
    text = ' '.join(text)
    #output = p(img, question)
    return text


demo = gr.Interface(
    qa,
    [gr.Textbox(label="Question"), PDF(label="Document")],
    gr.Textbox(),
)

if __name__ == "__main__":
    demo.launch()