/train/generate_query_pairs.py

#!/usr/bin/env python3
"""
Generate synthetic query-chunk pairs for contrastive learning.

Samples informative Wikipedia pages from kiwix tiles, sends screenshot chunks
to Gemini to generate factual Q&A pairs with source_type labels.

Output: JSONL with {query, answer, source_type, subject, source_sentence,
                     chunk_path, url, title, chunk_index, tiles_dir}

Prerequisites:
  - kiwix_tiles directory with Wikipedia screenshot tiles and index.jsonl
  - Google Cloud ADC (gcloud auth application-default login) for Vertex AI,
    OR set GOOGLE_API_KEY for direct Gemini API access

Usage:
  python generate_query_pairs.py \
      --tiles-dir /path/to/kiwix_tiles \
      --num-pages 1000 \
      --output batches/batch_000.jsonl

  # Batched generation (non-overlapping slices):
  python generate_query_pairs.py \
      --tiles-dir /path/to/kiwix_tiles \
      --batch-index 0 --total-batches 100 \
      --num-pages 2000 \
      --output batches/batch_000.jsonl

Ported from Vis-RAG/agent/scripts/contrastive/generate_query_pairs.py
"""

import json
import os
import re
import random
import base64
import asyncio
import argparse
import time
from pathlib import Path
from io import BytesIO
from collections import Counter

from PIL import Image
from google import genai
from google.genai.types import HttpOptions, GenerateContentConfig

# Vertex AI config — requires gcloud ADC (gcloud auth application-default login)
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", "wise-coyote-478119-h0")
os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "true")
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "global")

MODEL_PRICING = {
    "gemini-3.1-pro-preview": (1.25, 10.00),
    "gemini-2.5-pro-preview-03-25": (1.25, 10.00),
    "gemini-2.5-pro": (1.25, 10.00),
    "gemini-3.1-flash-lite-preview": (1.00, 4.00),
    "gemini-2.0-flash-001": (0.10, 0.40),
    "gemini-2.0-flash": (0.10, 0.40),
}
DEFAULT_MODEL = "gemini-3.1-pro-preview"
MAX_CONCURRENT = 40
IMG_JPEG_QUALITY = 85

# ── Page filtering ───────────────────────────────────────────────────
SKIP_PATTERNS = [
    r"disambiguation",
    r"category:",
    r"template:",
    r"wikipedia:",
    r"portal:",
    r"file:",
    r"help:",
    r"talk:",
    r"module:",
    r"draft:",
    r"_deaths$",
    r"_births$",
]

SKIP_CONTENT_PATTERNS = [
    r"\belection\b",
    r"\belections\b",
    r"\breferendum\b",
    r"\bprimary\b",
    r"\bby-election\b",
    r"\bcouncil election\b",
    r"^list of ",
    r"^lists of ",
    r"\bdiscography\b",
    r"\bfilmography\b",
    r"\btrack listing\b",
    r"_discography$",
    r"_filmography$",
    r"\bseason\b.*\bleague\b",
    r"\bleague season\b",
    r"\bfootball league\b",
    r"\bnba season\b",
    r"\bnfl season\b",
    r"\bcensus\b",
    r"\bdemographic\b",
    r"^list of .* episodes",
    r"\bepisodes of\b",
    r"\bgovernors? of\b",
    r"\bmayors? of\b",
    r"\bprime ministers? of\b",
    r"\bcareer statistics\b",
    r"\bplayer statistics\b",
]

SKIP_RE = [re.compile(p, re.IGNORECASE) for p in SKIP_PATTERNS]
SKIP_CONTENT_RE = [re.compile(p, re.IGNORECASE) for p in SKIP_CONTENT_PATTERNS]

