File size: 1,435 Bytes
fe0246c 8bc3f30 fe0246c 8bc3f30 fe0246c e22af09 d75400e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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)
if result:
top_answer = result[0] # get the highest-scoring answer
st.success(f"**Answer:** {top_answer['answer']} (score: {top_answer['score']:.2f})")
else:
st.warning("No answer found.")
|