import gradio as gr
from pypdf import PdfReader
from transformers import pipeline


pipe = pipeline("question-answering", model='distilbert-base-uncased-distilled-squad')

def predict(file_obj, question):
    docText = ""   
    reader = PdfReader(file_obj.name)
    print(len(reader.pages))
        
    for page in reader.pages:
        txt = page.extract_text()
        docText += f" {txt}"

    qaInput = {
    'question': question,
    'context': docText
    } 

    print(docText[-10:-1])
    
    return pipe(qaInput)

demo = gr.Interface(
  fn=predict,
  inputs=[gr.File(file_count="single", file_types=[".pdf"]), "text"],
  outputs=["text"]
)

demo.launch()