code/create_splits.py

"""
Create train/val/test splits from LLM annotations.
====================================================
Run this once before training. Outputs three CSV files:
    splits/train.csv
    splits/val.csv
    splits/test.csv

Each CSV has columns: text, label

Usage:
    python create_splits.py \
        --annotations ../audit_output/annotations_flagged_llama_100k.csv \
        --documents ../audit_output/flagged_documents.csv.gz \
        --random-annotations ../audit_output/annotations_random_100k_llama.csv \
        --random-documents ../audit_output/random_100k.csv \
        --output-dir splits
"""

import csv
import gzip
import argparse
from collections import Counter
from pathlib import Path

from sklearn.model_selection import train_test_split

from config import CLASSES, BENIGN_CLASS, BINARY_CLASSES, MISINFO_CLASS, TRAIN_RATIO, VAL_RATIO, SEED

csv.field_size_limit(10_000_000)


def load_annotations(annotations_path, documents_path):
    """Load annotation CSV and join with full texts from documents CSV.

    The annotation CSV has: url, domain, category, score, <prefix>_label, <prefix>_confidence, <prefix>_reason
    The documents CSV has: url, ..., full_text
    """
    # Step 1: read annotations to get url -> label mapping
    url_to_label = {}
    skipped = Counter()
    with open(annotations_path) as f:
        reader = csv.DictReader(f)
        label_col = None
        for col in reader.fieldnames:
            if col.endswith("_label") or col == "label":
                label_col = col
                break
        if not label_col:
            raise ValueError(f"No label column found. Fields: {reader.fieldnames}")

        for row in reader:
            label = row[label_col].strip().lower()
            if label not in CLASSES:
                skipped[label] += 1
                continue
            url_to_label[row["url"]] = label

    print(f"  {len(url_to_label):,} annotated URLs (skipped: {dict(skipped) if skipped else 'none'})")

    # Step 2: join with full texts from documents file
    texts, labels = [], []
    is_gz = str(documents_path).endswith(".gz")
    opener = gzip.open if is_gz else open
    mode = "rt" if is_gz else "r"

    with opener(documents_path, mode) as f:
        reader = csv.DictReader(f)
        for row in reader:
            url = row.get("url", "")
            if url in url_to_label:
                text = row.get("full_text", "")
                if text.strip():
                    texts.append(text)
                    labels.append(url_to_label[url])
                    del url_to_label[url]
            if not url_to_label:
                break

    print(f"  Loaded {len(texts):,} annotated examples with text")
    return texts, labels


def save_split(texts, labels, path):
    """Save a split as CSV with text,label columns."""
    with open(path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["text", "label"])
        for text, label in zip(texts, labels):
            writer.writerow([text, label])
    print(f"  Saved {len(texts):,} examples to {path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--annotations", required=True,
                        help="Llama annotation CSV for flagged docs (url, label, ...)")
    parser.add_argument("--documents", required=True,
                        help="Flagged documents CSV with full text (gzipped or plain)")
    parser.add_argument("--random-annotations", required=True,
                        help="LLM annotation CSV for random docs (url, label, ...)")
    parser.add_argument("--random-documents", required=True,
                        help="Random documents CSV with full text")
    parser.add_argument("--output-dir", default="splits",
                        help="Output directory for splits")
    parser.add_argument("--binary", action="store_true",
                        help="Collapse all misinfo classes into a single 'misinfo' label")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load flagged annotations joined with full texts (positives + benign-on-flagged-domains)
    print("Loading flagged annotations...")
    ann_texts, ann_labels = load_annotations(args.annotations, args.documents)

    # Load random annotations (use only benign as verified negatives)
    print("\nLoading random annotations (benign only)...")
    rand_texts, rand_labels = load_annotations(args.random_annotations, args.random_documents)
    benign_texts = [t for t, l in zip(rand_texts, rand_labels) if l == BENIGN_CLASS]
    benign_labels = [BENIGN_CLASS] * len(benign_texts)
    print(f"  Kept {len(benign_texts):,} benign from {len(rand_texts):,} random annotations")

    all_texts = ann_texts + benign_texts
    all_labels = ann_labels + benign_labels

    # Binary mode: collapse all misinfo classes into 'misinfo'
    if args.binary:
        all_labels = [MISINFO_CLASS if l != BENIGN_CLASS else l for l in all_labels]
        classes = BINARY_CLASSES
        print("\n[Binary mode] Collapsed all misinfo classes into 'misinfo'")
    else:
        classes = CLASSES

    print(f"\nTotal: {len(all_texts):,} examples")
    counts = Counter(all_labels)
    for cls in classes:
        print(f"  {cls:<25}: {counts.get(cls, 0):>6}")

    # Stratified split
    test_ratio = 1 - TRAIN_RATIO - VAL_RATIO
    print(f"\nSplitting {TRAIN_RATIO:.0%}/{VAL_RATIO:.0%}/{test_ratio:.0%}...")

    X_train, X_tmp, y_train, y_tmp = train_test_split(
        all_texts, all_labels,
        test_size=(1 - TRAIN_RATIO),
        stratify=all_labels,
        random_state=SEED,
    )
    val_frac = VAL_RATIO / (1 - TRAIN_RATIO)
    X_val, X_test, y_val, y_test = train_test_split(
        X_tmp, y_tmp,
        test_size=(1 - val_frac),
        stratify=y_tmp,
        random_state=SEED,
    )

    # Save
    print("\nSaving splits...")
    save_split(X_train, y_train, output_dir / "train.csv")
    save_split(X_val, y_val, output_dir / "val.csv")
    save_split(X_test, y_test, output_dir / "test.csv")

    # Summary
    print("\nSplit summary:")
    for name, labels in [("train", y_train), ("val", y_val), ("test", y_test)]:
        c = Counter(labels)
        print(f"  {name}: {len(labels):,} total")
        for cls in classes:
            print(f"    {cls:<25}: {c.get(cls, 0):>6}")


if __name__ == "__main__":
    main()