import tempfile
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
import random

config_name = "standard_prompt_only"
dataset = load_dataset("trl-internal-testing/zen", config_name, split="train")

def random_reward(completions, **kwargs):
    return [random.choice([0.0, 1.0]) for _ in completions]

with tempfile.TemporaryDirectory() as tmp_dir:
    training_args = GRPOConfig(
        output_dir=tmp_dir,
        learning_rate=0.1,  # increase the learning rate to speed up the test
        per_device_train_batch_size=6,  # reduce the batch size to reduce memory usage
        num_generations=3,  # reduce the number of generations to reduce memory usage
        max_completion_length=8,  # reduce the completion length to reduce memory usage
        log_completions=True,
        logging_steps=3,
    )
    trainer = GRPOTrainer(
        model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
        reward_funcs=random_reward,
        args=training_args,
        train_dataset=dataset,
    )

    trainer.train()