feat: add dataset loader for complexity classification

- Support Easy2Hard-Bench dataset with difficulty scores
- Support ARC dataset with Easy/Challenge splits
- Convert continuous difficulty to binary labels
- Create train/validation/test splits (70/15/15)

ml/data/load_dataset.py

@@ -0,0 +1,236 @@
+"""Load and preprocess Easy2Hard-Bench dataset for complexity classification."""
+
+import json
+from pathlib import Path
+from typing import Literal
+
+from datasets import Dataset, DatasetDict, concatenate_datasets, load_dataset
+
+
+def load_easy2hard_bench(
+    subset: Literal["all", "gsm8k", "arc", "winogrande"] = "all",
+    difficulty_threshold: float = 0.5,
+    max_samples: int | None = None,
+    seed: int = 42,
+) -> DatasetDict:
+    """
+    Load Easy2Hard-Bench dataset and convert to binary classification.
+
+    Args:
+        subset: Which subset to load ("all" for combined dataset)
+        difficulty_threshold: Score above this is "complex" (1), below is "simple" (0)
+        max_samples: Maximum samples to use (None for all)
+        seed: Random seed for shuffling
+
+    Returns:
+        DatasetDict with train/validation/test splits
+    """
+    print(f"Loading Easy2Hard-Bench dataset (subset={subset})...")
+
+    # Load the dataset from HuggingFace
+    dataset = load_dataset("furonghuang-lab/Easy2Hard-Bench")
+
+    # Get all available splits
+    all_data = []
+
+    for split_name in dataset.keys():
+        split_data = dataset[split_name]
+        all_data.append(split_data)
+
+    # Combine all splits
+    combined = concatenate_datasets(all_data)
+
+    print(f"Total examples loaded: {len(combined)}")
+
+    # Process the dataset
+    def process_example(example: dict) -> dict:
+        """Extract text and create binary label from difficulty score."""
+        # Get the question/prompt text
+        text = example.get("question", "") or example.get("prompt", "") or example.get("input", "")
+
+        # Get difficulty score (normalize to 0-1 if needed)
+        difficulty = example.get("difficulty", 0.5)
+
+        # Convert to binary label
+        label = 1 if difficulty >= difficulty_threshold else 0
+
+        return {
+            "text": str(text).strip(),
+            "label": label,
+            "difficulty_score": float(difficulty),
+        }
+
+    # Apply processing
+    processed = combined.map(
+        process_example,
+        remove_columns=combined.column_names,
+        desc="Processing examples",
+    )
+
+    # Filter out empty texts
+    processed = processed.filter(lambda x: len(x["text"]) > 0)
+
+    print(f"After filtering empty texts: {len(processed)}")
+
+    # Shuffle the dataset
+    processed = processed.shuffle(seed=seed)
+
+    # Limit samples if specified
+    if max_samples and len(processed) > max_samples:
+        processed = processed.select(range(max_samples))
+        print(f"Limited to {max_samples} samples")
+
+    # Create train/val/test splits (70/15/15)
+    train_test = processed.train_test_split(test_size=0.3, seed=seed)
+    val_test = train_test["test"].train_test_split(test_size=0.5, seed=seed)
+
+    dataset_dict = DatasetDict(
+        {
+            "train": train_test["train"],
+            "validation": val_test["train"],
+            "test": val_test["test"],
+        }
+    )
+
+    # Print statistics
+    print("\nDataset splits:")
+    for split_name, split_data in dataset_dict.items():
+        n_simple = sum(1 for x in split_data if x["label"] == 0)
+        n_complex = sum(1 for x in split_data if x["label"] == 1)
+        print(f"  {split_name}: {len(split_data)} total ({n_simple} simple, {n_complex} complex)")
+
+    return dataset_dict
+
+
+def load_arc_dataset(max_samples: int | None = None, seed: int = 42) -> DatasetDict:
+    """
+    Load ARC dataset with pre-defined Easy/Challenge splits.
+
+    This is an alternative to Easy2Hard-Bench that has explicit easy/hard labels.
+
+    Args:
+        max_samples: Maximum samples per split (None for all)
+        seed: Random seed for shuffling
+
+    Returns:
+        DatasetDict with train/validation/test splits
+    """
+    print("Loading ARC dataset (Easy + Challenge)...")
+
+    # Load both splits
+    arc_easy = load_dataset("allenai/ai2_arc", "ARC-Easy")
+    arc_challenge = load_dataset("allenai/ai2_arc", "ARC-Challenge")
+
+    def process_arc(example: dict, label: int) -> dict:
+        """Process ARC example."""
+        return {
+            "text": example["question"].strip(),
+            "label": label,
+            "difficulty_score": 0.25 if label == 0 else 0.75,
+        }
+
+    # Process and label
+    easy_data = arc_easy["train"].map(
+        lambda x: process_arc(x, 0),
+        remove_columns=arc_easy["train"].column_names,
+    )
+    challenge_data = arc_challenge["train"].map(
+        lambda x: process_arc(x, 1),
+        remove_columns=arc_challenge["train"].column_names,
+    )
+
+    # Combine
+    combined = concatenate_datasets([easy_data, challenge_data])
+    combined = combined.shuffle(seed=seed)
+
+    if max_samples and len(combined) > max_samples:
+        combined = combined.select(range(max_samples))
+
+    # Split
+    train_test = combined.train_test_split(test_size=0.3, seed=seed)
+    val_test = train_test["test"].train_test_split(test_size=0.5, seed=seed)
+
+    dataset_dict = DatasetDict(
+        {
+            "train": train_test["train"],
+            "validation": val_test["train"],
+            "test": val_test["test"],
+        }
+    )
+
+    print("\nDataset splits:")
+    for split_name, split_data in dataset_dict.items():
+        n_simple = sum(1 for x in split_data if x["label"] == 0)
+        n_complex = sum(1 for x in split_data if x["label"] == 1)
+        print(f"  {split_name}: {len(split_data)} total ({n_simple} simple, {n_complex} complex)")
+
+    return dataset_dict
+
+
+def save_dataset(dataset: DatasetDict, output_dir: str | Path) -> None:
+    """Save processed dataset to disk."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for split_name, split_data in dataset.items():
+        output_path = output_dir / f"{split_name}.jsonl"
+        with open(output_path, "w") as f:
+            for example in split_data:
+                f.write(json.dumps(example) + "\n")
+        print(f"Saved {split_name} to {output_path}")
+
+
+if __name__ == "__main__":
+    # Example usage
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Load complexity classification dataset")
+    parser.add_argument(
+        "--dataset",
+        choices=["easy2hard", "arc"],
+        default="easy2hard",
+        help="Dataset to load",
+    )
+    parser.add_argument(
+        "--max-samples",
+        type=int,
+        default=None,
+        help="Maximum samples to use",
+    )
+    parser.add_argument(
+        "--threshold",
+        type=float,
+        default=0.5,
+        help="Difficulty threshold for binary classification",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="ml/data/processed",
+        help="Output directory for processed data",
+    )
+    parser.add_argument(
+        "--save",
+        action="store_true",
+        help="Save processed dataset to disk",
+    )
+
+    args = parser.parse_args()
+
+    if args.dataset == "easy2hard":
+        dataset = load_easy2hard_bench(
+            difficulty_threshold=args.threshold,
+            max_samples=args.max_samples,
+        )
+    else:
+        dataset = load_arc_dataset(max_samples=args.max_samples)
+
+    if args.save:
+        save_dataset(dataset, args.output_dir)
+
+    # Show some examples
+    print("\nSample examples:")
+    for i, example in enumerate(dataset["train"].select(range(3))):
+        label_str = "complex" if example["label"] == 1 else "simple"
+        print(f"\n[{i+1}] ({label_str}, score={example['difficulty_score']:.2f})")
+        print(f"    {example['text'][:100]}...")