def train()

in train_rick/run_grpo.py [0:0]

• 29 lines of code
• 2 McCabe index (conditional complexity)

def train():
    dataset = load_dataset("qgallouedec/rick-physics-grpo", split="train")

    def format_dataset(example):
        return {"prompt": [{"role": "user", "content": example["question"]}]}

    dataset = dataset.map(format_dataset)

    args = GRPOConfig(
        max_completion_length=512,
        per_device_train_batch_size=64,
        gradient_accumulation_steps=4,
        num_train_epochs=10,
        num_generations=16,
        mask_truncated_completions=True,
        # Speedup and reduce memory
        gradient_checkpointing=True,
        bf16=True,
        use_vllm=True,
        output_dir="data/SmolLM2-360M-Rickified-GRPO",
        # Logging
        run_name="SmolLM2-360M-Rickified-GRPO",
        logging_steps=2,
        log_completions=True,
        num_completions_to_print=1,
    )

    trainer = GRPOTrainer(
        model="qgallouedec/SmolLM2-360M-Rickified",
        reward_funcs=[format_reward, correctness_reward],
        train_dataset=dataset,
        args=args,
    )
    trainer.train()
    trainer.push_to_hub(dataset_name="qgallouedec/rick-physics-grpo")