app.py

import hmac
import os
import tempfile

from colpali_engine.models.paligemma_colbert_architecture import ColPali
from colpali_engine.utils.colpali_processing_utils import process_images
from colpali_engine.utils.colpali_processing_utils import process_queries
import google.generativeai as genai
import numpy as np
import pdf2image
from PIL import Image
import requests
import streamlit as st
import torch
from torch.utils.data import DataLoader
from transformers import AutoProcessor


def check_password():
    """Returns `True` if the user had the correct password."""

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if hmac.compare_digest(st.session_state["password"], st.secrets["password"]):
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password.
        else:
            st.session_state["password_correct"] = False

    # Return True if the password is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show input for password.
    st.text_input(
        "Password", type="password", on_change=password_entered, key="password"
    )
    if "password_correct" in st.session_state:
        st.error("😕 Password incorrect")
    return False


if not check_password():
    st.stop()  # Do not continue if check_password is not True.


os.environ["TOKENIZERS_PARALLELISM"] = "false"
SS = st.session_state


def initialize_session_state():
    keys = [
        "colpali_model",
        "page_images",
        "page_embeddings",
        "retrieved_page_images",
        "retrieved_page_scores",
        "response",
    ]
    for key in keys:
        if key not in SS:
            SS[key] = None


def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    return device


def get_dtype(device: torch.device):
    if device == torch.device("cuda"):
        if torch.cuda.is_bf16_supported():
            dtype = torch.bfloat16
        else:
            dtype = torch.float16
    elif device == torch.device("mps"):
        dtype = torch.float32
    else:
        dtype = torch.float32
    return dtype


def load_colpali_model():
    paligemma_model_name = "google/paligemma-3b-mix-448"
    colpali_model_name = "vidore/colpali"
    device = get_device()
    dtype = get_dtype(device)

    model = ColPali.from_pretrained(
        paligemma_model_name,
        torch_dtype=dtype,
        token=st.secrets["hf_access_token"],
    ).eval()
    model.load_adapter(colpali_model_name)
    model.to(device)
    processor = AutoProcessor.from_pretrained(colpali_model_name)
    return model, processor


def embed_page_images(model, processor, page_images, batch_size=1):
    dataloader = DataLoader(
        page_images,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=lambda x: process_images(processor, x),
    )
    page_embeddings = []
    pbar = st.progress(0, text="embedding pages")
    for ibatch, batch in enumerate(dataloader):
        with torch.no_grad():
            batch = {k: v.to(model.device) for k, v in batch.items()}
            embeddings = model(**batch)
            page_embeddings.extend(list(torch.unbind(embeddings.to("cpu"))))
        pbar.progress((ibatch + 1) / len(page_images), text="embedding pages")
    return np.array([el.to(torch.float32) for el in page_embeddings])


def embed_query_texts(model, processor, query_texts, batch_size=1):
    # 448 is from the paligemma resolution we loaded
    dummy_image = Image.new("RGB", (448, 448), (255, 255, 255))
    dataloader = DataLoader(
        query_texts,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=lambda x: process_queries(processor, x, dummy_image),
    )
    query_embeddings = []
    for batch in dataloader:
        with torch.no_grad():
            batch = {k: v.to(model.device) for k, v in batch.items()}
            embeddings = model(**batch)
            query_embeddings.extend(list(torch.unbind(embeddings.to("cpu"))))
    return np.array([el.to(torch.float32) for el in query_embeddings])[0]


def get_pdf_page_images_from_bytes(
    pdf_bytes: bytes,
    use_tmp_dir=False,
):
    if use_tmp_dir:
        with tempfile.TemporaryDirectory() as tmp_path:
            page_images = pdf2image.convert_from_bytes(
                pdf_bytes, output_folder=tmp_path
            )
    else:
        page_images = pdf2image.convert_from_bytes(pdf_bytes)
    return page_images


def get_pdf_bytes_from_url(url: str) -> bytes | None:
    response = requests.get(url)
    if response.status_code == 200:
        return response.content
    else:
        print(f"failed to fetch {url}")
        print(response)
        return None


