ignaciaginting/donut_finance_agent

1import streamlit as st2from transformers import pipeline3import torch4from pdf2image import convert_from_bytes5from PIL import Image6 7st.set_page_config(page_title="Donut PDF QA", layout="centered")8 9@st.cache_resource10def load_model():11    return pipeline(12        task="document-question-answering",13        model="naver-clova-ix/donut-base-finetuned-docvqa",14        device=0 if torch.cuda.is_available() else -1,15        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float3216    )17 18qa_pipeline = load_model()19 20 21st.title("📄 Donut: PDF Question Answering")22 23uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])24question = st.text_input("Ask a question about the document")25 26if uploaded_file and question:27    st.write("Reading and converting PDF...")28    images = convert_from_bytes(uploaded_file.read(), dpi=200)29 30    page_number = st.number_input("Select page", min_value=1, max_value=len(images), value=1, step=1)31    page_image = images[page_number - 1]32    st.image(page_image, caption=f"Page {page_number}")33 34    with st.spinner("Finding answer..."):35        result = qa_pipeline(image=page_image, question=question)36        st.success("Answer:")37        st.write(result[0]['answer'])38