initial commit

app.py

@@ -0,0 +1,207 @@
+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
+
+
+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).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"])
+    #    gen_model = genai.GenerativeModel(model_name="gemini-1.5-flash")
+        gen_model = genai.GenerativeModel(model_name="gemini-1.5-pro")
+        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")

requirements.txt

@@ -0,0 +1,107 @@
+accelerate==0.32.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.7.0
+async-timeout==4.0.3
+attrs==23.2.0
+black==24.4.2
+blinker==1.8.2
+cachetools==5.4.0
+certifi==2024.7.4
+charset-normalizer==3.3.2
+click==8.1.7
+colpali_engine @ git+https://github.com/illuin-tech/colpali@8b01824546c62e46383ce26b439d9bfc6468f763
+datasets==2.20.0
+dill==0.3.8
+eval_type_backport==0.2.0
+filelock==3.15.4
+frozenlist==1.4.1
+fsspec==2024.5.0
+gitdb==4.0.11
+GitPython==3.1.43
+google-ai-generativelanguage==0.6.6
+google-api-core==2.19.1
+google-api-python-client==2.137.0
+google-auth==2.32.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.7.2
+googleapis-common-protos==1.63.2
+GPUtil==1.4.0
+grpcio==1.65.1
+grpcio-status==1.62.2
+httplib2==0.22.0
+huggingface-hub==0.24.0
+idna==3.7
+importlib_metadata==7.2.1
+Jinja2==3.1.4
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+mpmath==1.3.0
+mteb==1.12.85
+multidict==6.0.5
+multiprocess==0.70.16
+mypy-extensions==1.0.0
+networkx==3.3
+numpy==1.26.4
+packaging==23.2
+pandas==2.2.2
+pathspec==0.12.1
+pdf2image==1.17.0
+peft==0.11.1
+pillow==10.4.0
+platformdirs==4.2.2
+polars==1.2.1
+proto-plus==1.24.0
+protobuf==4.25.3
+psutil==6.0.0
+pyarrow==17.0.0
+pyarrow-hotfix==0.6
+pyasn1==0.6.0
+pyasn1_modules==0.4.0
+pydantic==2.8.2
+pydantic_core==2.20.1
+pydeck==0.9.1
+Pygments==2.18.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+pytrec_eval-terrier==0.5.6
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rich==13.7.1
+rpds-py==0.19.0
+rsa==4.9
+safetensors==0.4.3
+scikit-learn==1.5.1
+scipy==1.14.0
+sentence-transformers==3.0.1
+six==1.16.0
+smmap==5.0.1
+streamlit==1.31.1
+sympy==1.13.1
+tenacity==8.5.0
+threadpoolctl==3.5.0
+tokenizers==0.19.1
+toml==0.10.2
+tomli==2.0.1
+toolz==0.12.1
+torch==2.3.1
+tornado==6.4.1
+tqdm==4.66.4
+transformers==4.42.4
+typing_extensions==4.12.2
+tzdata==2024.1
+tzlocal==5.2
+uritemplate==4.1.1
+urllib3==2.2.2
+validators==0.33.0
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.19.2