social-data-attribution-w2/scripts/analysis/rq4_lexical_toolkit_profiles.py

#!/usr/bin/env python3
# pyright: reportArgumentType=false
"""RQ4 lexical enrichment CLI.

Modes:
  standalone : process all matching shards, write per-bin Parquet + CSV.
  worker     : process a sharded slice, write partial_bin_stats.json
               (+ optional per_doc Parquet partial).
  merge      : merge worker partials into per_bin_profiles.parquet/.csv,
               compute densities, z-scores, composites, baselines.

The pipeline reuses the legacy ``\\b\\w+\\b`` tokenizer and pronoun /
mental-state lexicons (mirrored as ``lexicons/rq4_pronouns.json`` and
``lexicons/rq4_mental_state_verbs.json``); see
``artifacts/rq4_lexical_profiles/README.md`` for the parity contract."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))

from data_attribution.rq4_lexical.aggregate import (  # noqa: E402
    add_composites,
    add_densities,
    add_summary_variables,
    add_text_ratios,
    add_zscores,
    stats_to_dataframe,
)
from data_attribution.rq4_lexical.config import load_lexical_config  # noqa: E402
from data_attribution.rq4_lexical.dimensions import build_dimension_counter  # noqa: E402
from data_attribution.rq4_lexical.pipeline import (  # noqa: E402
    list_shard_files,
    load_doc_to_bin,
    merge_partials_dir,
    process_shard,
    write_partial,
)

REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_PRONOUNS = REPO_ROOT / "lexicons" / "rq4_pronouns.json"
DEFAULT_MENTAL = REPO_ROOT / "lexicons" / "rq4_mental_state_verbs.json"
DEFAULT_DIALOGUE = REPO_ROOT / "lexicons" / "rq4_dialogue_patterns.yaml"
DEFAULT_EMPATH = REPO_ROOT / "configs" / "rq4_empath_categories.yaml"
DEFAULT_CROSSWALK = REPO_ROOT / "configs" / "rq4_liwc_empath_crosswalk.yaml"
DEFAULT_DIMENSIONS = REPO_ROOT / "configs" / "rq4_liwc_open_dimensions.yaml"
DEFAULT_LEXICONS_DIR = REPO_ROOT / "lexicons"


def _load_cfg(args: argparse.Namespace):
    return load_lexical_config(
        args.pronouns,
        args.mental_state,
        args.dialogue,
        args.empath_config,
        crosswalk_path=args.crosswalk_config,
        lexicons_dir=args.lexicons_dir,
        dimensions_path=args.dimensions_config,
    )


def _build_counter(cfg):
    return build_dimension_counter(
        list(cfg.empath_categories),
        {dim: list(cats) for dim, cats in cfg.crosswalk.items()},
        cfg.open_lexicons,
    )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="RQ4 lexical enrichment toolkit")
    p.add_argument(
        "--mode", choices=["standalone", "worker", "merge"], default="standalone"
    )
    p.add_argument("--manifest", help="Path to working_sample_manifest.parquet")
    p.add_argument("--shards-dir", help="Directory with shard_NNNN.jsonl[.zst] files")
    p.add_argument(
        "--shard-glob",
        default="shard_*.jsonl",
        help="Glob for shard files inside --shards-dir.",
    )
    p.add_argument("--max-shards", type=int, default=None)
    p.add_argument("--max-docs", type=int, default=None)
    p.add_argument("--chunk-index", type=int, default=None)
    p.add_argument("--chunk-count", type=int, default=None)
    p.add_argument("--worker-output-dir", help="Output dir for worker partial JSON")
    p.add_argument("--workers-dir", help="Parent dir of worker outputs (merge mode)")
    p.add_argument("--out", help="Output dir for merged Parquet/CSV (merge/standalone)")
    p.add_argument(
        "--write-per-doc",
        action="store_true",
        help="Worker also writes per-doc Parquet partial (for bootstrap).",
    )
    p.add_argument("--pronouns", type=Path, default=DEFAULT_PRONOUNS)
    p.add_argument("--mental-state", type=Path, default=DEFAULT_MENTAL)
    p.add_argument("--dialogue", type=Path, default=DEFAULT_DIALOGUE)
    p.add_argument("--empath-config", type=Path, default=DEFAULT_EMPATH)
    p.add_argument(
        "--crosswalk-config",
        type=Path,
        default=DEFAULT_CROSSWALK,
        help="LIWC->Empath crosswalk YAML (open parity layer).",
    )
    p.add_argument(
        "--dimensions-config",
        type=Path,
        default=DEFAULT_DIMENSIONS,
        help="Master LIWC-22 open dimension registry YAML.",
    )
    p.add_argument(
        "--lexicons-dir",
        type=Path,
        default=DEFAULT_LEXICONS_DIR,
        help="Directory holding the closed-class open lexicon JSONs.",
    )
    return p.parse_args()


def run_worker(args: argparse.Namespace) -> None:
    required = ["manifest", "shards_dir", "worker_output_dir", "chunk_count"]
    if (
        any(getattr(args, k) in (None, "") for k in required)
        or args.chunk_index is None
    ):
        sys.exit(
            "worker mode requires --manifest --shards-dir --chunk-index --chunk-count --worker-output-dir"
        )
    cfg = _load_cfg(args)
    counter = _build_counter(cfg)
    print(
        f"[worker {args.chunk_index}/{args.chunk_count}] empath cats={len(cfg.empath_categories)} labels={len(counter.labels)} dim_fields={len(cfg.dimension_count_fields)}",
        flush=True,
    )
    doc_to_bin = load_doc_to_bin(Path(args.manifest))
    print(f"[worker {args.chunk_index}] manifest docs: {len(doc_to_bin):,}", flush=True)
    shards = list_shard_files(
        Path(args.shards_dir),
        args.shard_glob,
        args.chunk_index,
        args.chunk_count,
        args.max_shards,
    )
    print(f"[worker {args.chunk_index}] processing {len(shards)} shards", flush=True)
    bin_stats: dict[tuple[str, str], dict[str, int]] = {}
    per_doc_rows: list[dict] | None = [] if args.write_per_doc else None
    docs_seen = 0
    for shard in shards:
        docs_seen = process_shard(
            shard,
            doc_to_bin,
            cfg,
            counter,
            bin_stats,
            per_doc_sink=per_doc_rows,
            max_docs=args.max_docs,
            docs_so_far=docs_seen,
        )
        if args.max_docs is not None and docs_seen >= args.max_docs:
            break
    out_dir = Path(args.worker_output_dir)
    write_partial(bin_stats, out_dir)
    if per_doc_rows is not None:
        per_doc_path = out_dir / "per_doc_features.parquet"
        pd.DataFrame(per_doc_rows).to_parquet(per_doc_path, index=False)
        print(
            f"[worker {args.chunk_index}] per_doc rows: {len(per_doc_rows):,} -> {per_doc_path}",
            flush=True,
        )


def _write_merged(
    bin_stats: dict[tuple[str, str], dict[str, int]],
    cfg,
    out_dir: Path,
) -> None:
    df = stats_to_dataframe(bin_stats, cfg)
    df = add_densities(df, cfg)
    df, baselines = add_zscores(df, cfg)
    df = add_composites(df, cfg)
    df = add_text_ratios(df, cfg)
    df = add_summary_variables(df, cfg)
    out_dir.mkdir(parents=True, exist_ok=True)
    df.to_parquet(out_dir / "per_bin_profiles.parquet", index=False)
    df.to_csv(out_dir / "per_bin_profiles.csv", index=False)
    with (out_dir / "corpus_baselines.json").open("w") as fh:
        json.dump({"baselines_per_1k_density": baselines}, fh, indent=2)
    summary = {
        "n_bins": int(df.shape[0]),
        "n_docs": int(df["n_docs"].sum()) if "n_docs" in df else None,
        "total_words": int(df["total_words"].sum()) if "total_words" in df else None,
    }
    with (out_dir / "run_summary.json").open("w") as fh:
        json.dump(summary, fh, indent=2)
    print(
        f"merge: wrote {df.shape[0]} bins to {out_dir}/per_bin_profiles.*", flush=True
    )


def run_merge(args: argparse.Namespace) -> None:
    if not args.workers_dir or not args.out:
        sys.exit("merge mode requires --workers-dir and --out")
    cfg = _load_cfg(args)
    bin_stats = merge_partials_dir(Path(args.workers_dir), cfg)
    _write_merged(bin_stats, cfg, Path(args.out))


def run_standalone(args: argparse.Namespace) -> None:
    if not (args.manifest and args.shards_dir and args.out):
        sys.exit("standalone mode requires --manifest --shards-dir --out")
    cfg = _load_cfg(args)
    counter = _build_counter(cfg)
    doc_to_bin = load_doc_to_bin(Path(args.manifest))
    shards = list_shard_files(
        Path(args.shards_dir), args.shard_glob, None, None, args.max_shards
    )
    print(
        f"standalone: {len(shards)} shards, {len(doc_to_bin):,} doc->bin entries",
        flush=True,
    )
    bin_stats: dict[tuple[str, str], dict[str, int]] = {}
    per_doc_rows: list[dict] | None = [] if args.write_per_doc else None
    docs_seen = 0
    for shard in shards:
        docs_seen = process_shard(
            shard,
            doc_to_bin,
            cfg,
            counter,
            bin_stats,
            per_doc_sink=per_doc_rows,
            max_docs=args.max_docs,
            docs_so_far=docs_seen,
        )
        if args.max_docs is not None and docs_seen >= args.max_docs:
            break
    _write_merged(bin_stats, cfg, Path(args.out))
    if per_doc_rows is not None:
        out_dir = Path(args.out)
        per_doc_path = out_dir / "per_doc_features.parquet"
        pd.DataFrame(per_doc_rows).to_parquet(per_doc_path, index=False)
        print(
            f"standalone: per_doc rows {len(per_doc_rows):,} -> {per_doc_path}",
            flush=True,
        )


def main() -> None:
    args = parse_args()
    if args.mode == "worker":
        run_worker(args)
    elif args.mode == "merge":
        run_merge(args)
    else:
        run_standalone(args)


if __name__ == "__main__":
    main()