src/cpptai/datasets.py

"""Dataset loaders for broader benchmarking (GSM8K, MATH, HumanEval, SciBench).
Currently provides stubs and synthetic generators, as actual datasets require external files.
"""

from pathlib import Path
from typing import List, Dict, Optional
import json
import logging
import random

try:
    from datasets import load_dataset
    HAS_HF_DATASETS = True
except ImportError:
    HAS_HF_DATASETS = False

logger = logging.getLogger(__name__)


def _categorize_gsm8k(question: str) -> str:
    """Categorize a GSM8K problem by type.

    Returns one of:
    - 'one_step': single arithmetic operation
    - 'multi_step': multiple operations
    - 'word_problem': story-based with multiple entities
    - 'comparison': comparing quantities
    - 'fraction': involving fractions or percentages
    """
    q = question.lower()
    if any(w in q for w in ["how many more", "how many fewer", "difference", "compare"]):
        return "comparison"
    if any(w in q for w in ["fraction", "percent", "%", "half", "quarter"]):
        return "fraction"
    if any(w in q for w in ["total", "altogether", "combined", "sum", "together"]):
        word_count = len(q.split())
        if word_count < 30:
            return "one_step"
        return "multi_step"
    word_count = len(q.split())
    if word_count < 30:
        return "one_step"
    return "word_problem"


