more stuff

app.py

@@ -23,8 +23,9 @@ def initialize_session_state():
     keys = [
         "colpali_model",
         "page_images",
+        "page_embeddings",
         "retrieved_page_images",
-        "response",
+        "retrieved_page_scores" "response",
     ]
     for key in keys:
         if key not in SS:
@@ -68,7 +69,7 @@ def load_colpali_model():
     return model, processor
 
 
-def embed_page_images(model, processor, page_images, batch_size=2):
+def embed_page_images(model, processor, page_images, batch_size=1):
     dataloader = DataLoader(
         page_images,
         batch_size=batch_size,
@@ -76,11 +77,13 @@ def embed_page_images(model, processor, page_images, batch_size=2):
         collate_fn=lambda x: process_images(processor, x),
     )
     page_embeddings = []
-    for batch in dataloader:
+    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(page_embeddings)
 
 
@@ -102,14 +105,15 @@ def embed_query_texts(model, processor, query_texts, batch_size=1):
     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)
+            page_images = pdf2image.convert_from_bytes(
+                pdf_bytes, output_folder=tmp_path
+            )
     else:
         page_images = pdf2image.convert_from_bytes(pdf_bytes)
     return page_images
@@ -125,13 +129,17 @@ def get_pdf_bytes_from_url(url: str) -> bytes | None:
         return None
 
 
-def display_pages(page_images, key):
+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]:
-            st.image(page_image)
+            if captions is not None:
+                caption = captions[ii_page]
+            else:
+                caption = None
+            st.image(page_image, caption=caption)
 
 
 initialize_session_state()
@@ -142,24 +150,59 @@ if SS["colpali_model"] is None:
 
 
 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)
+    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)
 
-    if st.button("embed pages"):
-        SS["page_embeddings"] = embed_page_images(
-            SS["colpali_model"],
-            SS["processor"],
-            SS["page_images"],
+        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")
-    top_k = st.slider("num pages to retrieve", min_value=1, max_value=8, value=3, step=1)
+
+    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"],
@@ -171,7 +214,7 @@ with st.container(border=True):
         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["page_embeddings"][ipage],
                 SS["query_embeddings"].T,
             )
             max_sim_score = patch_query_scores.max(axis=0).sum()
@@ -181,25 +224,23 @@ with st.container(border=True):
         i_ranked_pages = np.argsort(-page_query_scores)
 
         page_images = []
-        for ii in range(top_k):
+        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 = [
-            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
+        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-flash"
         genai_model_name = "gemini-1.5-pro"
         gen_model = genai.GenerativeModel(
             model_name=genai_model_name,
             generation_config=genai.GenerationConfig(
-                temperature=0.1,
+                temperature=0.0,
             ),
         )
         response = gen_model.generate_content(prompt)
@@ -208,12 +249,16 @@ with st.container(border=True):
 
 
 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")
-
+    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 PDF Pages")
+    st.header("All Pages")
     display_pages(SS["page_images"], "all_pages")