BAD_QUESTION_PATTERNS = [
    r"\baccording to the\b",
    r"\baccording to this\b",
    r"\bvisible\b",
    r"\btrack listing\b",
    r"\blisted (here|above|below|in the table)\b",
    r"\bthe table (shows|lists|above|below)\b",
    r"\bin the (following|above) table\b",
    r"\bshown in\b",
    r"^what is (the )?listed",
    r"^what (are|is) listed",
    r"\bpage (shows|lists|includes)\b",
    r"\bthe film\b(?!\s+[A-Z\"])",
    r"\bthe song\b(?!\s+[A-Z\"])",
    r"\bthe album\b(?!\s+[A-Z\"])",
    r"\bthe book\b(?!\s+[A-Z\"])",
    r"\bthe team\b(?!\s+[A-Z\"])",
    r"\bthe show\b(?!\s+[A-Z\"])",
    r"\bthe series\b(?!\s+[A-Z\"])",
    r"\bthe station\b(?!\s+[A-Z\"])",
    r"\bthe school\b(?!\s+[A-Z\"])",
    r"\bthe match\b(?!\s+[A-Z\"])",
    r"\bthe game\b(?!\s+[A-Z\"])",
    r"\bthe competition\b(?!\s+[A-Z\"])",
    r"\bthe episode\b(?!\s+[A-Z\"\d])",
    r"\bthe production\b(?!\s+[A-Z\"])",
    r"\bthe tournament\b(?!\s+[A-Z\d\"])",
    r"^(when|where|who|what|how|why) (was|is|were|did|does|has|have) (it|they|this|that|he|she)\b",
]
BAD_Q_RE = [re.compile(p, re.IGNORECASE) for p in BAD_QUESTION_PATTERNS]


def get_page_chunk_count(entry: dict, tiles_root: Path) -> int:
    cached = entry.get("_chunk_count")
    if cached is not None:
        return cached

    tiles_dir = tiles_root / entry["tiles_dir"]
    chunks_json = tiles_dir / "chunks.json"
    if not chunks_json.exists():
        entry["_chunk_count"] = 0
        return 0

    with open(chunks_json) as f:
        meta = json.load(f)

    chunk_count = len(meta.get("chunks", []))
    entry["_chunk_count"] = chunk_count
    return chunk_count


def is_informative_page(entry: dict) -> bool:
    if entry.get("page_height", 0) < 3000:
        return False
    if entry.get("num_tiles", 0) < 1:
        return False
    if not entry.get("complete", False):
        return False

    title_lower = entry["title"].lower()
    url_lower = entry.get("url", "").lower()
    check = title_lower + " " + url_lower

    for pat in SKIP_RE:
        if pat.search(check):
            return False
    for pat in SKIP_CONTENT_RE:
        if pat.search(check):
            return False

    return True


def is_natural_question(qa: dict) -> bool:
    q = qa.get("query", "")
    for pat in BAD_Q_RE:
        if pat.search(q):
            return False
    a = qa.get("answer", "")
    if a and a[-1] in ("→", "…", "–", "/", "(", ","):
        return False
    s = qa.get("source_sentence", "") or ""
    src_type = qa.get("source_type", "prose")
    if not s:
        return False
    if s.rstrip()[-1:] in ("(", ",", "–", "/", "→", "…"):
        return False
    if src_type == "prose" and len(s.split()) < 10:
        return False
    if src_type in ("infobox", "table") and len(s.split()) < 3:
        return False
    return True


def load_and_sample_pages(
    index_path: Path,
    n: int,
    batch_index: int = 0,
    total_batches: int = 1,
) -> list:
    MASTER_SEED = 0
    print(f"Loading index from {index_path}...")
    candidates: list = []

    with open(index_path) as f:
        for line in f:
            entry = json.loads(line)
            if is_informative_page(entry):
                candidates.append(entry)

    print(f"Total eligible: {len(candidates):,} pages")

    rng = random.Random(MASTER_SEED)
    rng.shuffle(candidates)

    slice_size = len(candidates) // total_batches
    start = batch_index * slice_size
    end = start + slice_size if batch_index < total_batches - 1 else len(candidates)
    pool = candidates[start:end]

    print(
        f"Batch {batch_index}/{total_batches}: pages [{start}:{end}] ({len(pool):,} in pool)"
    )
    selected = pool[:n] if n <= len(pool) else pool
    return selected


def filter_selected_pages_by_chunk_count(
    pages: list[dict],
    tiles_root: Path,
    min_page_chunks: int | None = None,
    max_page_chunks: int | None = None,
) -> list[dict]:
    if min_page_chunks is None and max_page_chunks is None:
        return pages

    kept = []
    for entry in pages:
        chunk_count = get_page_chunk_count(entry, tiles_root)
        if min_page_chunks is not None and chunk_count < min_page_chunks:
            continue
        if max_page_chunks is not None and chunk_count > max_page_chunks:
            continue
        kept.append(entry)
    return kept


