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})")