Seed Dataset Preparation
Download Recipe
1 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 # SPDX-License-Identifier: Apache-2.0 3 # /// script 4 # requires-python = ">=3.10" 5 # dependencies = [ 6 # "datasets", 7 # "pymupdf", 8 # "pandas", 9 # "pyarrow", 10 # ] 11 # /// 12 """Long-Document Understanding Seed Dataset Preparation 13 14 This script uses HuggingFace's FinePDFs dataset (HuggingFaceFW/finepdfs) as 15 an example data source to demonstrate how to prepare seed data for the rest 16 of the recipes. It downloads the original PDFs, renders each page to a PNG 17 image, and produces three seed parquet files: 18 19 1. **per-page seed** – one row per page, with a ``png_images_base64`` 20 column containing a JSON array with a single base64-encoded PNG of 21 that page. Suitable for single-page recipes (02 through 06). 22 2. **windowed seed** – one row per window of consecutive pages, with a 23 ``png_images_base64`` column containing a JSON array of base64-encoded 24 PNGs for the pages in that window. The window size adapts to document 25 length (2 pages for short documents up to 8 for long ones). 26 Suitable for the multi-page windowed recipe (07). 27 3. **whole-document seed** – one row per document, with a 28 ``png_images_base64`` column containing a JSON array of base64-encoded 29 PNGs for all pages. Suitable for the whole-document recipe (08). 30 31 Prerequisites: 32 - Internet access to download PDFs from their original URLs. 33 34 Run: 35 # Prepare seeds from 10 English PDFs (default) 36 uv run 01-seed-dataset-preparation.py --output-dir ./seed_data 37 38 # Prepare seeds from 50 PDFs 39 uv run 01-seed-dataset-preparation.py --output-dir ./seed_data --num-docs 50 40 41 # Use a different language subset 42 uv run 01-seed-dataset-preparation.py --output-dir ./seed_data --subset fra_Latn 43 44 # Skip documents that fail to download (default behaviour) and set 45 # a custom timeout 46 uv run 01-seed-dataset-preparation.py --output-dir ./seed_data --timeout 30 47 48 # For help 49 uv run 01-seed-dataset-preparation.py --help 50 """ 51 52 from __future__ import annotations 53 54 import base64 55 import json 56 import logging 57 import os 58 import urllib.request 59 from argparse import ArgumentParser 60 from pathlib import Path 61 62 import fitz # pymupdf 63 import pandas as pd 64 from datasets import load_dataset 65 66 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") 67 log = logging.getLogger(__name__) 68 69 logging.getLogger("httpx").setLevel(logging.WARNING) 70 logging.getLogger("httpcore").setLevel(logging.WARNING) 71 logging.getLogger("huggingface_hub").setLevel(logging.WARNING) 72 logging.getLogger("datasets").setLevel(logging.WARNING) 73 logging.getLogger("fsspec").setLevel(logging.WARNING) 74 75 DPI = 144 76 FINEPDFS_REPO = "HuggingFaceFW/finepdfs" 77 78 79 def download_pdf(url: str, timeout: int = 20) -> bytes | None: 80 """Download a PDF from *url*, returning raw bytes or None on failure.""" 81 try: 82 req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"}) 83 with urllib.request.urlopen(req, timeout=timeout) as resp: 84 return resp.read() 85 except Exception as exc: 86 log.warning("Failed to download %s: %s", url, exc) 87 return None 88 89 90 def render_pages(pdf_bytes: bytes, dpi: int = DPI) -> list[bytes]: 91 """Render every page of *pdf_bytes* to PNG, returning a list of raw PNG bytes.""" 92 doc = fitz.open(stream=pdf_bytes, filetype="pdf") 93 pages: list[bytes] = [] 94 for page in doc: 95 pix = page.get_pixmap(dpi=dpi) 96 pages.append(pix.tobytes("png")) 97 doc.close() 98 return pages 99 100 101 def png_to_base64(png_bytes: bytes) -> str: 102 """Encode raw PNG bytes as a base64 string.""" 103 return base64.b64encode(png_bytes).decode("ascii") 104 105 106 def adaptive_window_size(n_pages: int) -> int: 107 """Choose a window size that scales with document length. 108 109 Short documents get small windows (2 pages) so multi-page questions 110 remain feasible; longer documents get larger windows (up to 8) to 111 cover more context per seed row. 112 """ 113 if n_pages > 10 and n_pages < 20: 114 return 3 115 elif n_pages > 20 and n_pages < 30: 116 return 4 117 elif n_pages > 30 and n_pages < 40: 118 return 5 119 elif n_pages > 40 and n_pages < 50: 120 return 6 121 elif n_pages > 50 and n_pages < 60: 122 return 7 123 elif n_pages > 60: 124 return 8 125 return 2 126 127 128 def main() -> None: 129 parser = ArgumentParser(description="Prepare seed parquets from FinePDFs") 130 parser.add_argument( 131 "--output-dir", 132 type=str, 133 required=True, 134 help="Directory for output parquet files", 135 ) 136 parser.add_argument( 137 "--num-docs", 138 type=int, 139 default=10, 140 help="Number of PDF documents to process (default: 10)", 141 ) 142 parser.add_argument( 143 "--subset", 144 type=str, 145 default="eng_Latn", 146 help="FinePDFs language subset (default: eng_Latn)", 147 ) 148 parser.add_argument( 149 "--timeout", 150 type=int, 151 default=20, 152 help="HTTP download timeout in seconds (default: 20)", 153 ) 154 parser.add_argument( 155 "--dpi", 156 type=int, 157 default=DPI, 158 help=f"Render resolution in DPI (default: {DPI})", 159 ) 160 parser.add_argument( 161 "--max-pages", 162 type=int, 163 default=50, 164 help="Skip documents with more pages than this (default: 50)", 165 ) 166 parser.add_argument( 167 "--min-window-pages", 168 type=int, 169 default=2, 170 help="Minimum pages in a window; documents shorter than this are skipped for windowed output (default: 2)", 171 ) 172 args = parser.parse_args() 173 174 output_dir = Path(args.output_dir) 175 output_dir.mkdir(parents=True, exist_ok=True) 176 177 log.info( 178 "Streaming %d documents from %s (subset=%s)", 179 args.num_docs, 180 FINEPDFS_REPO, 181 args.subset, 182 ) 183 184 ds = load_dataset( 185 FINEPDFS_REPO, 186 name=args.subset, 187 split="train", 188 streaming=True, 189 ) 190 191 per_page_rows: list[dict] = [] 192 windowed_rows: list[dict] = [] 193 whole_doc_rows: list[dict] = [] 194 195 docs_processed = 0 196 for row in ds: 197 if docs_processed >= args.num_docs: 198 break 199 200 doc_id = row.get("id", f"doc_{docs_processed:06d}") 201 url = row["url"] 202 date = row.get("date", "") 203 204 pdf_bytes = download_pdf(url, timeout=args.timeout) 205 if pdf_bytes is None: 206 continue 207 208 try: 209 page_pngs = render_pages(pdf_bytes, dpi=args.dpi) 210 except Exception as exc: 211 log.warning("Failed to render %s: %s", url, exc) 212 continue 213 214 if len(page_pngs) == 0: 215 log.warning("No pages rendered for %s, skipping", url) 216 continue 217 218 if len(page_pngs) > args.max_pages: 219 log.info( 220 "Skipping %s (%d pages > --max-pages %d)", 221 url, 222 len(page_pngs), 223 args.max_pages, 224 ) 225 continue 226 227 page_b64s: list[str] = [] 228 229 for page_idx, png_bytes in enumerate(page_pngs): 230 b64 = png_to_base64(png_bytes) 231 page_b64s.append(b64) 232 233 per_page_rows.append( 234 { 235 "id": doc_id, 236 "url": url, 237 "date": date, 238 "page_number": page_idx, 239 "total_pages": len(page_pngs), 240 "png_images_base64": json.dumps([b64]), 241 } 242 ) 243 244 whole_doc_rows.append( 245 { 246 "id": doc_id, 247 "url": url, 248 "date": date, 249 "total_pages": len(page_pngs), 250 "png_images_base64": json.dumps(page_b64s), 251 } 252 ) 253 254 n_pages = len(page_b64s) 255 win_size = adaptive_window_size(n_pages) 256 n_windows = n_pages // win_size 257 for i in range(n_windows): 258 win_start = i * win_size 259 win_end = win_start + win_size 260 if win_end - win_start < args.min_window_pages: 261 break 262 windowed_rows.append( 263 { 264 "id": doc_id, 265 "url": url, 266 "date": date, 267 "total_pages": n_pages, 268 "start_page": win_start, 269 "end_page": win_end, 270 "window_size": win_end - win_start, 271 "png_images_base64": json.dumps(page_b64s[win_start:win_end]), 272 } 273 ) 274 275 docs_processed += 1 276 log.info( 277 "[%d/%d] %s — %d pages", 278 docs_processed, 279 args.num_docs, 280 url, 281 len(page_pngs), 282 ) 283 284 if not per_page_rows: 285 log.error("No documents were successfully processed. Exiting.") 286 return 287 288 per_page_path = output_dir / "seed_per_page.parquet" 289 windowed_path = output_dir / "seed_windowed.parquet" 290 whole_doc_path = output_dir / "seed_whole_document.parquet" 291 292 pd.DataFrame(per_page_rows).to_parquet(per_page_path, index=False) 293 if windowed_rows: 294 pd.DataFrame(windowed_rows).to_parquet(windowed_path, index=False) 295 pd.DataFrame(whole_doc_rows).to_parquet(whole_doc_path, index=False) 296 297 log.info("Per-page seed: %s (%d rows)", per_page_path, len(per_page_rows)) 298 log.info("Windowed seed: %s (%d rows)", windowed_path, len(windowed_rows)) 299 log.info("Whole-document seed: %s (%d rows)", whole_doc_path, len(whole_doc_rows)) 300 301 302 if __name__ == "__main__": 303 main() 304 # Force-exit to avoid hanging on background threads from datasets/fsspec. 305 os._exit(0)