def pick_random_chunk(entry: dict, tiles_root: Path) -> tuple:
    tiles_dir = tiles_root / entry["tiles_dir"]
    chunks_json = tiles_dir / "chunks.json"

    if not chunks_json.exists():
        return None, None

    with open(chunks_json) as f:
        meta = json.load(f)

    chunks = meta.get("chunks", [])
    if not chunks:
        return None, None

    usable = chunks[: max(1, int(len(chunks) * 0.7))]
    chunk = random.choice(usable)
    chunk_path = tiles_dir / chunk["file"]

    if not chunk_path.exists():
        return None, None

    return str(chunk_path), chunk["chunk_index"]


def encode_image(path: str) -> str:
    img = Image.open(path)
    buf = BytesIO()
    img.convert("RGB").save(buf, format="JPEG", quality=IMG_JPEG_QUALITY)
    return base64.b64encode(buf.getvalue()).decode()


QUERY_PROMPT = """\
You are generating a query–evidence pair for training a visual retrieval model over Wikipedia screenshot chunks.

TASK: Given this screenshot chunk, generate ONE factual question whose answer is explicitly and completely visible in this chunk.

━━━ STYLE — write questions like real search queries, not templates ━━━
Good questions sound like something a curious person would actually search for online.
Vary the phrasing — use "how much", "in what year", "which", "who", "where", "what caused", "how long", etc.

Here are examples of the STYLE we want (from SimpleQA benchmark):
  ✓ "How much money, in euros, was the surgeon held responsible for paying in the Olivia Puls case?"
  ✓ "In what city was the 2010 FIFA World Cup opening ceremony held?"
  ✓ "How many days did the 1906 San Francisco earthquake fire burn?"
  ✓ "Which award did Fullmetal Alchemist win at the American Anime Awards in 2007?"
  ✓ "Who was the first Black female judge appointed to the Cook County Circuit Court?"
  ✓ "What was the name of the ship that sank during the 1994 Estonia ferry disaster?"

━━━ EVIDENCE SOURCE — be diverse ━━━
You may draw from ANY visible content: prose text, infobox fields, table cells, image captions, diagrams, or photographs.
Do not always default to infobox — choose whichever source produces the most natural, interesting question.

━━━ HARD RULES ━━━
1. SELF-CONTAINED: The question must be fully understandable on its own — no page title, no external context needed.
   Every entity in the question must be named explicitly.

   ✗ "Who composed the music for the film?"          → missing film name
   ✗ "What is Rideaux's occupation?"                 → surname only, who is Rideaux?
   ✗ "On what date was Lerew awarded the DFC?"       → surname only + unexplained acronym
   ✗ "Which medication is listed as a SARM in the provided table?"  → depends on "the provided table"
   ✗ "Who played the actress in the 2013 film Horns?"  → "the actress" not identified
   ✗ "When was the SMN founded?"                     → unexplained acronym
   ✗ "Who was the spouse of John Houston?"           → too ambiguous, which John Houston?
   ✗ "Which cyclist placed second in the Tempo race?" → missing event/year context
   ✗ "What is listed in the infobox?"               → references page layout
   ✗ "Which mission is shown in the screenshot?"     → depends on visual layout
   ✗ "Which former Cleveland Indians player did the Seattle Mariners sign on December 20?" → missing year
   ✗ "Who was Kesha Rogers' opponent in the general election?"  → missing year and race context

   ✓ "Who composed the music for Once Upon a Time in Hong Kong?"
   ✓ "What is the occupation of Rabbi Shmuel Kamenetsky?"
   ✓ "In what year did photographer Clarence Rideaux found the agency PicturePerfect?"
   ✓ "On what date was RAF pilot Arthur Lerew awarded the Distinguished Flying Cross in World War II?"
   ✓ "Which former Cleveland Indians pitcher did the Seattle Mariners sign on December 20, 2004?"

2. EVIDENCE COMPLETE: The answer must be fully visible in this chunk — not guessed or inferred.
   The source sentence (S:) must be a complete, untruncated span.

3. DISTINCTIVE: Include enough specifics (names, dates, locations, titles) to distinguish this chunk from similar pages.

━━━ ANSWER ━━━
Prefer a single concise entity: name, date, place, number, title, or short phrase.

━━━ SKIP if any is true ━━━
- Raw vote counts, track listings, census tables, or episode lists
- Answer not fully visible or requires external context
- Cannot write a self-contained question naming all entities
- Source sentence is truncated or a fragment

Write exactly: SKIP

source_type: image | table | infobox | prose
subject: science | medicine | history | geography | technology | education | culture | politics | economics | biology | sports | entertainment | other

Output format (5 lines only):
Q: <natural, self-contained question>
A: <concise answer>
S: <verbatim complete span from the chunk>
T: <source_type>
C: <subject>"""


