feat: use m-a-p/Code-Feedback dataset for GRPO training

train_grpo.ipynb

@@ -6,7 +6,8 @@
       "source": [
         "# GRPO Training with CodeArena RL Benchmark\n",
         "\n",
-        "This notebook demonstrates how to connect our custom `codearena-rl-benchmark` environment to HuggingFace's `trl.GRPOTrainer`."
+        "This notebook demonstrates how to connect our custom `codearena-rl-benchmark` environment to HuggingFace's `trl.GRPOTrainer`.\n",
+        "It uses the `m-a-p/Code-Feedback` dataset to train the LLM for coding debugging and improving time complexity."
       ]
     },
     {
@@ -27,7 +28,7 @@
       "outputs": [],
       "source": [
         "import torch\n",
-        "from datasets import Dataset\n",
+        "from datasets import load_dataset\n",
         "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
         "from trl import GRPOConfig, GRPOTrainer\n",
         "import httpx\n",
@@ -86,13 +87,24 @@
         "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
         "tokenizer.pad_token = tokenizer.eos_token\n",
         "\n",
-        "# Sample training dataset (prompts extracted from tasks)\n",
-        "# In a real setup, you'd reset the env for each prompt to get the initial buggy_code.\n",
-        "dataset = Dataset.from_dict({\n",
-        "    \"prompt\": [\n",
-        "        \"Fix this Python code:\\ndef average_list(numbers)\\n    if length(numbers) == 0:\\n        return 0\\n    return sum(numbers) / length(numbers)\"\n",
-        "    ]\n",
-        "})\n",
+        "# Load dataset for Coding Debugging and Time Complexity Optimization\n",
+        "dataset = load_dataset(\"m-a-p/Code-Feedback\", split=\"train\")\n",
+        "\n",
+        "def format_prompt(example):\n",
+        "    # m-a-p/Code-Feedback contains 'messages' with user and assistant roles\n",
+        "    messages = example.get('messages', [])\n",
+        "    user_query = \"\"\n",
+        "    if messages and len(messages) > 0 and messages[0].get('role') == 'user':\n",
+        "        user_query = messages[0].get('content', '')\n",
+        "    \n",
+        "    prompt = f\"Optimize and debug this code to improve time complexity:\\n{user_query}\"\n",
+        "    return {\"prompt\": prompt}\n",
+        "\n",
+        "dataset = dataset.map(format_prompt)\n",
+        "# Keep only the prompt column for the trainer\n",
+        "dataset = dataset.select_columns([\"prompt\"])\n",
+        "# Limit for demo purposes\n",
+        "dataset = dataset.select(range(100))\n",
         "\n",
         "# Initialize GRPO Trainer\n",
         "training_args = GRPOConfig(\n",