# Synthetic Data Generation Pipeline End-to-end pipeline for generating the `screenshot-training-natural-filtered-v2` training dataset used to fine-tune `Qwen3-VL-Embedding-2B` for visual document retrieval. The pipeline produces ~115K high-quality query→screenshot-chunk pairs with hard negatives, starting from raw Wikipedia screenshot tiles. --- ## Pipeline Overview ``` ┌─────────────────────────────────────────────────────────────────────┐ │ VISUAL QUERY PIPELINE │ │ │ │ Wikipedia pages (kiwix_tiles/) │ │ │ │ │ ▼ │ │ ① generate_query_pairs.py Gemini reads screenshot chunks, │ │ │ generates Q/A pairs │ │ ▼ │ │ ② filter_self_contained.py GPT-4o removes non-self- │ │ │ contained queries │ │ ▼ │ │ ③ mine_hard_negatives.py Search API retrieves confusable │ │ │ chunks as hard negatives │ │ ▼ │ │ ④ filter_hard_negatives_vqa.py VLM removes false negatives │ │ │ (chunks that actually answer │ │ │ the query) │ │ ▼ │ │ ⑤ clean_queries_simpleqa_style.py Score naturalness + factoid │ │ │ style fit via Gemini │ │ ▼ │ │ ⑥ export_natural_filtered_v2.py Keep only naturalness≥4 & │ │ │ style_fit≥4 │ │ ▼ │ │ ⑦ split_first5_chunks.py 90/5/5 train/eval/test split │ │ │ │ │ ▼ │ │ ⑧ prepare_hf_dataset.py Package for Hugging Face │ │ package_hf_image_shards.py │ │ upload_hf_dataset.py │ └─────────────────────────────────────────────────────────────────────┘ ┌─────────────────────────────────────────────────────────────────────┐ │ TEXT WARMUP PIPELINE │ │ │ │ Text passage DB (text_baseline.db) │ │ │ │ │ ▼ │ │ ① generate_text_query_pairs.py Gemini reads text passages, │ │ │ generates Q/A pairs │ │ ▼ │ │ ② filter_self_contained.py Same filter as visual pipeline │ │ │ │ │ ▼ │ │ ③ mine_text_hard_negatives.py Text search API retrieves │ │ │ confusable passages │ │ ▼ │ │ ④ filter_text_hard_negatives_llm.py LLM removes false negatives │ │ │ │ Output: text-qa-pair dataset (used for --text-warmup-steps) │ └─────────────────────────────────────────────────────────────────────┘ ``` --- ## Prerequisites ### Data sources | Resource | Purpose | Approx. size | |----------|---------|-------------| | `kiwix_tiles/` directory | Wikipedia screenshot tiles with `index.jsonl` | ~500 GB | | `text_baseline.db` (SQLite) | Wikipedia text passages for text warmup | ~20 GB | | Search API (`:30888`) | Hard negative mining (retrieval over the tile index) | running service | ### API keys | Key | Used by | Purpose | |-----|---------|---------| | Google Cloud ADC | `generate_query_pairs.py`, `generate_text_query_pairs.py`, `clean_queries_simpleqa_style.py` | Gemini models via Vertex AI | | `OPENAI_API_KEY` | `filter_self_contained.py`, `filter_hard_negatives_vqa.py` | GPT-4o for query filtering and VQA false-negative removal | ```bash # Google Cloud (one-time setup) gcloud auth application-default login # OpenAI export OPENAI_API_KEY="sk-proj-..." export OPENAI_BASE_URL="https://us.api.openai.com/v1" # if your key requires this ``` --- ## Step-by-Step Reproduction ### Step 1: Generate Visual Query Pairs Send screenshot chunks to Gemini, which generates factual Q/A pairs. **Script:** `generate_query_pairs.py` **What it does:** 1. Loads `index.jsonl` from the kiwix_tiles directory 2. Filters out infrastructure pages (disambiguation, elections, lists, etc.) 3. For each eligible page, picks a random chunk from the first 70% of chunks 4. Sends the chunk image to Gemini with a detailed prompt 5. Gemini generates a Q/A pair with source_type and subject labels 6. Post-generation filters reject layout-bound or unnatural questions **Single batch:** ```bash uv run python generate_query_pairs.py \ --tiles-dir /opt/dlami/nvme/kiwix_tiles \ --num-pages 2000 \ --model gemini-3.1-flash-lite-preview \ --output batches/batch_000.jsonl \ --seed 0 ``` **Full-scale generation (120 non-overlapping batches):** ```bash bash run_generate_batches.sh /opt/dlami/nvme/kiwix_tiles 0 119 gemini-3.1-flash-lite-preview ``` This produces ~195K raw pairs across 120 batch files. **Output format:** ```json { "query": "In what year did Henryk Skarżyński perform the first cochlear implantation in Poland?", "answer": "1992", "source_sentence": "He performed the first operation of cochlear implantation in Poland...", "source_type": "prose", "subject": "medicine", "chunk_path": "shard_400/shard_00010/3316848.png.tiles/chunk_0000_00.png", "url": "https://en.wikipedia.org/wiki/Henryk_Skarżyński", "title": "Henryk Skarżyński", "chunk_index": 0, "tiles_dir": "shard_400/shard_00010/3316848.png.tiles" } ``` **Cost estimate:** ~$100–200 for 195K pairs with Flash Lite, ~$2000 with Pro. **Merge all batches:** ```bash cat batches/batch_*.jsonl > raw_query_pairs.jsonl wc -l raw_query_pairs.jsonl # expect ~195K ``` ### Step 2: Filter Self-Contained Queries Remove queries that reference the page layout ("in the table") or use vague references ("the team", "the film") without naming the entity. **Script:** `filter_self_contained.py` ```bash OPENAI_API_KEY=sk-... uv run python filter_self_contained.py \ --input raw_query_pairs.jsonl \ --output filtered_query_pairs.jsonl \ --model gpt-4o ``` **Expected results:** - Drop rate: ~15% (29K dropped from 195K) - Cost: ~$23 with GPT-4o - Runtime: ~8 minutes **Test run first:** ```bash OPENAI_API_KEY=sk-... uv run python filter_self_contained.py \ --input raw_query_pairs.jsonl \ --output /dev/null \ --model gpt-4o \ --test-first 100 ``` ### Step 3: Mine Hard Negatives For each (query, positive_chunk) pair, retrieve top-K candidates from the search API. Non-positive results become hard negatives. **Script:** `mine_hard_negatives.py` (already in this repo) **Prerequisite:** Search API running at `localhost:30888` (see `serve/` and `deploy/`). ```bash uv run python mine_hard_negatives.py \ --input filtered_query_pairs.jsonl \ --output filtered_query_pairs_hn.jsonl \ --num-negatives 7 \ --n-docs 50 \ --filter-mode margin \ --margin 0.95 ``` **Output adds:** ```json { "query": "...", "chunk_path": "...", "neg_chunk_paths": ["/path/to/neg1.png", "/path/to/neg2.png", "..."], "retrieve_top20": [{"rank": 1, "path": "...", "score": 0.61}, "..."], "positive_score": 0.65, "positive_rank": 1 } ``` ### Step 4: Filter Hard Negatives (VQA-based) Remove false negatives — mined "negatives" that actually do answer the query. A VLM answers the query from each candidate image; if it gets the answer right, that candidate is NOT a true hard negative and is removed. **Script:** `filter_hard_negatives_vqa.py` / `run_filter_hard_negatives_chunks.py` For large files, use the chunked runner for resumability: ```bash uv run python run_filter_hard_negatives_chunks.py \ --input-jsonl filtered_query_pairs_hn.jsonl \ --output-dir training/data/filtered_hn_chunks \ --chunk-size 10000 \ --skip-existing ``` Each chunk produces: - `filtered_hn.jsonl` — the cleaned hard negatives - `candidate_reviews.jsonl` — per-candidate VLM judgments - `summary.json` — statistics ### Step 5: Clean Queries (Naturalness Scoring) Score each query on naturalness (1–5) and factoid style fit (1–5) using Gemini. This step does NOT look at images — only the query text is judged. **Script:** `clean_queries_simpleqa_style.py` ```bash uv run python clean_queries_simpleqa_style.py \ --model gemini-2.0-flash-001 \ --target-count 50000 \ --batch-size 20 \ --concurrency 8 \ --dedupe-query \ --output training/data/cleaned/simpleqa_style_cleaned.jsonl \ --reviews-output training/data/cleaned/simpleqa_style_cleaned.reviews.jsonl \ --summary-output training/data/cleaned/simpleqa_style_cleaned.summary.json ``` ### Step 6: Export Strict Retained Rows Keep only rows with `naturalness >= 4` and `simpleqa_style_fit >= 4`. **Script:** `export_natural_filtered_v2.py` ```bash uv run python export_natural_filtered_v2.py \ --reviews training/data/cleaned/simpleqa_style_cleaned.reviews.jsonl \ --output-dir training/data/natural_filtered_v2 \ --output-name filtered_hn.jsonl \ --min-naturalness 4 \ --min-style-fit 4 ``` **Expected:** ~77% keep rate → ~115K rows from 150K input. ### Step 7: Split into Train/Eval/Test **Script:** `split_first5_chunks.py` ```bash uv run python split_first5_chunks.py \ --inputs training/data/natural_filtered_v2/filtered_hn.jsonl \ --output-dir training/data/natural_filtered_v2/split \ --seed 42 \ --train-ratio 0.9 \ --eval-ratio 0.05 \ --test-ratio 0.05 ``` Output: - `split/train_hn.jsonl` (~104K) - `split/eval_hn.jsonl` (~5.8K) - `split/test_hn.jsonl` (~5.8K) ### Step 8: Package and Upload to Hugging Face ```bash # Materialize images (hardlinks from kiwix_tiles) uv run python prepare_hf_dataset.py \ --split-dir training/data/natural_filtered_v2/split \ --image-root /opt/dlami/nvme/kiwix_tiles \ --output-dir /opt/dlami/nvme/hf_export/screenshot-training-natural-filtered-v2 \ --repo-id Chrisyichuan/screenshot-training-natural-filtered-v2 # Create tar shards (1000 shards, ~93.5 GB total) uv run python package_hf_image_shards.py \ --source-dir /opt/dlami/nvme/hf_export/screenshot-training-natural-filtered-v2 \ --output-dir /opt/dlami/nvme/hf_export_sharded/screenshot-training-natural-filtered-v2 \ --overwrite # Upload uv run python upload_hf_dataset.py \ --repo-id Chrisyichuan/screenshot-training-natural-filtered-v2 \ --local-dir /opt/dlami/nvme/hf_export_sharded/screenshot-training-natural-filtered-v2 \ --repo-type dataset \ --skip-create ``` --- ## Text Warmup Data Pipeline The text warmup pipeline generates `text-qa-pair` data used for the `--text-warmup-steps` phase of training. It follows a parallel track. ### Generate Text Query Pairs **Script:** `generate_text_query_pairs.py` ```bash uv run python generate_text_query_pairs.py \ --db-path /opt/dlami/nvme/text_embeddings_1024/text_baseline.db \ --num-articles 52000 \ --model gemini-3.1-flash-lite-preview \ --output text_query_pairs.jsonl \ --min-article-chunks 6 \ --min-paragraph-words 60 ``` The script: 1. Queries the SQLite DB for articles with ≥6 chunks 2. For each article, selects a chunk with a long natural prose paragraph 3. Sends the passage to Gemini with 5-shot examples 4. Filters for natural, self-contained prose questions 5. Verifies the supporting span exists verbatim in the passage **Expected yield:** ~40% (21K pairs from 52K articles). **Cost:** ~$110 with Flash Lite. ### Filter + Mine + Package ```bash # Filter self-contained (same script as visual pipeline) OPENAI_API_KEY=sk-... uv run python filter_self_contained.py \ --input text_query_pairs.jsonl \ --output text_query_pairs_filtered.jsonl \ --model gpt-4o # Mine text hard negatives (requires text search API on :30889) uv run python mine_text_hard_negatives.py \ --input text_query_pairs_filtered.jsonl \ --output text_query_pairs_hn.jsonl # Filter text hard negatives uv run python run_filter_text_hard_negatives_chunks.py \ --input-jsonl text_query_pairs_hn.jsonl \ --output-dir training/data/text_filtered_hn_chunks ``` --- ## Script Inventory ### Query Generation (Steps 1–2) | Script | Input | Output | API | |--------|-------|--------|-----| | `generate_query_pairs.py` | kiwix_tiles/ + index.jsonl | raw visual Q/A pairs (JSONL) | Gemini (Vertex AI) | | `generate_text_query_pairs.py` | text_baseline.db (SQLite) | raw text Q/A pairs (JSONL) | Gemini + OpenAI fallback | | `filter_self_contained.py` | any Q/A JSONL | filtered JSONL (non-self-contained removed) | OpenAI (GPT-4o) | | `run_generate_batches.sh` | kiwix_tiles/ | batched generation wrapper | — | ### Hard Negative Mining (Step 3) | Script | Input | Output | API | |--------|-------|--------|-----| | `mine_hard_negatives.py` | filtered Q/A JSONL | JSONL with `neg_chunk_paths` | Search API (:30888) | | `mine_text_hard_negatives.py` | filtered text Q/A JSONL | JSONL with `neg_passages` | Text search API (:30889) | ### Hard Negative Filtering (Step 4) | Script | Input | Output | API | |--------|-------|--------|-----| | `filter_hard_negatives_vqa.py` | HN JSONL | cleaned HN JSONL | OpenAI/Gemini (VLM) | | `run_filter_hard_negatives_chunks.py` | large HN JSONL | chunked output dir | — | | `filter_text_hard_negatives_llm.py` | text HN JSONL | cleaned text HN JSONL | OpenAI/Gemini (LLM) | | `run_filter_text_hard_negatives_chunks.py` | large text HN JSONL | chunked output dir | — | ### Query Cleaning & Export (Steps 5–7) | Script | Input | Output | API | |--------|-------|--------|-----| | `clean_queries_simpleqa_style.py` | filtered HN JSONL | reviews JSONL with scores | Gemini | | `export_natural_filtered_v2.py` | reviews JSONL | strict-filtered JSONL | — | | `split_first5_chunks.py` | filtered JSONL | train/eval/test split | — | ### Dataset Packaging (Step 8) | Script | Input | Output | API | |--------|-------|--------|-----| | `prepare_hf_dataset.py` | split dir + image root | HF dataset folder | — | | `package_hf_image_shards.py` | HF dataset folder | tar-sharded images | — | | `upload_hf_dataset.py` | sharded dataset | HF Hub upload | HF API | | `extract_hf_image_shards.py` | downloaded shards | extracted images/ | — | --- ## Cost Estimates (Full Pipeline) | Step | Model | Volume | Est. Cost | |------|-------|--------|-----------| | Generate visual queries | gemini-3.1-flash-lite | 195K pairs | ~$150 | | Filter self-contained | gpt-4o | 195K queries | ~$23 | | Generate text queries | gemini-3.1-flash-lite + gpt-4o-mini fallback | 52K articles | ~$110 | | Filter text self-contained | gpt-4o | 21K queries | ~$3 | | Filter HN (VQA) | gpt-4o / gemini | 150K × 7 candidates | ~$200 | | Clean queries (naturalness) | gemini-2.0-flash | 150K queries | ~$15 | | **Total** | | | **~$500** | --- ## Data Formats ### Raw Q/A pair (after Step 1) ```json { "query": "What is the population of Tokyo?", "answer": "13.96 million", "source_sentence": "Tokyo has a population of 13.96 million...", "source_type": "prose", "subject": "geography", "chunk_path": "shard_400/shard_00010/350170.png.tiles/chunk_0000_00.png", "url": "https://en.wikipedia.org/wiki/Tokyo", "title": "Tokyo", "chunk_index": 0, "tiles_dir": "shard_400/shard_00010/350170.png.tiles" } ``` ### With hard negatives (after Step 3) ```json { "query": "...", "chunk_path": "...", "neg_chunk_paths": ["/path/to/neg1.png", "/path/to/neg2.png"], "retrieve_top20": [{"rank": 1, "path": "...", "score": 0.61}], "positive_score": 0.65, "positive_rank": 1 } ``` ### Published HF format (after Step 8) ```json { "query": "...", "chunk_path": "images/shard_000/...", "neg_chunk_paths": ["images/shard_001/...", "images/shard_002/..."], "source_positive_rank": 1, "source_positive_score": 0.63 } ``` --- ## Quality Control ### Page-level filtering (Step 1) - Skip infrastructure pages (disambiguation, category, template, portal) - Skip low-quality content (elections, census, track listings, episode lists) - Require page_height ≥ 3000px, ≥1 tile, complete capture ### Query-level filtering | Filter | What it catches | Drop rate | |--------|----------------|-----------| | `is_natural_question()` (Step 1) | Layout references, dangling entities, truncated answers | ~30% of raw Gemini output | | `filter_self_contained.py` (Step 2) | Unnamed subjects, document-structure references | ~15% | | `clean_queries_simpleqa_style.py` (Step 5) | Templatic phrasing, poor naturalness | ~23% | ### Hard negative quality (Step 4) The VQA filter ensures hard negatives are truly *confusable but wrong*: 1. VLM tries to answer the query from the candidate image 2. A judge grades the answer as CORRECT / WRONG / CANNOT_ANSWER 3. Only WRONG or CANNOT_ANSWER candidates are kept as hard negatives 4. If the positive image itself fails the VQA check, the entire row is skipped --- ## Differences from Published Dataset The published `screenshot-training-natural-filtered-v2` was generated with: - **Visual queries:** `gemini-3.1-flash-lite-preview` (120 batches, 195K raw → 165K filtered) - **Self-contained filter:** `gpt-4o` (15% drop rate, zero false keeps) - **Naturalness cleaning:** `gemini-2.0-flash-001` - **Naturalness thresholds:** `naturalness ≥ 4`, `simpleqa_style_fit ≥ 4` - **Final size:** 115,593 rows → 104K train / 5.8K eval / 5.8K test