image_qa / app.py
dwh1142's picture
Update app.py
502fd5c
raw
history blame contribute delete
No virus
666 Bytes
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()