app.py

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


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


def initialize_session_state():
    keys = [
        "colpali_model",
        "page_images",
        "retrieved_page_images",
        "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"):
        dtype = torch.bfloat16
    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=2):
    dataloader = DataLoader(
        page_images,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=lambda x: process_images(processor, x),
    )
    page_embeddings = []
    for batch in 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"))))
    return np.array(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(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):
    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]:
            st.image(page_image)


initialize_session_state()


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


with st.sidebar:
    url = st.text_input("arxiv url", "https://arxiv.org/pdf/2112.01488.pdf")

    if st.button("load paper"):
        pdf_bytes = get_pdf_bytes_from_url(url)
        SS["page_images"] = get_pdf_page_images_from_bytes(pdf_bytes)


    if st.button("embed pages"):
        SS["page_embeddings"] = embed_page_images(
            SS["colpali_model"],
            SS["processor"],
            SS["page_images"],
        )


with st.container(border=True):
    query = st.text_area("query")
    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 = []
        for ii in range(top_k):
            page_images.append(SS["page_images"][i_ranked_pages[ii]])
        SS["retrieved_page_images"] = page_images


        prompt = [
            query +
            " Think through your answer step by step. "
            "Support your answer with descriptions of the images. "
            "Do not infer information that is not in the images.",
        ] + 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.1,
            ),
        )
        response = gen_model.generate_content(prompt)
        text = response.candidates[0].content.parts[0].text
        SS["response"] = text


if SS["response"] is not None:
    st.write(SS["response"])
    st.header("Retrieved Pages")
    display_pages(SS["retrieved_page_images"], "retrieved_pages")



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