def display_pages(page_images, key, captions=None):
    n_cols = st.slider("ncol", min_value=1, max_value=8, value=4, step=1, key=key)
    cols = st.columns(n_cols)
    for ii_page, page_image in enumerate(page_images):
        ii_col = ii_page % n_cols
        with cols[ii_col]:
            if captions is not None:
                caption = captions[ii_page]
            else:
                caption = None
            st.image(page_image, caption=caption)


initialize_session_state()


if SS["colpali_model"] is None:
    SS["colpali_model"], SS["processor"] = load_colpali_model()


with st.sidebar:

    with st.container(border=True):
        st.header("Load PDF (URL or Upload)")
        st.write("When a PDF is loaded, each page will be turned into an image.")

        url = st.text_input("Provide a URL", "https://arxiv.org/pdf/2404.15549v2")
        if st.button("load paper from url"):
            pdf_bytes = get_pdf_bytes_from_url(url)
            SS["page_images"] = get_pdf_page_images_from_bytes(pdf_bytes)

        uploaded_file = st.file_uploader("Upload a file", type=["pdf"])
        if uploaded_file is not None:
            pdf_bytes = uploaded_file.getvalue()
            SS["page_images"] = get_pdf_page_images_from_bytes(pdf_bytes)

    with st.container(border=True):
        st.header("Embed Page Images")
        st.write(
            "In order to retrieve relevant images for a query, we must first embed the images."
        )
        if st.button("embed pages"):
            SS["page_embeddings"] = embed_page_images(
                SS["colpali_model"],
                SS["processor"],
                SS["page_images"],
            )

    if SS["page_images"] is not None:
        st.write("Num Page Images: {}".format(len(SS["page_images"])))

    if SS["page_embeddings"] is not None:
        st.write("Page Embeddings Shape: {}".format(SS["page_embeddings"].shape))


with st.container(border=True):
    query = st.text_area("query")

    prompt_template_default = """Your goal is to answer queries based on the provided images. Each image is one page from a single PDF document. Provide answers that are at least 3 sentences long. Clearly explain the reasoning behind your answer. Create trustworthy answers by referencing the material in the PDF pages. Do not reference page numbers unless they appear on the page images.

---

{query}"""

    with st.expander("Prompt Template"):
        prompt_template = st.text_area(
            "Customize the prompt template",
            prompt_template_default,
            height=200,
        )

    top_k = st.slider(
        "num pages to retrieve", min_value=1, max_value=8, value=3, step=1
    )
    if st.button("answer query"):
        SS["query_embeddings"] = embed_query_texts(
            SS["colpali_model"],
            SS["processor"],
            [query],
        )

        page_query_scores = []
        for ipage in range(len(SS["page_embeddings"])):
            # for every query token find the max_sim with every page patch
            patch_query_scores = np.dot(
                SS["page_embeddings"][ipage],
                SS["query_embeddings"].T,
            )
            max_sim_score = patch_query_scores.max(axis=0).sum()
            page_query_scores.append(max_sim_score)

        page_query_scores = np.array(page_query_scores)
        i_ranked_pages = np.argsort(-page_query_scores)

        page_images = []
        page_scores = []
        num_pages = len(SS["page_images"])
        for ii in range(min(top_k, num_pages)):
            page_images.append(SS["page_images"][i_ranked_pages[ii]])
            page_scores.append(page_query_scores[i_ranked_pages[ii]])
        SS["retrieved_page_images"] = page_images
        SS["retrieved_page_scores"] = page_scores

        prompt = [prompt_template.format(query=query)] + page_images

        genai.configure(api_key=st.secrets["google_genai_api_key"])
        #        genai_model_name = "gemini-1.5-flash"
        genai_model_name = "gemini-1.5-pro"
        gen_model = genai.GenerativeModel(
            model_name=genai_model_name,
            generation_config=genai.GenerationConfig(
                temperature=0.0,
            ),
        )
        response = gen_model.generate_content(prompt)
        text = response.candidates[0].content.parts[0].text
        SS["response"] = text


if SS["response"] is not None:
    st.header("Response")
    st.write(SS["response"])
    st.header("Retrieved Pages")
    display_pages(
        SS["retrieved_page_images"],
        "retrieved_pages",
        captions=[f"Score={el:.2f}" for el in SS["retrieved_page_scores"]],
    )


if SS["page_images"] is not None:
    st.header("All Pages")
    display_pages(SS["page_images"], "all_pages")