scripts/split_dataset.py

#!/usr/bin/env python3
"""
split_dataset.py

Splits the training file into train/validation sets.

Usage:
    python split_dataset.py --ratio 0.8  # 80% train, 20% valid
    python split_dataset.py --count 100  # Keep 100 examples for valid
    python split_dataset.py --stratify  # Preserve distribution across task_type

The validation set is written to datasets/mythos_coder_valid.jsonl.
"""

import argparse
import json
import random
from collections import defaultdict
from pathlib import Path


def load_examples(file_path):
    """Load all examples from a JSONL file."""
    examples = []
    if not file_path.exists():
        return examples

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    examples.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return examples


def save_examples(examples, file_path):
    """Save examples to a JSONL file."""
    file_path.parent.mkdir(parents=True, exist_ok=True)
    with open(file_path, "w", encoding="utf-8") as f:
        for ex in examples:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")


def simple_split(examples, ratio, seed=None):
    """Split examples randomly by ratio."""
    if seed is not None:
        random.seed(seed)

    shuffled = examples.copy()
    random.shuffle(shuffled)

    split_idx = int(len(shuffled) * ratio)
    train = shuffled[:split_idx]
    valid = shuffled[split_idx:]

    return train, valid


def stratified_split(examples, ratio, seed=None):
    """Split examples preserving distribution across task_type."""
    if seed is not None:
        random.seed(seed)

    # Group by task_type
    by_type = defaultdict(list)
    for ex in examples:
        task_type = ex.get("task_type", "unknown")
        by_type[task_type].append(ex)

    train = []
    valid = []

    for task_type, type_examples in by_type.items():
        shuffled = type_examples.copy()
        random.shuffle(shuffled)

        split_idx = int(len(shuffled) * ratio)
        train.extend(shuffled[:split_idx])
        valid.extend(shuffled[split_idx:])

    # Shuffle again to mix types
    random.shuffle(train)
    random.shuffle(valid)

    return train, valid


def count_split(examples, valid_count, seed=None):
    """Split keeping specific count for validation."""
    if seed is not None:
        random.seed(seed)

    shuffled = examples.copy()
    random.shuffle(shuffled)

    valid_count = min(valid_count, len(shuffled))
    valid = shuffled[:valid_count]
    train = shuffled[valid_count:]

    return train, valid


def print_distribution(examples, label):
    """Print distribution statistics."""
    by_type = defaultdict(int)
    by_difficulty = defaultdict(int)

    for ex in examples:
        by_type[ex.get("task_type", "unknown")] += 1
        by_difficulty[ex.get("difficulty", "unknown")] += 1

    print(f"\n{label} ({len(examples)} examples):")
    print(f"  By task_type: {dict(by_type)}")
    print(f"  By difficulty: {dict(by_difficulty)}")


def main():
    parser = argparse.ArgumentParser(description="Split dataset into train/validation")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--ratio", type=float, help="Train ratio (e.g., 0.8 for 80%% train)")
    group.add_argument("--count", type=int, help="Number of examples for validation set")
    parser.add_argument("--stratify", action="store_true", help="Stratify by task_type")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--in-place", action="store_true", help="Overwrite training file with reduced set")
    args = parser.parse_args()

    project_root = Path(__file__).parent.parent
    train_path = project_root / "datasets" / "mythos_coder_train.jsonl"
    valid_path = project_root / "datasets" / "mythos_coder_valid.jsonl"

    # Load all examples
    examples = load_examples(train_path)
    if not examples:
        print(f"No examples found in {train_path}")
        return

    print(f"Loaded {len(examples)} examples from {train_path}")
    print_distribution(examples, "Original distribution")

    # Perform split
    if args.stratify:
        if args.ratio is None:
            print("Error: --stratify requires --ratio")
            return
        train, valid = stratified_split(examples, args.ratio, args.seed)
    elif args.count is not None:
        train, valid = count_split(examples, args.count, args.seed)
    else:
        train, valid = simple_split(examples, args.ratio, args.seed)

    # Save files
    if args.in_place:
        save_examples(train, train_path)
        print(f"\nOverwrote {train_path} with {len(train)} examples")
    else:
        print(f"\nKeeping all {len(examples)} examples in training file")

    save_examples(valid, valid_path)
    print(f"Saved {len(valid)} examples to {valid_path}")

    # Print distributions
    if args.in_place:
        print_distribution(train, "Training set")
    print_distribution(valid, "Validation set")


if __name__ == "__main__":
    main()