class DatasetLoader:
    @staticmethod
    def load_gsm8k(n: int = 1319, force_stub: bool = False) -> List[Dict]:
        """Loads GSM8K (Grade School Math 8K) problems — ALL 1319 test problems.

        Args:
            n: Max problems to load. Default 1319 = full test set.
            force_stub: If True, skip HuggingFace and use synthetic data.

        Returns:
            List of problems with id, prompt, expected, dataset.
        """
        problems = []
        if not force_stub and HAS_HF_DATASETS:
            try:
                ds = load_dataset("gsm8k", "main", split="test", streaming=True)
                iterator = iter(ds)
                for i in range(n):
                    try:
                        item = next(iterator)
                        answer = item["answer"]
                        # Extract the numeric answer after "####"
                        if "####" in answer:
                            numeric_answer = answer.split("####")[-1].strip()
                        else:
                            numeric_answer = answer.strip()
                        problems.append({
                            "id": f"gsm8k_{i+1}",
                            "prompt": item["question"],
                            "expected": [numeric_answer],
                            "dataset": "gsm8k",
                            "metadata": {
                                "full_answer": answer,
                                "category": _categorize_gsm8k(item["question"]),
                            }
                        })
                    except StopIteration:
                        break
                print(f"GSM8K: loaded {len(problems)} real problems from HuggingFace")
                return problems
            except Exception as e:
                print(f"GSM8K HF load failed: {e}")

        # Fallback: synthetic stubs
        print(f"GSM8K: generating {n} synthetic problems")
        for i in range(1, n + 1):
            problems.append({
                "id": f"gsm8k_{i}",
                "prompt": f"Natalia sold clips to {i*10} of her friends in April...",
                "expected": [str(i*10 + (i*10)//2)],
                "dataset": "gsm8k",
                "metadata": {"category": "synthetic", "is_stub": True}
            })
        return problems

    @staticmethod
    def load_math(n: int = 100) -> List[Dict]:
        """Loads MATH (Mathematics for Machine Learning) problems.

        Tries multiple dataset sources in order:
        1. 'hendrycks/competition_math' — correct mirror
        2. 'competition_math' — original (may not exist)
        3. LOCAL FILE: checks 'data/math_test.json' if available
        4. Stub fallback

        Args:
            n: Max problems to load. Default 100 (level 1-3).
        """
        problems = []

        sources = [
            "hendrycks/competition_math",
            "competition_math",
        ]

        if HAS_HF_DATASETS:
            for source in sources:
                try:
                    ds = load_dataset(source, split="test", streaming=True)
                    iterator = iter(ds)
                    loaded = 0
                    for i in range(n * 3):  # Load extra to filter by difficulty
                        try:
                            item = next(iterator)
                            difficulty = item.get("difficulty", 1)
                            if difficulty > 3:  # Skip level 4-5
                                continue
                            problems.append({
                                "id": f"math_{len(problems)+1}",
                                "prompt": item["problem"],
                                "expected": [item["solution"]],
                                "dataset": "math",
                                "metadata": {
                                    "difficulty": difficulty,
                                    "subject": item.get("subject", "unknown"),
                                }
                            })
                            loaded += 1
                            if loaded >= n:
                                break
                        except StopIteration:
                            break
                    if problems:
                        print(f"MATH: loaded {len(problems)} problems from {source}")
                        return problems
                except Exception as e:
                    print(f"MATH: {source} failed: {e}")
                    continue

        # Try local file
        local_path = Path("data/math_test.json")
        if local_path.exists():
            try:
                data = json.loads(local_path.read_text())
                for i, item in enumerate(data[:n]):
                    problems.append({
                        "id": f"math_{i+1}",
                        "prompt": item["problem"],
                        "expected": [item["solution"]],
                        "dataset": "math",
                        "metadata": {"difficulty": item.get("difficulty", 1)}
                    })
                print(f"MATH: loaded {len(problems)} from local file")
                return problems
            except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
                logger.warning("MATH local file load failed: %s", e)
                pass

        # Final fallback: synthetic
        print(f"MATH: generating {n} synthetic problems (all loaders failed)")
        for i in range(1, n + 1):
            problems.append({
                "id": f"math_{i}",
                "prompt": f"Solve for x: {i}x + {i*2} = {i*3 + i*2}",
                "expected": [str((i*3 + i*2 - i*2) / i) if i != 0 else "0"],
                "dataset": "math",
                "metadata": {"is_stub": True}
            })
        return problems

    @staticmethod
    def load_humaneval(n: int = 10) -> List[Dict]:
        """Loads HumanEval (Python coding problems)."""
        problems = []
        if HAS_HF_DATASETS:
            try:
                ds = load_dataset("openai_humaneval", split="test", streaming=True)
                iterator = iter(ds)
                for i in range(n):
                    try:
                        item = next(iterator)
                        problems.append({
                            "id": f"he_{i+1}",
                            "prompt": item["prompt"],
                            "expected": [item["canonical_solution"]],
                            "dataset": "humaneval",
                            "test": item.get("test", ""),
                            "entry_point": item.get("entry_point", ""),
                            "task_id": item.get("task_id", f"he_{i+1}"),
                        })
                    except StopIteration:
                        break
                return problems
            except Exception as e:
                logger.warning(f"Failed to load HumanEval from HuggingFace: {e}. Falling back to stubs.")

        # Fallback Stub
        for i in range(1, n + 1):
            problems.append({
                "id": f"he_{i}",
                "prompt": f"def add_numbers(a, b):\n    \"\"\" Add two numbers {i} times. \"\"\"",
                "expected": ["return (a + b)"],
                "dataset": "humaneval",
                "test": "def check(f): assert f(1, 2) == 3; assert f(-1, 1) == 0",
                "entry_point": "add_numbers",
                "task_id": f"he_{i}",
            })
        return problems

    @staticmethod
    def load_scibench(n: int = 10) -> List[Dict]:
        """Simulates loading SciBench (Scientific reasoning)."""
        # SciBench is not standard on HF or requires specific access/config usually.
        # Keeping as stub for reliability unless we find a specific HF path.
        # Attempting 'mit-han-lab/scibench' sometimes works but can be flaky.
        # We will stick to stub for now to avoid errors, or try a generic science dataset.
        problems = []
        for i in range(1, n + 1):
            problems.append({
                "id": f"scibench_{i}",
                "prompt": f"Calculate the kinetic energy of a {i}kg object moving at 10 m/s.",
                "expected": [str(0.5 * i * 100)],
                "dataset": "scibench"
            })
        return problems

def get_all_datasets(n_per_set: int = 5) -> List[Dict]:
    return (
        DatasetLoader.load_gsm8k(n_per_set) +
        DatasetLoader.load_math(n_per_set) +
        DatasetLoader.load_humaneval(n_per_set) +
        DatasetLoader.load_scibench(n_per_set)
    )