src/pseudo_docs.py

#!/usr/bin/env python3
"""
Build the corpus: canonical cross-lingual passage IDs + pseudo-documents.

EVERYTHING DOWNSTREAM HANGS OFF THIS FILE.

Two ideas, both load-bearing:

1. CANONICAL PASSAGE IDs
   MSMARCO-XI is parallel: one passage exists in English plus 14 translations.
   The join key is (query_id, passage_index) -- the same slot across every
   language view. That single key gives us:
     - one chunk_id valid in all 15 languages
     - cross-lingual score fusion (a 15-member ensemble on one query embedding)
     - References Completeness computed on ENGLISH and projected to all languages,
       sidestepping the English-only coreference limitation in arXiv:2603.25333

2. PSEUDO-DOCUMENTS
   MS MARCO passages are ~55-60 words -- roughly one paragraph. The chunking
   methods that win in the literature (Paragraph Group Chunking, arXiv:2603.06976)
   need multi-paragraph documents and degenerate to a no-op here. Chauhan & Hegde
   (DATA 2026) measured chunking at 0.9-2.0% of variance on exactly this corpus,
   plausibly because their 256/512-token chunks exceeded their ~80-token documents.

   Grouping passages back into pseudo-documents RECONSTRUCTS the conditions under
   which chunking matters. It also revives Block Integrity as a metric: original
   passage boundaries become the gold blocks, so BI = "did we split a passage?"

   This is a TESTED HYPOTHESIS, not a preference. Run the ANOVA on raw passages
   AND on pseudo-documents; if chunking's eta-squared rises, you have empirically
   identified when chunking starts to matter -- the open question across all six
   papers in the review.

GROUPING STRATEGIES
  url       group by source URL          (truest to original web documents)
  query     group by query_id            (always available; topically coherent)
  cluster   embedding k-means            (best coherence, needs GPU embeddings)

Outputs (parquet, under $VOICERAG_ROOT/data):
  passages.parquet     one row per (canonical_id, lang)
  pseudo_docs.parquet  one row per (pseudo_doc_id, lang), with block offsets
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.schema_utils import iter_passages, norm_lang, default_root, load_report  # noqa: E402

TARGET_WORDS = 3000   # ~4-8k tokens: a realistic web-document size
MIN_WORDS = 800
MAX_PASSAGES_PER_DOC = 60


def canonical_id(query_id: str | int, passage_index: int) -> str:
    """Language-independent passage identity. Stable across runs and machines."""
    raw = f"{query_id}::{passage_index}"
    return hashlib.blake2b(raw.encode(), digest_size=8).hexdigest()


def load_schema(root: Path) -> tuple[dict, dict, list[str]]:
    report = load_report(root)
    return report["field_mapping"], report["passage_mapping"], report["files"]


def explode_passages(root: Path, limit_files: int | None) -> "object":
    """Flatten nested passages into one row per (canonical_id, lang)."""
    import polars as pl

    fmap, pmap, files = load_schema(root)
    pcol = fmap["passages"]
    qid_col, lang_col = fmap.get("query_id"), fmap.get("lang")
    t_key, en_key = pmap.get("text"), pmap.get("text_en")
    sel_key, url_key = pmap.get("is_selected"), pmap.get("url")

    if not (pcol and t_key and qid_col):
        raise SystemExit("schema_report.json is missing passages/text/query_id — rerun 02_inspect_schema.py")

    cols = [c for c in (qid_col, lang_col, pcol) if c]
    frames = []

    for fp in (files[:limit_files] if limit_files else files):
        try:
            df = pl.read_parquet(fp, columns=cols)
        except Exception as exc:
            print(f"    skip {Path(fp).name}: {exc}")
            continue

        rows = []
        langs = df[lang_col].to_list() if lang_col else [None] * len(df)
        for qid, lang, plist in zip(df[qid_col].to_list(), langs, df[pcol].to_list()):
            lang = norm_lang(lang)
            for idx, text, text_en, sel, url in iter_passages(plist, t_key, en_key, sel_key, url_key):
                if not isinstance(text, str) or not text.strip():
                    continue
                rows.append({
                    "canonical_id": canonical_id(qid, idx),
                    "query_id": str(qid),
                    "passage_index": idx,
                    "lang": lang,
                    "text": text,
                    "text_en": text_en,
                    "is_selected": sel,
                    "url": url,
                    "n_words": len(text.split()),
                })
        if rows:
            frames.append(pl.DataFrame(rows))
            print(f"    {Path(fp).name[:48]:48s} +{len(rows):>8,} passages")

    if not frames:
        raise SystemExit("no passages extracted — check the passage mapping in schema_report.json")

    out = pl.concat(frames, how="vertical_relaxed")
    before = len(out)
    out = out.unique(subset=["canonical_id", "lang"], keep="first")
    print(f"    deduped {before:,} -> {len(out):,} unique (canonical_id, lang)")
    return out


def build_pseudo_docs(passages, strategy: str, target_words: int):
    """Group passages into documents. Grouping is decided on the ENGLISH view and
    applied identically to every language, so pseudo_doc_id is also canonical."""
    import polars as pl

    if strategy == "url" and passages["url"].null_count() < len(passages):
        key = "url"
    else:
        if strategy == "url":
            print("    !! no usable url column, falling back to query grouping")
        key = "query_id"

    # Decide grouping once, on English (or whatever the largest view is).
    langs = passages["lang"].value_counts().sort("count", descending=True)
    pivot_lang = "en" if "en" in passages["lang"].unique().to_list() else langs["lang"][0]
    pivot = passages.filter(pl.col("lang") == pivot_lang).sort([key, "query_id", "passage_index"])
    print(f"    grouping on '{key}' using pivot language '{pivot_lang}' ({len(pivot):,} passages)")

    assignment: dict[str, tuple[str, int]] = {}   # canonical_id -> (doc_id, block_index)
    doc_n = 0
    cur_key, cur_words, cur_block = None, 0, 0
    doc_id = None

    for cid, gkey, nw in zip(pivot["canonical_id"], pivot[key], pivot["n_words"]):
        start_new = (
            doc_id is None
            or gkey != cur_key
            or cur_words >= target_words
            or cur_block >= MAX_PASSAGES_PER_DOC
        )
        if start_new:
            doc_n += 1
            doc_id = f"pd{doc_n:08d}"
            cur_key, cur_words, cur_block = gkey, 0, 0
        assignment[cid] = (doc_id, cur_block)
        cur_words += nw
        cur_block += 1

    print(f"    formed {doc_n:,} pseudo-documents")

    amap = pl.DataFrame({
        "canonical_id": list(assignment.keys()),
        "pseudo_doc_id": [v[0] for v in assignment.values()],
        "block_index": [v[1] for v in assignment.values()],
    })
    # Left join: passages present in a language but absent from the pivot get null doc ids.
    tagged = passages.join(amap, on="canonical_id", how="left")

    orphans = tagged["pseudo_doc_id"].null_count()
    if orphans:
        print(f"    !! {orphans:,} passages had no pivot-language counterpart (dropped from pseudo-docs)")

    docs = (
        tagged.filter(pl.col("pseudo_doc_id").is_not_null())
        .sort(["pseudo_doc_id", "lang", "block_index"])
        .group_by(["pseudo_doc_id", "lang"])
        .agg([
            pl.col("text").alias("blocks"),                 # gold blocks -> Block Integrity metric
            pl.col("canonical_id").alias("block_ids"),
            pl.col("is_selected").alias("block_labels"),
            pl.col("n_words").sum().alias("n_words"),
            pl.len().alias("n_blocks"),
        ])
        .with_columns(pl.col("blocks").list.join("\n\n").alias("text"))
        .filter(pl.col("n_words") >= MIN_WORDS)
    )
    return tagged, docs


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None,
                    help="data root; defaults to $VOICERAG_ROOT or <repo>/../voicerag_data")
    ap.add_argument("--strategy", choices=["url", "query", "cluster"], default="url")
    ap.add_argument("--target-words", type=int, default=TARGET_WORDS)
    ap.add_argument("--limit-files", type=int, default=None, help="for a fast smoke test")
    args = ap.parse_args()

    root = (args.root.expanduser().resolve() if args.root else default_root())
    print(f"==> data root: {root}")
    data = root / "data"
    data.mkdir(parents=True, exist_ok=True)

    if args.strategy == "cluster":
        raise SystemExit("cluster grouping needs the embedding index — use 'url' or 'query' first")

    print("==> exploding passages")
    passages = explode_passages(root, args.limit_files)

    print("\n==> per-language passage counts")
    for lang, n in passages["lang"].value_counts().sort("count", descending=True).iter_rows():
        print(f"    {lang:>4s}  {n:>12,}")

    print(f"\n==> building pseudo-documents (strategy={args.strategy}, target={args.target_words} words)")
    passages, docs = build_pseudo_docs(passages, args.strategy, args.target_words)

    p_out, d_out = data / "passages.parquet", data / "pseudo_docs.parquet"
    passages.write_parquet(p_out, compression="zstd")
    docs.write_parquet(d_out, compression="zstd")

    import polars as pl
    print(f"\n==> wrote {p_out}  ({len(passages):,} rows, {p_out.stat().st_size/2**20:.0f} MiB)")
    print(f"==> wrote {d_out}  ({len(docs):,} rows, {d_out.stat().st_size/2**20:.0f} MiB)")
    print(f"\n    unique canonical passages : {passages['canonical_id'].n_unique():,}")
    print(f"    language views            : {passages['lang'].n_unique()}")
    print(f"    labelled positives        : {passages.filter(pl.col('is_selected') == 1).height:,}")
    print(f"    pseudo-doc mean words     : {docs['n_words'].mean():.0f}")
    print(f"    pseudo-doc mean blocks    : {docs['n_blocks'].mean():.1f}")
    print("\n    Next: chunking portfolio runs against BOTH passages.parquet and")
    print("          pseudo_docs.parquet — that pair is the ANOVA 'corpus form' factor.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())