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