Merge pull request #1672 from albertdbio/fix/update-colpali-demo-pdf-scraping

Update ColPali demo pdf scraping

Author: thomasht86

visual-retrieval-colpali/prepare_feed_deploy.py

@@ -4,7 +4,7 @@
 # We will use ColPali as the model to extract patch vectors from images of pdf pages.
 # At query time, we use MaxSim to retrieve and/or (based on the configuration) rank the page results.
 #
-# To see the application in action, visit TODO:
+# To see the application in action, visit https://huggingface.co/spaces/vespa-engine/colpali-vespa-visual-retrieval
 #
 # The web application is written in FastHTML, meaning the complete application is written in python.
 #
@@ -169,20 +169,45 @@
 links = []
 url_to_year = {}
 
-# Find all 'div's with id starting with 'year-'
-for year_div in soup.find_all("div", id=lambda x: x and x.startswith("year-")):
-    year_id = year_div.get("id", "")
+# Find all the years
+for year_section in soup.find_all("section", attrs={"data-name": "report-year"}):
+    year_id = year_section.get("data-filter-section", "")
     year = year_id.replace("year-", "")
-
-    # Within this div, find all 'a' elements with the specific classes
-    for a_tag in year_div.select("a.button.button--download-secondary[href]"):
-        href = a_tag["href"]
-        full_url = urljoin(url, href)
-        # exclude non-pdf links
-        if full_url.endswith(".pdf"):
-            links.append(full_url)
-            url_to_year[full_url] = year
-links, url_to_year
+    
+    # Find the yearly report section
+    report_section = year_section.find("div", attrs={
+        "data-filter-section": "year",
+        "data-name": "report-type"
+    })
+    if not report_section:
+        continue
+        
+    # Get the first link
+    report_link = report_section.select_one("ul.link-list a")
+    if not report_link:
+        continue
+        
+    report_url = urljoin(url, report_link["href"])
+    
+    # Visit the report page to find PDF download links
+    try:
+        report_response = requests.get(report_url)
+        report_response.raise_for_status()
+        report_soup = BeautifulSoup(report_response.text, "html.parser")
+        
+        # Find only the first PDF download link with the specific class
+        pdf_link = report_soup.select_one("a.btn.btn-secondary[data-right-icon='download']")
+        if not pdf_link or not pdf_link["href"].endswith(".pdf"):
+            continue
+            
+        pdf_url = urljoin(report_url, pdf_link["href"])
+        links.append(pdf_url)
+        url_to_year[pdf_url] = year
+        print(f"Found PDF: {pdf_url} (Year: {year})")
+    except Exception as e:
+        print(f"Error fetching report page {report_url}: {e}")
+
+print(f"Found {len(links)} PDF links")
 # -
 
 # Limit the number of PDFs to download
@@ -471,6 +496,7 @@ def generate_queries(image, prompt_text, pydantic_model):
 
 # +
 # write title, url, page_no, text, queries, not image to JSON
+os.makedirs("output", exist_ok=True)
 with open("output/pdf_pages.json", "w") as f:
     to_write = [{k: v for k, v in pdf.items() if k != "image"} for pdf in pdf_pages]
     json.dump(to_write, f, indent=2)
@@ -590,7 +616,6 @@ def float_to_binary_embedding(float_query_embedding: dict) -> dict:
 
 
 # Save vespa_feed to vespa_feed.json
-os.makedirs("output", exist_ok=True)
 with open("output/vespa_feed.json", "w") as f:
     vespa_feed_to_save = []
     for page in vespa_feed:

visual-retrieval-colpali/pyproject.toml

@@ -3,7 +3,7 @@ name = "visual-retrieval-colpali"
 version = "0.1.0"
 description = "Visual retrieval with ColPali"
 readme = "README.md"
-requires-python = ">=3.10, <3.13" 
+requires-python = ">=3.10, <3.12"
 license = { text = "Apache-2.0" }
 dependencies = [
     "python-fasthtml",
@@ -20,23 +20,19 @@ dependencies = [
     "shad4fast>=1.2.1",
     "google-generativeai>=0.7.2",
     "spacy",
-    "pip"
+    "pip",
 ]
 
 # dev-dependencies
 [project.optional-dependencies]
-dev = [
-    "ruff",
-    "python-dotenv",
-    "huggingface_hub[cli]"
-]
+dev = ["ruff", "python-dotenv", "huggingface_hub[cli]"]
 feed = [
     "ipykernel",
     "jupytext",
     "pydantic",
     "beautifulsoup4",
     "pdf2image",
-    "google-generativeai"
+    "google-generativeai",
 ]
 [tool.ruff]
 # Exclude a variety of commonly ignored directories.
@@ -115,4 +111,4 @@ docstring-code-format = false
 #
 # This only has an effect when the `docstring-code-format` setting is
 # enabled.
-docstring-code-line-length = "dynamic"
+docstring-code-line-length = "dynamic"