async def generate_qa(
    client: genai.Client,
    model: str,
    chunk_path: str,
    semaphore: asyncio.Semaphore,
    token_counter: dict,
) -> dict | None:
    async with semaphore:
        b64 = encode_image(chunk_path)

        contents = [
            {
                "role": "user",
                "parts": [
                    {"text": QUERY_PROMPT},
                    {"inline_data": {"mime_type": "image/jpeg", "data": b64}},
                ],
            }
        ]
        config = GenerateContentConfig(temperature=0.7, max_output_tokens=1024)

        for attempt in range(5):
            try:
                t0 = time.time()
                resp = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: client.models.generate_content(
                        model=model,
                        contents=contents,
                        config=config,
                    ),
                )
                elapsed = time.time() - t0

                usage = resp.usage_metadata
                if usage:
                    token_counter["input"] += getattr(usage, "prompt_token_count", 0)
                    token_counter["output"] += getattr(
                        usage, "candidates_token_count", 0
                    )
                    token_counter["calls"] += 1
                    token_counter["total_time"] += elapsed

                text = ""
                for p in resp.candidates[0].content.parts:
                    raw = getattr(p, "text", None)
                    if raw and not getattr(p, "thought", False):
                        text = raw.strip()

                if not text or text.strip() == "SKIP":
                    return None

                fields = {}
                for line in text.split("\n"):
                    line = line.strip()
                    for prefix, key in [
                        ("Q:", "query"),
                        ("A:", "answer"),
                        ("S:", "source_sentence"),
                        ("T:", "source_type"),
                        ("C:", "subject"),
                    ]:
                        if line.startswith(prefix):
                            fields[key] = line[len(prefix) :].strip()
                            break

                q = fields.get("query")
                a = fields.get("answer")
                if not q or not a or len(a) < 2:
                    return None

                qa = {
                    "query": q,
                    "answer": a,
                    "source_sentence": fields.get("source_sentence"),
                    "source_type": fields.get("source_type", "prose"),
                    "subject": fields.get("subject", "other"),
                }

                if not is_natural_question(qa):
                    return None

                return qa

            except Exception as e:
                err = str(e)
                if "429" in err or "RESOURCE_EXHAUSTED" in err:
                    wait = 2**attempt * 15 + random.uniform(1, 5)
                    print(f"  Rate limited, waiting {wait:.0f}s...")
                    await asyncio.sleep(wait)
                elif attempt < 4:
                    await asyncio.sleep(2)
                else:
                    print(f"  Failed after 5 attempts: {e}")
                    return None

    return None


