""" 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, _label, _confidence, _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()