data/prepare_dataset.py

"""
Prepare datasets for Style-SFT and DPO training.

Datasets:
- HC3 (Hello-SimpleAI/HC3): ~25K human vs ChatGPT answer pairs across 6 domains
- Output: JSONL files for Style-SFT, uploaded to HF

DPO preference pairs are generated by the DPO Modal app (needs GPU for detector scoring).
"""

import argparse
import json
import os
import sys
from pathlib import Path


def download_hc3_dataset(
    output_path: str,
    max_samples: int = 25000,
    train_split: float = 0.9,
) -> dict:
    """Download HC3 from Hugging Face and convert to style transfer pairs.

    HC3 contains questions answered by BOTH humans and ChatGPT.
    Each pair is: {"ai_text": chatgpt_answer, "human_text": human_answer, "domain": source}

    Sources available: finance, medicine, law, psychology, open_qa, wiki_csai
    """
    try:
        from datasets import load_dataset
    except ImportError:
        print("[Data] ERROR: datasets library not installed. Run: pip install datasets")
        return {"status": "error", "message": "datasets not installed"}

    pairs = []
    domains_seen = set()

    print("[Data] Downloading HC3 from Hugging Face...")
    try:
        # HC3 has multiple configs (one per domain). Try loading the combined 'all' split.
        hc3 = load_dataset("Hello-SimpleAI/HC3", split="train", trust_remote_code=True)
    except Exception as e:
        print(f"[Data] Cannot load HC3 combined split ({e}), trying individual configs...")
        # Fallback: try individual configs
        configs = ["finance", "medicine", "open_qa", "reddit_eli5", "wiki_csai"]
        all_samples = []
        for cfg in configs:
            try:
                ds = load_dataset("Hello-SimpleAI/HC3", cfg, split="train", trust_remote_code=True)
                print(f"[Data]   Loaded {cfg}: {len(ds)} samples")
                # Tag each sample with its config as domain
                for s in ds:
                    s["domain"] = cfg
                    all_samples.append(s)
            except Exception as e2:
                print(f"[Data]   Skipping {cfg}: {e2}")
        if not all_samples:
            print("[Data] ERROR: Could not load any HC3 config. Check HuggingFace access.")
            return {"status": "error", "message": "No HC3 data loaded"}
        # Use list directly
        hc3 = all_samples

    print(f"[Data] HC3 loaded: {len(hc3)} raw samples")

    # Detect if hc3 has a 'config' attribute (Dataset) vs being a plain list
    current_domain = "unknown"
    if hasattr(hc3, 'config_name'):
        current_domain = hc3.config_name

    # Track which config each sample came from
    if isinstance(hc3, list):
        sample_iter = hc3
    else:
        sample_iter = hc3

    for sample in sample_iter:
        if isinstance(sample, dict):
            # HC3 format: {"id": ..., "question": ..., "human_answers": [...], "chatgpt_answers": [...]}
            human_answers = sample.get("human_answers", [])
            ai_answers = sample.get("chatgpt_answers", [])
            # Domain comes from the config name (or tagged during fallback loading)
            source = sample.get("domain", sample.get("source", current_domain))

            if human_answers and ai_answers:
                # Take the first answer from each
                human_text = human_answers[0] if isinstance(human_answers, list) else human_answers
                ai_text = ai_answers[0] if isinstance(ai_answers, list) else ai_answers

                # Filter: skip very short texts (< 50 chars) and very long (> 2000 chars)
                if len(human_text) < 50 or len(ai_text) < 50:
                    continue
                if len(human_text) > 2000 or len(ai_text) > 2000:
                    continue

                pairs.append({
                    "ai_text": ai_text.strip(),
                    "human_text": human_text.strip(),
                    "domain": source,
                })
                domains_seen.add(source)

            if len(pairs) >= max_samples:
                break

        elif isinstance(sample, str):
            # Fallback for string samples
            pairs.append({
                "ai_text": sample,
                "human_text": sample,  # placeholder
                "domain": "unknown",
            })

    print(f"[Data] Filtered to {len(pairs)} valid pairs across {len(domains_seen)} domains: {sorted(domains_seen)}")

    if len(pairs) < 100:
        print("[Data] WARNING: Very few pairs found. Check HC3 format.")
        # Fall back to synthetic pairs plus any real ones we found
        pairs.extend(_get_synthetic_pairs())
        print(f"[Data] Added synthetic pairs, total: {len(pairs)}")

    # Train/val split
    split_idx = int(len(pairs) * train_split)
    train_pairs = pairs[:split_idx]
    val_pairs = pairs[split_idx:]

    # Save
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    train_path = output_path.replace(".jsonl", "_train.jsonl")
    val_path = output_path.replace(".jsonl", "_val.jsonl")

    for path, data in [(train_path, train_pairs), (val_path, val_pairs)]:
        with open(path, "w", encoding="utf-8") as f:
            for pair in data:
                f.write(json.dumps(pair, ensure_ascii=False) + "\n")

    print(f"[Data] Saved {len(train_pairs)} train pairs to {train_path}")
    print(f"[Data] Saved {len(val_pairs)} val pairs to {val_path}")

    # Also save a text file with the Bitcoin example + other texts for CoPA testing
    test_texts_path = os.path.join(os.path.dirname(output_path), "test_texts.json")
    test_texts = {
        "ai_texts": [
            pair["ai_text"] for pair in val_pairs[:50]
        ],
        "human_texts": [
            pair["human_text"] for pair in val_pairs[:50]
        ],
    }
    with open(test_texts_path, "w", encoding="utf-8") as f:
        json.dump(test_texts, f, indent=2, ensure_ascii=False)

    return {
        "status": "ok",
        "train_path": train_path,
        "val_path": val_path,
        "num_train": len(train_pairs),
        "num_val": len(val_pairs),
        "domains": sorted(domains_seen),
    }


