source/grpo_data.py

from __future__ import annotations

from collections import Counter


def build_logic_tasks() -> tuple[list[dict], list[dict]]:
    examples: list[dict] = []
    for number in range(100):
        examples.append(
            {
                "task": "parity",
                "prompt": (
                    "Output 1 if the number is even, otherwise output 0. "
                    f"Number: {number}. Answer:"
                ),
                "answer": int(number % 2 == 0),
                "key": number,
            }
        )
    for left in range(20):
        for right in range(20):
            if left == right:
                continue
            examples.append(
                {
                    "task": "less_than",
                    "prompt": (
                        "Output 1 if the first number is less than the second, "
                        f"otherwise output 0. First: {left}. Second: {right}. Answer:"
                    ),
                    "answer": int(left < right),
                    "key": left * 23 + right * 7,
                }
            )
            examples.append(
                {
                    "task": "sum_at_least_20",
                    "prompt": (
                        "Output 1 if the sum is at least 20, otherwise output 0. "
                        f"Numbers: {left} and {right}. Answer:"
                    ),
                    "answer": int(left + right >= 20),
                    "key": left * 31 + right * 11,
                }
            )
    train = [
        example for example in examples if (example["key"] + len(example["task"])) % 5 != 0
    ]
    evaluation = [
        example for example in examples if (example["key"] + len(example["task"])) % 5 == 0
    ]
    return train, evaluation


def class_counts(examples: list[dict]) -> dict[int, int]:
    return dict(Counter(example["answer"] for example in examples))