Spaces:
Running
Running
| 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") | |