def _get_synthetic_pairs() -> list[dict]:
    """Minimal synthetic pairs as fallback if HC3 download fails."""
    return [
        {
            "ai_text": (
                "The implementation of machine learning algorithms has "
                "demonstrated significant improvements in various domains. "
                "These systems leverage large datasets to identify patterns "
                "and make predictions with high accuracy."
            ),
            "human_text": (
                "So I tried using ML for this project and honestly it worked "
                "way better than I expected. You feed it a bunch of data and it "
                "somehow figures out patterns you'd never spot manually."
            ),
            "domain": "tech",
        },
        {
            "ai_text": (
                "Climate change represents a critical challenge that requires "
                "coordinated international action. Rising temperatures have led "
                "to more frequent extreme weather events and ecosystem disruption."
            ),
            "human_text": (
                "Look, the climate thing is getting really scary. Every summer "
                "feels hotter than the last one, and these crazy storms keep "
                "hitting places that never used to get them."
            ),
            "domain": "science",
        },
        {
            "ai_text": (
                "The Renaissance was a period of profound cultural transformation "
                "in European history, characterized by renewed interest in "
                "classical learning and artistic innovation."
            ),
            "human_text": (
                "The Renaissance was basically Europe waking up after the Middle "
                "Ages. Artists started caring about making things look real, and "
                "people got really into old Greek and Roman stuff again."
            ),
            "domain": "history",
        },
        {
            "ai_text": (
                "Bitcoin, often called BTC, is the first and most well-known "
                "cryptocurrency. It operates on a decentralized network called "
                "blockchain and has a limited supply of 21 million coins."
            ),
            "human_text": (
                "So Bitcoin — or BTC if you're fancy — is basically the OG "
                "crypto that started it all. Nobody controls it, it runs on "
                "this thing called blockchain, and there'll only ever be "
                "21 million of them. That's it."
            ),
            "domain": "finance",
        },
        {
            "ai_text": (
                "Regular physical exercise provides numerous health benefits "
                "including improved cardiovascular function, enhanced mental "
                "wellbeing, and reduced risk of chronic diseases."
            ),
            "human_text": (
                "Look, going for a run or hitting the gym isn't just about "
                "looking good. It makes your heart stronger, clears your head, "
                "and honestly just makes you feel better overall. Plus it keeps "
                "all those scary health problems away."
            ),
            "domain": "health",
        },
    ]


def main():
    parser = argparse.ArgumentParser(description="Prepare evasion-detection datasets")
    parser.add_argument("--output-dir", default="data/processed")
    parser.add_argument("--max-samples", type=int, default=25000)
    parser.add_argument("--train-split", type=float, default=0.9)
    parser.add_argument("--upload-hf", action="store_true",
                        help="Upload to HuggingFace after preparation")
    parser.add_argument("--hf-repo", default="simonlesaumon/evasion-detection-artifacts")
    args = parser.parse_args()

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

    result = download_hc3_dataset(
        str(out / "style_transfer_pairs.jsonl"),
        max_samples=args.max_samples,
        train_split=args.train_split,
    )

    if result["status"] == "ok":
        print(f"\n[Data] Done! {result['num_train']} train + {result['num_val']} val pairs")
        print(f"[Data] Domains: {result['domains']}")
        print(f"[Data] Train: {result['train_path']}")
        print(f"[Data] Val: {result['val_path']}")

        if args.upload_hf:
            print("\n[Data] Uploading to Hugging Face...")
            try:
                import subprocess
                for local_path, hf_path in [
                    (result["train_path"], f"datasets/style_transfer_train.jsonl"),
                    (result["val_path"], f"datasets/style_transfer_val.jsonl"),
                ]:
                    subprocess.run([
                        "hf", "upload", args.hf_repo, local_path, hf_path
                    ], check=True)
                    print(f"[Data]   Uploaded {hf_path}")
                print(f"[Data] Upload complete: https://huggingface.co/{args.hf_repo}")
            except Exception as e:
                print(f"[Data] Upload skipped: {e}")
                print("[Data] Run manually: hf upload <repo> data/processed/style_transfer_pairs_train.jsonl datasets/")
    else:
        print(f"\n[Data] ERROR: {result}")
        sys.exit(1)


if __name__ == "__main__":
    main()