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

#!/usr/bin/env python3
from __future__ import annotations

"""RQ4: Within-bin text characterization for the 10K/bin working sample.

Computes per-bin lexical features (pronoun density, mental-state verb frequency,
document length stats) by streaming through the Bergson JSONL shards and joining
with the sample manifest for bin labels.

Supports three modes for distributed execution on SLURM:
  standalone  - process all shards in one process (default, for local testing)
  worker      - process a chunk of shards, write partial JSON
  merge       - combine worker partials into final CSV

Standalone usage:
    python3 rq4_bin_characterization.py \
        --manifest /path/to/working_sample_manifest.parquet \
        --shards-dir /path/to/shards/ \
        --output /path/to/rq4_bin_features.csv

Distributed usage (via launch_rq4.sh):
    # Worker (one per SLURM array task):
    python3 rq4_bin_characterization.py --mode worker \
        --manifest ... --shards-dir ... \
        --chunk-index $SLURM_ARRAY_TASK_ID --chunk-count 32 \
        --worker-output-dir /path/to/workers/$SLURM_ARRAY_TASK_ID

    # Merge (after all workers complete):
    python3 rq4_bin_characterization.py --mode merge \
        --workers-dir /path/to/workers \
        --output /path/to/rq4_bin_features.csv
"""

import argparse
import csv
import json
import re
import sys
from collections import defaultdict
from pathlib import Path

import pandas as pd

# --- Lexicon definitions ---
FIRST_PERSON = re.compile(
    r"\b(i|me|my|mine|myself|we|us|our|ours|ourselves)\b", re.IGNORECASE
)
SECOND_PERSON = re.compile(
    r"\b(you|your|yours|yourself|yourselves)\b", re.IGNORECASE
)
THIRD_PERSON = re.compile(
    r"\b(he|him|his|himself|she|her|hers|herself|they|them|their|theirs|themselves)\b",
    re.IGNORECASE,
)
MENTAL_STATE_VERBS = re.compile(
    r"\b(think|thinks|thinking|thought|believe|believes|believed|believing|"
    r"feel|feels|felt|feeling|want|wants|wanted|wanting|"
    r"know|knows|knew|knowing|understand|understands|understood|understanding|"
    r"realize|realizes|realized|realizing|expect|expects|expected|expecting|"
    r"hope|hopes|hoped|hoping|fear|fears|feared|fearing|"
    r"wish|wishes|wished|wishing|assume|assumes|assumed|assuming|"
    r"suspect|suspects|suspected|suspecting|doubt|doubts|doubted|doubting)\b",
    re.IGNORECASE,
)
WORD_RE = re.compile(r"\b\w+\b")

FIELDS = ("n_docs", "total_words", "first_person", "second_person",
          "third_person", "mental_state")


def count_features(text: str) -> dict:
    words = WORD_RE.findall(text)
    n_words = len(words)
    if n_words == 0:
        return {"n_words": 0, "first_person": 0, "second_person": 0,
                "third_person": 0, "mental_state": 0}
    return {
        "n_words": n_words,
        "first_person": len(FIRST_PERSON.findall(text)),
        "second_person": len(SECOND_PERSON.findall(text)),
        "third_person": len(THIRD_PERSON.findall(text)),
        "mental_state": len(MENTAL_STATE_VERBS.findall(text)),
    }


def load_manifest(manifest_path: str) -> dict[str, tuple[str, str]]:
    print(f"Loading manifest from {manifest_path}...", flush=True)
    df = pd.read_parquet(manifest_path, columns=["doc_id", "bin_topic", "bin_format"])
    doc_to_bin: dict[str, tuple[str, str]] = {}
    for _, row in df.iterrows():
        doc_to_bin[row["doc_id"]] = (row["bin_topic"], row["bin_format"])
    print(f"  Loaded {len(doc_to_bin):,} document-to-bin mappings", flush=True)
    return doc_to_bin


def process_shards(
    shard_files: list[Path],
    doc_to_bin: dict[str, tuple[str, str]],
) -> dict[tuple[str, str], dict[str, int]]:
    bin_stats: dict[tuple[str, str], dict[str, int]] = defaultdict(
        lambda: {f: 0 for f in FIELDS}
    )
    n_matched = 0
    n_unmatched = 0

    for i, shard_path in enumerate(shard_files):
        print(f"  Processing shard {i+1}/{len(shard_files)}: {shard_path.name}...",
              end="", flush=True)
        shard_matched = 0
        with open(shard_path, "r") as f:
            for line in f:
                doc = json.loads(line)
                doc_id = doc["id"]
                bin_key = doc_to_bin.get(doc_id)
                if bin_key is None:
                    n_unmatched += 1
                    continue
                n_matched += 1
                shard_matched += 1

                features = count_features(doc["text"])
                stats = bin_stats[bin_key]
                stats["n_docs"] += 1
                stats["total_words"] += features["n_words"]
                stats["first_person"] += features["first_person"]
                stats["second_person"] += features["second_person"]
                stats["third_person"] += features["third_person"]
                stats["mental_state"] += features["mental_state"]

        print(f" {shard_matched:,} docs matched", flush=True)

    print(f"\nTotal matched: {n_matched:,}, unmatched: {n_unmatched:,}", flush=True)
    return dict(bin_stats)


