Create ecoeval/datasets.py

ecoeval/datasets.py

@@ -0,0 +1,91 @@
+# ecoeval/datasets.py
+from typing import List
+from datasets import Dataset
+
+
+def _tiny_python_benchmark() -> Dataset:
+    """
+    Very small internal dataset of Python tasks with prompts + unit tests.
+    Each row has: prompt, test_code
+    """
+    tasks = [
+        {
+            "prompt": "Write a Python function add(a, b) that returns their sum.",
+            "test_code": "assert add(1, 2) == 3\nassert add(-1, 5) == 4",
+        },
+        {
+            "prompt": "Write a Python function is_prime(n) that returns True if n is a prime number, False otherwise.",
+            "test_code": (
+                "assert is_prime(2)\n"
+                "assert is_prime(3)\n"
+                "assert not is_prime(4)\n"
+                "assert is_prime(17)\n"
+                "assert not is_prime(1)"
+            ),
+        },
+        {
+            "prompt": "Write a Python function factorial(n) that returns n! for non-negative n.",
+            "test_code": (
+                "assert factorial(0) == 1\n"
+                "assert factorial(3) == 6\n"
+                "assert factorial(5) == 120"
+            ),
+        },
+        {
+            "prompt": "Write a Python function reverse_string(s) that returns the reversed string.",
+            "test_code": (
+                "assert reverse_string('abc') == 'cba'\n"
+                "assert reverse_string('') == ''\n"
+                "assert reverse_string('racecar') == 'racecar'"
+            ),
+        },
+        {
+            "prompt": "Write a Python function fibonacci(n) that returns a list of the first n Fibonacci numbers, starting from 0.",
+            "test_code": (
+                "assert fibonacci(1) == [0]\n"
+                "assert fibonacci(5) == [0, 1, 1, 2, 3]\n"
+                "assert fibonacci(0) == []"
+            ),
+        },
+        {
+            "prompt": "Write a Python function count_vowels(s) that returns the number of vowels in the string s.",
+            "test_code": (
+                "assert count_vowels('hello') == 2\n"
+                "assert count_vowels('xyz') == 0\n"
+                "assert count_vowels('AEiou') == 5"
+            ),
+        },
+        {
+            "prompt": "Write a Python function flatten(lst) that takes a list of lists of integers and returns a single flattened list.",
+            "test_code": (
+                "assert flatten([[1,2],[3,4]]) == [1,2,3,4]\n"
+                "assert flatten([]) == []\n"
+                "assert flatten([[1],[2],[3]]) == [1,2,3]"
+            ),
+        },
+        {
+            "prompt": "Write a Python function is_palindrome(s) that returns True if s is a palindrome, ignoring case and spaces.",
+            "test_code": (
+                "assert is_palindrome('racecar')\n"
+                "assert is_palindrome('RaceCar')\n"
+                "assert is_palindrome('nurses run')\n"
+                "assert not is_palindrome('hello')"
+            ),
+        },
+    ]
+    return Dataset.from_list(tasks)
+
+
+_AVAILABLE_DATASETS = {
+    "tiny-python-benchmark": _tiny_python_benchmark,
+}
+
+
+def list_available_datasets() -> List[str]:
+    return sorted(_AVAILABLE_DATASETS.keys())
+
+
+def load_dataset_by_name(name: str) -> Dataset:
+    if name not in _AVAILABLE_DATASETS:
+        raise ValueError(f"Unknown dataset: {name}")
+    return _AVAILABLE_DATASETS[name]()