async def main():
    parser = argparse.ArgumentParser(
        description="Generate synthetic query-chunk pairs from Wikipedia screenshot tiles"
    )
    parser.add_argument(
        "--tiles-dir",
        type=Path,
        required=True,
        help="Root directory of kiwix_tiles (containing index.jsonl)",
    )
    parser.add_argument("--num-pages", type=int, default=10)
    parser.add_argument("--output", type=str, default="query_pairs.jsonl")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--model",
        type=str,
        default=DEFAULT_MODEL,
        help="Gemini model (e.g. gemini-3.1-pro-preview, gemini-2.0-flash-001)",
    )
    parser.add_argument(
        "--batch-index",
        type=int,
        default=0,
        help="Which non-overlapping batch to process (0-based)",
    )
    parser.add_argument(
        "--total-batches",
        type=int,
        default=1,
        help="Total number of batches the candidate pool is divided into",
    )
    parser.add_argument("--max-concurrent", type=int, default=MAX_CONCURRENT)
    parser.add_argument("--postfilter-min-page-chunks", type=int, default=None)
    parser.add_argument("--postfilter-max-page-chunks", type=int, default=None)
    args = parser.parse_args()

    random.seed(args.seed)
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    tiles_root = args.tiles_dir
    index_path = tiles_root / "index.jsonl"
    if not index_path.exists():
        raise FileNotFoundError(f"index.jsonl not found at {index_path}")

    model = args.model
    price_input, price_output = MODEL_PRICING.get(model, (1.25, 10.00))

    print(
        f"Model: {model} via Vertex AI (project={os.environ.get('GOOGLE_CLOUD_PROJECT', 'N/A')})"
    )

    pages = load_and_sample_pages(
        index_path, args.num_pages, args.batch_index, args.total_batches
    )
    pages_before_postfilter = len(pages)
    pages = filter_selected_pages_by_chunk_count(
        pages,
        tiles_root,
        min_page_chunks=args.postfilter_min_page_chunks,
        max_page_chunks=args.postfilter_max_page_chunks,
    )
    if (
        args.postfilter_min_page_chunks is not None
        or args.postfilter_max_page_chunks is not None
    ):
        chunk_filter = []
        if args.postfilter_min_page_chunks is not None:
            chunk_filter.append(f"chunks>={args.postfilter_min_page_chunks}")
        if args.postfilter_max_page_chunks is not None:
            chunk_filter.append(f"chunks<={args.postfilter_max_page_chunks}")
        print(
            "Post-slice page filter: "
            + ", ".join(chunk_filter)
            + f" -> kept {len(pages)}/{pages_before_postfilter} pages"
        )
    print(f"Selected {len(pages)} pages\n")

    semaphore = asyncio.Semaphore(args.max_concurrent)
    token_counter = {"input": 0, "output": 0, "calls": 0, "total_time": 0.0}
    results = []

    client = genai.Client(http_options=HttpOptions(api_version="v1"))

    async def process_one(page, chunk_path, chunk_idx):
        qa = await generate_qa(client, model, chunk_path, semaphore, token_counter)
        if not qa:
            return None
        rel_path = str(Path(chunk_path).relative_to(tiles_root))
        return {
            **qa,
            "chunk_path": rel_path,
            "url": page["url"],
            "title": page["title"],
            "chunk_index": chunk_idx,
            "tiles_dir": page["tiles_dir"],
        }

    work_items = []
    for page in pages:
        chunk_path, chunk_idx = pick_random_chunk(page, tiles_root)
        if chunk_path is None:
            continue
        work_items.append((page, chunk_path, chunk_idx))

    print(f"Generating Q&A for {len(work_items)} chunks...\n")
    t_start = time.time()

    tasks = [asyncio.ensure_future(process_one(p, cp, ci)) for p, cp, ci in work_items]

    with open(args.output, "w") as outf:
        for coro in asyncio.as_completed(tasks):
            result = await coro
            if result:
                results.append(result)
                outf.write(json.dumps(result, ensure_ascii=False) + "\n")
                outf.flush()
                st = result.get("source_type", "?")
                subj = result.get("subject", "?")
                print(f"  [{st:8s}|{subj:14s}] {result['title']}")
                print(f"    Q: {result['query'][:90]}")
                print(f"    A: {result['answer'][:90]}")
                print()

    wall_time = time.time() - t_start

    type_dist = Counter(r["source_type"] for r in results)
    subj_dist = Counter(r["subject"] for r in results)

    in_tok = token_counter["input"]
    out_tok = token_counter["output"]
    calls = token_counter["calls"]
    cost = (in_tok / 1e6 * price_input) + (out_tok / 1e6 * price_output)

    print(f"\n{'=' * 60}")
    print(f"Wrote {len(results)} Q&A pairs to {args.output}")
    print(f"Source types: {dict(type_dist)}")
    print(f"Subjects:     {dict(subj_dist)}")
    print(f"\n--- Throughput & Cost ({model}) ---")
    print(f"  Calls:        {calls}")
    print(f"  Input tokens: {in_tok:,}  ({in_tok / max(calls, 1):.0f} avg/call)")
    print(f"  Output tokens:{out_tok:,}  ({out_tok / max(calls, 1):.0f} avg/call)")
    print(f"  Wall time:    {wall_time:.1f}s  ({wall_time / max(calls, 1):.1f}s/call)")
    print(f"  Est. cost:    ${cost:.4f} for {calls} calls")
    print(f"  Per 10 pairs: ${cost / max(calls, 1) * 10:.4f}")
    print(f"  Per 1K pairs: ${cost / max(calls, 1) * 1000:.2f}")
    print(f"  Per 50K pairs:${cost / max(calls, 1) * 50000:.0f}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    asyncio.run(main())