def write_partial(bin_stats: dict, output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    out_path = output_dir / "partial_bin_stats.json"
    serializable = {}
    for (topic, fmt), counts in bin_stats.items():
        key = f"{topic}|||{fmt}"
        serializable[key] = counts
    with open(out_path, "w") as f:
        json.dump(serializable, f)
    print(f"Wrote partial stats ({len(bin_stats)} bins) to {out_path}", flush=True)


def read_partial(path: Path) -> dict[tuple[str, str], dict[str, int]]:
    with open(path) as f:
        raw = json.load(f)
    result: dict[tuple[str, str], dict[str, int]] = {}
    for key, counts in raw.items():
        topic, fmt = key.split("|||", 1)
        result[(topic, fmt)] = counts
    return result


def merge_partials(workers_dir: Path) -> dict[tuple[str, str], dict[str, int]]:
    merged: dict[tuple[str, str], dict[str, int]] = defaultdict(
        lambda: {f: 0 for f in FIELDS}
    )
    partial_files = sorted(workers_dir.glob("*/partial_bin_stats.json"))
    if not partial_files:
        print(f"ERROR: No partial files found in {workers_dir}/*/", file=sys.stderr)
        sys.exit(1)

    print(f"Merging {len(partial_files)} partial files...", flush=True)
    for pf in partial_files:
        partial = read_partial(pf)
        for bin_key, counts in partial.items():
            for field in FIELDS:
                merged[bin_key][field] += counts[field]
        print(f"  Merged {pf.parent.name}: {len(partial)} bins", flush=True)

    print(f"Total bins after merge: {len(merged)}", flush=True)
    return dict(merged)


def write_csv(bin_stats: dict[tuple[str, str], dict[str, int]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "bin_topic", "bin_format", "n_docs", "total_words",
            "mean_doc_words", "first_person_per_1k", "second_person_per_1k",
            "third_person_per_1k", "mental_state_per_1k",
        ])
        for (topic, fmt), stats in sorted(bin_stats.items()):
            n = stats["n_docs"]
            w = stats["total_words"]
            if w == 0:
                continue
            writer.writerow([
                topic, fmt, n, w,
                f"{w / n:.1f}",
                f"{stats['first_person'] / w * 1000:.2f}",
                f"{stats['second_person'] / w * 1000:.2f}",
                f"{stats['third_person'] / w * 1000:.2f}",
                f"{stats['mental_state'] / w * 1000:.2f}",
            ])
    n_written = sum(1 for s in bin_stats.values() if s["total_words"] > 0)
    total_docs = sum(s["n_docs"] for s in bin_stats.values())
    print(f"\nWrote {n_written} bins ({total_docs:,} total docs) to {output_path}", flush=True)


def get_shard_files(shards_dir: str, max_shards: int | None = None,
                    chunk_index: int | None = None,
                    chunk_count: int | None = None) -> list[Path]:
    all_shards = sorted(Path(shards_dir).glob("shard_*.jsonl"))
    if chunk_index is not None and chunk_count is not None:
        all_shards = all_shards[chunk_index::chunk_count]
    if max_shards:
        all_shards = all_shards[:max_shards]
    return all_shards


def main():
    parser = argparse.ArgumentParser(description="RQ4 bin characterization")
    parser.add_argument("--mode", choices=["standalone", "worker", "merge"],
                        default="standalone")
    parser.add_argument("--manifest", help="Path to working_sample_manifest.parquet")
    parser.add_argument("--shards-dir", help="Directory with shard_NNNN.jsonl files")
    parser.add_argument("--output", help="Output CSV path")
    parser.add_argument("--max-shards", type=int, default=None,
                        help="Process only N shards (for testing)")
    parser.add_argument("--chunk-index", type=int, help="Worker chunk index")
    parser.add_argument("--chunk-count", type=int, help="Total number of chunks")
    parser.add_argument("--worker-output-dir", help="Directory for worker partial output")
    parser.add_argument("--workers-dir", help="Parent directory containing worker outputs")
    args = parser.parse_args()

    if args.mode == "worker":
        if not all([args.manifest, args.shards_dir, args.chunk_index is not None,
                    args.chunk_count, args.worker_output_dir]):
            parser.error("worker mode requires --manifest, --shards-dir, "
                         "--chunk-index, --chunk-count, --worker-output-dir")
        doc_to_bin = load_manifest(args.manifest)
        shard_files = get_shard_files(args.shards_dir, args.max_shards,
                                      args.chunk_index, args.chunk_count)
        print(f"Worker {args.chunk_index}/{args.chunk_count}: "
              f"processing {len(shard_files)} shards", flush=True)
        bin_stats = process_shards(shard_files, doc_to_bin)
        write_partial(bin_stats, Path(args.worker_output_dir))

    elif args.mode == "merge":
        if not all([args.workers_dir, args.output]):
            parser.error("merge mode requires --workers-dir and --output")
        bin_stats = merge_partials(Path(args.workers_dir))
        write_csv(bin_stats, Path(args.output))

    else:
        if not all([args.manifest, args.shards_dir, args.output]):
            parser.error("standalone mode requires --manifest, --shards-dir, --output")
        doc_to_bin = load_manifest(args.manifest)
        shard_files = get_shard_files(args.shards_dir, args.max_shards)
        bin_stats = process_shards(shard_files, doc_to_bin)
        write_csv(bin_stats, Path(args.output))


if __name__ == "__main__":
    main()