ignaciaginting's picture
Updates
fe0246c verified
raw
history blame
1.28 kB
import streamlit as st
from transformers import pipeline
from PIL import Image
import tempfile
import fitz # PyMuPDF
# Load the model
@st.cache_resource
def load_model():
return pipeline("document-question-answering", model="impira/layoutlm-document-qa")
qa_pipeline = load_model()
st.title("πŸ“„ Document Question Answering App")
st.write("Upload a PDF file, enter a question, and get answers from the document.")
# Upload PDF
pdf_file = st.file_uploader("Upload PDF", type=["pdf"])
# Ask a question
question = st.text_input("Ask a question about the document:")
if pdf_file and question:
# Convert first page of PDF to image using PyMuPDF
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
tmp_file.write(pdf_file.read())
pdf_path = tmp_file.name
doc = fitz.open(pdf_path)
page = doc.load_page(0) # just first page for now
pix = page.get_pixmap(dpi=150)
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
# Show the rendered page
st.image(img, caption="Page 1 of PDF")
# Run the pipeline
with st.spinner("Searching for the answer..."):
result = qa_pipeline(img, question)
st.success(f"**Answer:** {result['answer']} (score: {result['score']:.2f})")