code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- omit in toc -->\n",
    "# Train DeepSeek R1 from Scratch\n",
    "\n",
    "![Python](https://img.shields.io/badge/Python-3.10-blue) [![Medium](https://img.shields.io/badge/Medium-Read%20Now-red?logo=medium)](https://medium.com/@fareedkhandev) ![Contributions](https://img.shields.io/badge/Contributions-Welcome-yellow) \n",
    "\n",
    "\n",
    "The entire training process of DeepSeek R1 is nothing but using different way of reinforcement learning on top of their base model (i.e. [deepseek V3](https://huggingface.co/deepseek-ai/DeepSeek-V3))\n",
    "\n",
    "To make everything easy to understand we will use hand drawn flowcharts along with the code and will follow the step by step implementation of deepseek technical report and will build our own model using a tiny base model that you can also run locally.\n",
    "\n",
    "We will also be covering theory next to each step, so in case you are not aware of DeepSeek R1 this blog will cover that too.\n",
    "\n",
    "I explained **DeepSeek R1** with **hand-drawn diagrams** for non-tech folks. [Read here](https://levelup.gitconnected.com/drawing-deepseek-r1-architecture-and-training-process-from-scratch-72043da33955).\n",
    "\n",
    "<!-- omit in toc -->\n",
    "## Table of Contents\n",
    "- [Setting up the Stage](#setting-up-the-stage)\n",
    "- [Our Training Dataset](#our-training-dataset)\n",
    "- [DeepSeek R1 Training Quick Overview](#deepseek-r1-training-quick-overview)\n",
    "- [Choosing our Base Model](#choosing-our-base-model)\n",
    "- [Policy Model (R) In RL Setup](#policy-model-r-in-rl-setup)\n",
    "- [GRPO Algorithm for R1 Zero](#grpo-algorithm-for-r1-zero)\n",
    "- [Prompt Template](#prompt-template)\n",
    "- [Preprocessing Training Data](#preprocessing-training-data)\n",
    "- [Reward Functions](#reward-functions)\n",
    "  - [Accuracy Reward](#accuracy-reward)\n",
    "  - [Format Reward](#format-reward)\n",
    "  - [Reasoning Steps Reward](#reasoning-steps-reward)\n",
    "  - [Cosine Scaled Reward](#cosine-scaled-reward)\n",
    "  - [Repetition Penalty Reward](#repetition-penalty-reward)\n",
    "- [Training Configurations for R1 Zero](#training-configurations-for-r1-zero)\n",
    "- [GRPO Training Loop](#grpo-training-loop)\n",
    "- [Saving Tiny R1 Zero LLM](#saving-tiny-r1-zero-llm)\n",
    "- [Two main problems with R1 Zero](#two-main-problems-with-r1-zero)\n",
    "- [Preparing Cold Start Data for SFT](#preparing-cold-start-data-for-sft)\n",
    "- [Few-shot Prompting with Long CoT](#few-shot-prompting-with-long-cot)\n",
    "- [Direct Prompting](#direct-prompting)\n",
    "- [Post Processing Refinement](#post-processing-refinement)\n",
    "- [SFT Stage 1 With Cold Start Data](#sft-stage-1-with-cold-start-data)\n",
    "- [Stage 1 SFT Trainer Configs for R1](#stage-1-sft-trainer-configs-for-r1)\n",
    "- [Stage 1 STF Training Loop](#stage-1-stf-training-loop)\n",
    "- [Saving Tiny R1 LLM](#saving-tiny-r1-llm)\n",
    "- [Reasoning-Oriented RL](#reasoning-oriented-rl)\n",
    "- [Rejection Sampling](#rejection-sampling)\n",
    "- [SFT Stage 2 Training](#sft-stage-2-training)\n",
    "- [Distillation](#distillation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Stage\n",
    "\n",
    "Once you install the required libraries, you can start by importing the necessary libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import re\n",
    "import math\n",
    "from dataclasses import dataclass, field\n",
    "from typing import List, Optional\n",
    "\n",
    "# Import PyTorch and Hugging Face Transformers\n",
    "import torch\n",
    "import transformers\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    HfArgumentParser,\n",
    "    TrainingArguments,\n",
    "    set_seed,\n",
    "    TrainerCallback,\n",
    "    TrainerControl,\n",
    "    TrainerState,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "\n",
    "# Import dataset utilities\n",
    "import datasets\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Import libraries from TRL (Transformers Reinforcement Learning)\n",
    "from trl import (\n",
    "    AutoModelForCausalLMWithValueHead, \n",
    "    PPOConfig, \n",
    "    PPOTrainer, \n",
    "    GRPOTrainer, \n",
    "    GRPOConfig, \n",
    "    SFTTrainer\n",
    ")\n",
    "\n",
    "# Import math-related utilities\n",
    "from latex2sympy2_extended import NormalizationConfig\n",
    "from math_verify import LatexExtractionConfig, parse, verify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Our Training Dataset\n",
    "\n",
    "Although the paper doesn’t specify the exact initial dataset for RL pre-training, we assume it should be reasoning focused.\n",
    "\n",
    "So to stay as close as possible to the original replication, we will use these two open-source reasoning Hugging Face datasets:\n",
    "\n",
    " 1. [NuminaMath-TIR](https://huggingface.co/datasets/AI-MO/NuminaMath-TIR) (For R1 Zero Training)\n",
    "\n",
    " 2. [Bespoke-Stratos-17k](https://huggingface.co/datasets/bespokelabs/Bespoke-Stratos-17k) (For R1 Training)\n",
    "\n",
    "AI-MO/NuminaMath-TIR contains 70K math problems with messages column showing the COT (chain of though) reasoning behind the solution.\n",
    "\n",
    "| Field    | Description |  \n",
    "|----------|------------|  \n",
    "| problem  | The math problem |  \n",
    "| solution | Step-by-step solution |  \n",
    "| messages    | Chat to solve the problem |\n",
    "\n",
    "Take a look at its sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the \"AI-MO/NuminaMath-TIR\" dataset from DigitalLearningGmbH\n",
    "MATH_le = load_dataset(\"AI-MO/NuminaMath-TIR\", \"default\")  \n",
    "\n",
    "# Access the first sample in the training set\n",
    "MATH_le['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While Bespoke-Stratos contains 17K problems focused on math and code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the \"Bespoke-Stratos-17k\" dataset from bespokelabs\n",
    "bespoke_rl = load_dataset(\"bespokelabs/Bespoke-Stratos-17k\", \"default\") \n",
    "\n",
    "# Access the first sample in the training set\n",
    "bespoke_rl['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It’s not necessary to choose these datasets only, you can select any of your choice as long as it is reasoning focused (**a problem and its step-by-step solution**)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DeepSeek R1 Training Quick Overview\n",
    "\n",
    "So, before going into the technical implementation, a quick overview is that DeepSeek-R1 isn’t trained from scratch, like, from nothing. Instead, they started with a pretty smart LLM they already had [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) but they wanted to make it a reasoning superstar.\n",
    "\n",
    "![DeepSeek R1 Implementation Quick Overview (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/5872/1*XhE5MVuPfOrrbBxgYtHMsg.png)\n",
    "\n",
    "To do that, they used **Reinforcement Learning**, or RL for short where you reward the LLM when it does something good for reasoning while punish it otherwise.\n",
    "\n",
    "But it’s not just one simple training session. It’s like a whole bunch of steps, a pipeline they call it. They first tried just pure **RL** to see if reasoning would pop up by itself **that was DeepSeek-R1-Zero**, kinda an experiment. Then for the **real DeepSeek-R1**, they made it more organized with different stages. They give it some starting data to get it going, then do RL, then more data, then more RL… it’s like leveling up, step by step!\n",
    "\n",
    "The whole point is to make these language models way better at thinking through problems.\n",
    ">So yeah, that’s the super short version before we look into the crazy details of each step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Choosing our Base Model\n",
    "\n",
    "Since DeepSeek team chose DeepSeek-V3 as their base model to create R1 Zero and R1, but it’s quite huge **685 GB 💀 in size** which is obviously not in our reach.\n",
    "\n",
    "To keep it simple, we will use a much smaller base model [Qwen/Qwen2.5–0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct) (0.9 GB in size). If you have a higher GPU RAM that can even load unquantized LLMs, you can go for a bigger model, such as [Qwen/Qwen2.5–7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct).\n",
    "\n",
    "Let’s take a look at some of the specification of our base model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
    "OUTPUT_DIR = \"data/Qwen-GRPO-training\" # For saving our trained model\n",
    "\n",
    "# Create output directory if it doesn't exist\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# Initialize tokenizer with chat template\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    trust_remote_code=True,\n",
    "    padding_side=\"right\"\n",
    ")\n",
    "\n",
    "# Set pad token if not set\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "print(f\"Vocabulary size: {len(tokenizer)}\")\n",
    "print(f\"Model max length: {tokenizer.model_max_length}\")\n",
    "print(f\"Pad token: {tokenizer.pad_token}\")\n",
    "print(f\"EOS token: {tokenizer.eos_token}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are some basic info about the model, take a look at the total number of parameters our base model has."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize base model\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    trust_remote_code=True,\n",
    "    torch_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "print(f\"Model parameters: {model.num_parameters():,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Close to 0.5B params, let’s print a simple response from it and then we will move on to next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check CUDA availability\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Move model to the appropriate device\n",
    "model.to(device)\n",
    "\n",
    "# Test basic inference\n",
    "def test_model_inference(user_input: str):\n",
    "    \"\"\"Test basic model inference with the loaded model and tokenizer.\"\"\"\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are Qwen, a helpful assistant.\"},\n",
    "        {\"role\": \"user\", \"content\": user_input}\n",
    "    ]\n",
    "\n",
    "    # Apply chat template\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True\n",
    "    )\n",
    "\n",
    "    # Tokenize and generate\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\").to(device)\n",
    "    outputs = model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=100,\n",
    "        do_sample=True,\n",
    "        temperature=0.7\n",
    "    )\n",
    "\n",
    "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    return response\n",
    "\n",
    "# Test the model\n",
    "test_input = \"how are you?\"\n",
    "response = test_model_inference(test_input)\n",
    "print(f\"Test Input: {test_input}\")\n",
    "print(f\"Model Response: {response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, the output of this tiny model is quite reliable and suitable for our DeepSeek lookalike model training for sure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policy Model (R) In RL Setup\n",
    "\n",
    "Now that we have selected our base model, next we need to understand how a basic RL setup works for training an LLM.\n",
    "\n",
    "For DeepSeek R1 their starting point was (DeepSeek V3) base model and in our case we are starting with Qwen2.5–0.5B-Instruct. By a starting point I meant that **it has created the DeepSeek R1 zero version**, an initial version which has some errors in it before the final version was created.\n",
    "\n",
    "The initial version (R1 Zero) was created using Reinforcement Learning where (DeepSeek v3/Qwen2.5–0.5B) acts as an RL agent (actor who takes action). Let’s first visualize how it works.\n",
    "\n",
    "![Qwen 2.5 as an agent workflow (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/5410/1*S6YIXu1vIVmQFl-DgRFktg.png)\n",
    "\n",
    "The RL agent (DeepSeek V3/Qwen2–0.5B) starts by taking an **Action**, which means it generates an answer and some reasoning for a given problem that’s put into its **Environment**. The Environment, in this case, is simply the reasoning task itself.\n",
    "\n",
    "After taking an action, the Environment gives back a **Reward**. This Reward is like feedback, it tells our base model (DeepSeek V3/Qwen2–0.5B) how good its action was. A positive Reward means it did something right, maybe got the answer correct or reasoned well. This feedback signal then goes back to our base model, helping it learn and adjust how it takes actions in the future to get even better Rewards.\n",
    "> In the next section, we will be discussing this methodology in more detail"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GRPO Algorithm for R1 Zero\n",
    "\n",
    "So that we have understand a basic RL flow now we need to learn what exact RL algorithm DeepSeek uses for R1-Zero.\n",
    "\n",
    "There are many RL algos available, but traditional RL use something called a **“critic” **to help the main decision making part (“actor” i.e. DeepSeek-V3/Qwen2-0.5B). This critic is usually just as big and complex as the actor itself, which basically doubles the amount of computational cost.\n",
    "\n",
    "But DeepSeek uses GRPO for training their initial (R1 Zero), **GRPO** does things differently because it figures out a baseline, a kind of reference point for good actions directly from the results it gets from a **group** of actions. Because of this, GRPO doesn’t need a separate critic model at all. This saves a lot of computation and makes things more efficient.\n",
    "\n",
    "Let’s draw a flowchart of how GRPO is being used for R1 Zero training, and then we will **interpretate** it.\n",
    "\n",
    "![GRPO Flow for DeepSeek R1 Zero (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/6404/1*8mfNzi-gvasR7mSaseswmg.png)\n",
    "\n",
    "Let’s understand how DeepSeek GRPO implementation works with our base model (Qwen2–0.5B). \n",
    "\n",
    "First, the **Problem Input (A)** is given to the **Qwen Model (B)**, Qwen attempts to generate an answer through **Generate Completion (C)**. The final result, called the **Completion Output (D)**, includes reasoning steps in <think> tags and the final solution in <answer> tags.\n",
    "\n",
    "Next, the **Problem Input (A)** and the **Ground Truth Solution (E)** are fed into the **Reward Functions (F)**, acting as intelligent graders. These functions compare Qwen **Completion Output (D)** with the correct solution and evaluate different aspects such as:\n",
    "\n",
    " 1. **Accuracy** (is the answer correct?)\n",
    "\n",
    " 2. **Format** (are the <think> and <answer> tags used properly?)\n",
    "\n",
    " 3. **Reasoning Steps** (is the logic clear?)\n",
    "\n",
    " 4. **Cosine Scaling** (is the response concise?)\n",
    "\n",
    " 5. **Repetition Penalty** (is there unnecessary repetition?).\n",
    "\n",
    "These evaluations produce **Reward Scores (G)**, which are then passed to the **GRPO Trainer (H)**. The trainer uses gradients to adjust the **Qwen Model (B)**, fine-tuning how it generates answers. This process is called **Gradient Reward Policy Optimization** because it optimizes Qwen responses using **gradients**, **reward feedback**, and **policy adjustments** to maximize performance.\n",
    "\n",
    "Finally, the updated **Qwen Model (B)** is tested again on new problems, continuously refining itself through repeated cycles. With each iteration, Qwen becomes a better problem solver.\n",
    "\n",
    "> In the upcoming section we will start preprocessing our training dataset for GRPO training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Template\n",
    "\n",
    "We are using the same thinking prompt template that DeepSeek uses for the GRPO algorithm to build R1 Zero, so let’s define that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DeepSeek system prompt for GRPO based training\n",
    "SYSTEM_PROMPT = (\n",
    "    \"A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant \"\n",
    "    \"first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning \"\n",
    "    \"process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., \"\n",
    "    \"<think> reasoning process here </think><answer> answer here </answer>\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This **system prompt** tells the base model (Qwen2–0.5B) its role as a helpful assistant who reasons step-by-step before answering.\n",
    "\n",
    "The `<think>` and `<answer>` tags are used to structure the model response, separating its internal reasoning from the final answer for better evaluation and reward.\n",
    "\n",
    "## Preprocessing Training Data\n",
    "\n",
    "Now that we have our system prompt ready, we need to transform our training data according to our template.\n",
    "\n",
    "![Preprocessing dataset overview (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/6160/1*XnM7v4dPD4LtyAh2MLuInA.png)\n",
    "\n",
    "We need to create the make_conversation function that will handle the conversation for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to structure the training data\n",
    "def make_conversation(example):\n",
    "    \"\"\"Convert dataset examples into conversation format.\"\"\"\n",
    "    return {\n",
    "        \"prompt\": [\n",
    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "            {\"role\": \"user\", \"content\": example[\"problem\"]},\n",
    "        ],\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It will take each problem column value from our training dataset and return a dictionary with the system prompt and the appended problem question for each row. Let’s create this function that will prepare our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and prepare dataset\n",
    "def load_math_dataset():\n",
    "    \"\"\"Load and prepare the mathematics dataset.\"\"\"\n",
    "    dataset = load_dataset(\n",
    "        \"AI-MO/NuminaMath-TIR\",\n",
    "        name=\"default\",\n",
    "        split=['train', 'test']\n",
    "    )\n",
    "    \n",
    "    # Convert splits into dictionary\n",
    "    dataset = {\n",
    "        'train': dataset[0],\n",
    "        'test': dataset[1]\n",
    "    }\n",
    "    \n",
    "    # Apply conversation format\n",
    "    for split in dataset:\n",
    "        dataset[split] = dataset[split].map(make_conversation)\n",
    "\n",
    "        # Remove 'messages' column if exists\n",
    "        if \"messages\" in dataset[split].column_names:\n",
    "            dataset[split] = dataset[split].remove_columns(\"messages\")\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have everything ready, let’s transform our training data into the required format and print the training and test size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load our training dataset and printing train/test size\n",
    "dataset = load_math_dataset()\n",
    "\n",
    "print(f\"Train set size: {len(dataset['train'])}\")\n",
    "print(f\"Test set size: {len(dataset['test'])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have split our training dataset, we need to validate our dataset (**Check if user/assistant conversation exist**) before moving to the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_dataset(dataset):\n",
    "    \"\"\"Perform basic validation checks on the dataset.\"\"\"\n",
    "    \n",
    "    # Define the required fields for the dataset\n",
    "    required_fields = [\"problem\", \"prompt\"]\n",
    "\n",
    "    # Loop through the 'train' and 'test' splits of the dataset\n",
    "    for split in ['train', 'test']:\n",
    "        print(f\"\\nValidating {split} split:\")\n",
    "\n",
    "        # Retrieve column names from the dataset\n",
    "        fields = dataset[split].column_names\n",
    "\n",
    "        # Check if any required fields are missing\n",
    "        missing = [field for field in required_fields if field not in fields]\n",
    "        if missing:\n",
    "            print(f\"Warning: Missing fields: {missing}\")  # Warn if fields are missing\n",
    "        else:\n",
    "            print(\"✓ All required fields present\")  # Confirm all fields are present\n",
    "\n",
    "        # Retrieve the first sample from the dataset split\n",
    "        sample = dataset[split][0]\n",
    "\n",
    "        # Extract the 'prompt' field, which contains a list of messages\n",
    "        messages = sample['prompt']\n",
    "\n",
    "        # Validate the prompt format:\n",
    "        # - It should contain at least two messages\n",
    "        # - The first message should be from the 'system' role\n",
    "        # - The second message should be from the 'user' role\n",
    "        if (len(messages) >= 2 and\n",
    "            messages[0]['role'] == 'system' and\n",
    "            messages[1]['role'] == 'user'):\n",
    "            print(\"✓ Prompt format is correct\")  # Confirm correct format\n",
    "        else:\n",
    "            print(\"Warning: Incorrect prompt format\")  # Warn if format is incorrect\n",
    "\n",
    "# Validate dataset\n",
    "validate_dataset(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our training dataset is validated successfully 🙌, it means we have successfully transformed our dataset for training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Functions\n",
    "\n",
    "We already saw in GRPO section that it evaluate the answer of base model through five different ways:\n",
    "\n",
    "![Reward Functions (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/7474/1*kJln8i6Tv4aspnTfMoRW-Q.png)\n",
    "\n",
    " 1. **Accuracy** (is the answer correct?)\n",
    "\n",
    " 2. **Format** (are the `<think>` and `<answer>` tags used properly?)\n",
    "\n",
    " 3. **Reasoning Steps** (is the logic clear?)\n",
    "\n",
    " 4. **Cosine Scaling** (is the response concise?)\n",
    "\n",
    " 5. **Repetition Penalty** (is there unnecessary repetition?).\n",
    "\n",
    "Each of these are functions will calculate the reward for each response, and we need to code them. So, let’s do that first."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Accuracy Reward\n",
    "\n",
    "Accuracy reward is the most easy to understand but requires a bit complex code. In this reward model we want to check if mathematically our base model response is equivalent to the ground truth solution.\n",
    "\n",
    "![Accuracy Reward (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/7860/1*A3tW-OZSZ4m10EEzogjy8Q.png)\n",
    "\n",
    "If the model answer is mathematically correct, we assign a reward of **1.0**. If it is incorrect, the reward is **0.0**. In cases where the ground truth solution cannot be parsed, we assign a neutral reward of **0.5** to avoid unfair penalties.\n",
    "\n",
    "Now, let’s implement the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy_reward(completions, **kwargs):\n",
    "    \"\"\"\n",
    "    Reward function to check if the model's response is mathematically \n",
    "    equivalent to the ground truth solution.\n",
    "    Uses latex2sympy2 for parsing and math_verify for validation.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Extract responses\n",
    "    contents = [completion[0][\"content\"] for completion in completions]\n",
    "    rewards = []\n",
    "\n",
    "    solutions = kwargs.get(\"solution\") # Get solutions from kwargs\n",
    "    \n",
    "    for content, sol in zip(contents, solutions):\n",
    "        # Parse the ground truth solution\n",
    "        gold_parsed = parse(sol, extraction_mode=\"first_match\", \n",
    "                            extraction_config=[LatexExtractionConfig()])\n",
    "        \n",
    "        if gold_parsed:  # Check if parsing was successful\n",
    "            # Parse the model's answer with relaxed normalization\n",
    "            answer_parsed = parse(\n",
    "                content,\n",
    "                extraction_config=[\n",
    "                    LatexExtractionConfig(\n",
    "                        normalization_config=NormalizationConfig(\n",
    "                            nits=False,\n",
    "                            malformed_operators=False,\n",
    "                            basic_latex=True,\n",
    "                            equations=True,\n",
    "                            boxed=\"all\",\n",
    "                            units=True,\n",
    "                        ),\n",
    "                        boxed_match_priority=0,\n",
    "                        try_extract_without_anchor=False,\n",
    "                    )\n",
    "                ],\n",
    "                extraction_mode=\"first_match\",\n",
    "            )\n",
    "\n",
    "            # Reward 1.0 if correct, 0.0 if incorrect\n",
    "            reward = float(verify(answer_parsed, gold_parsed))\n",
    "        else:\n",
    "            # If ground truth cannot be parsed, assign neutral reward (0.5)\n",
    "            reward = 0.5\n",
    "            print(\"Warning: Failed to parse gold solution:\", sol)\n",
    "\n",
    "        rewards.append(reward)\n",
    "    \n",
    "    return rewards"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this function, we check whether the model response is **equivalent** to the correct answer. Instead of comparing raw text, we:\n",
    "\n",
    " 1. Convert the solution into a structured mathematical format using **latex2sympy2**.\n",
    "\n",
    " 2. If parsing fails, assign a neutral reward of **0.5**.\n",
    "\n",
    " 3. Extract the model output and normalize it for better robustness.\n",
    "\n",
    " 4. Use **math_verify** to check if the parsed response matches the parsed solution.\n",
    "\n",
    " 5. If correct assign **1,** if incorrect assign **0**.\n",
    "\n",
    "This ensures that accuracy evaluation is not just about textual similarity but **true mathematical correctness.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Format Reward\n",
    "\n",
    "Format Reward is all about making sure our model follows instructions and structures its output correctly. We asked it to put its reasoning in `<think>` tags and the final answer in `<answer>` tags, right? This reward function checks exactly that!\n",
    "\n",
    "![Forward Reward (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/6620/1*DbUraziwiOoAj6SvtSJmpw.png)\n",
    "\n",
    "If the model uses those tags correctly, we give it a reward of 1. If it messes up the format, it gets 0. Simple as that! This encourages the model to pay attention to the output structure we want.\n",
    "\n",
    "Let’s code this up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement Format Reward Function\n",
    "def format_reward(completions, **kwargs):\n",
    "    \"\"\"\n",
    "    Reward function to check if the completion has the correct format:\n",
    "    <think>...</think> <answer>...</answer>.\n",
    "    \"\"\"\n",
    "    # Define the regex pattern for the desired format\n",
    "    pattern = r\"^<think>.*?</think>\\s*<answer>.*?</answer>$\"\n",
    "\n",
    "    # Extract the content from each completion\n",
    "    completion_contents = [completion[0][\"content\"] for completion in completions]\n",
    "\n",
    "    # Check if each completion matches the pattern\n",
    "    matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE)\n",
    "               for content in completion_contents]\n",
    "\n",
    "    # Reward 1.0 for correct format, 0.0 otherwise\n",
    "    return [1.0 if match else 0.0 for match in matches]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this function:\n",
    "\n",
    "* We define a pattern using regular expressions (regex). This pattern basically says “the content should *start* with <think>, have *anything* inside until </think>, then some *spaces*, then <answer>, *anything* inside until </answer>, and then *end* there”.\n",
    "\n",
    "* We get the actual text content from each model completion.\n",
    "\n",
    "* Then we use use re.match to see if each content perfectly matches our pattern. re.DOTALL helps the . in regex match newlines too, and re.MULTILINE makes ^ and $ match the start/end of the whole string, not just lines.\n",
    "\n",
    "* Finally, we give a reward 1 if it matched the format perfectly, 0 if it didn’t. This is a strict on/off reward for format correctness."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reasoning Steps Reward\n",
    "\n",
    "Reasoning Steps Reward is a bit clever. We want to encourage our model to show its **“thinking process”**. So, we are going to reward it for including things that *look like* reasoning steps.\n",
    "\n",
    "![Reasoning Steps Reward Encouragement (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/5406/1*hx0sAVnY58WOYw6rGF64ug.png)\n",
    "\n",
    "We will look for keywords and patterns that usually show up in step-by-step reasoning, like:\n",
    "\n",
    "* Step 1, Step 2, etc.\n",
    "\n",
    "* Numbered lists like 1, 2\n",
    "\n",
    "* Bullet points like - or *\n",
    "\n",
    "* Transition words like First, Second, Next, Finally\n",
    "\n",
    "The more of these it includes, the better the reward. It’s like giving points for showing its work!\n",
    "\n",
    "Let’s code this reasoning encouraging function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reasoning_steps_reward(completions, **kwargs):\n",
    "    r\"\"\"\n",
    "    Reward function to encourage clear step-by-step reasoning.\n",
    "    It looks for patterns like \"Step 1:\", numbered lists, bullet points,\n",
    "    and transition words.\n",
    "    \"\"\"\n",
    "    # Regex pattern to find indicators of reasoning steps\n",
    "    pattern = r\"(Step \\d+:|^\\d+\\.|\\n-|\\n\\*|First,|Second,|Next,|Finally,)\"\n",
    "\n",
    "    # Extract completion contents\n",
    "    completion_contents = [completion[0][\"content\"] for completion in completions]\n",
    "\n",
    "    # Count the number of reasoning step indicators in each completion\n",
    "    matches = [len(re.findall(pattern, content, re.MULTILINE))\n",
    "               for content in completion_contents]\n",
    "\n",
    "    # Reward is proportional to the number of reasoning steps, maxing out at 1.0\n",
    "    # We're using a \"magic number\" 3 here - encourage at least 3 steps for full reward\n",
    "    return [min(1.0, count / 3) for count in matches]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "We create a pattern that’s a bit more complex regex. It looks for all those reasoning indicator things we listed above.\n",
    "\n",
    "We use re.findall to find *all* the matches of our pattern within each content. `len(re.findall(…))` then gives us the *count* of these indicators.\n",
    "\n",
    "The reward is calculated as min(1.0, count / 3). This means\n",
    "\n",
    "* If it finds 3 or more reasoning indicators ( count >= 3), the reward is 1.0 (max reward).\n",
    "\n",
    "* If it finds fewer (e.g., count = 1 or 2), it gets a *partial* reward (like 1/3 or 2/3).\n",
    "\n",
    "* If it finds none (count = 0), the reward is 0.0.\n",
    "\n",
    "The / 3 is a bit of a magic number here. We’re saying **“aim for about 3 reasoning steps to get full credit”** You can tweak this number if you want to encourage more or fewer steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cosine Scaled Reward\n",
    "\n",
    "Cosine Scaled Reward is a bit more advanced. It’s about encouraging *conciseness* in correct answers and being *less harsh* on longer incorrect answers.\n",
    "\n",
    "![Cosine Scaling Concept (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/7094/1*WmG8r1OVeU4R3jObAy0yCg.png)\n",
    "\n",
    "Think of it like this:\n",
    "\n",
    "* **For correct answers:** We want to reward *shorter*, more direct solutions more than long, rambling ones. A short, correct answer is often better.\n",
    "\n",
    "* **For incorrect answers:** A short, wrong answer is probably worse than a longer, wrong answer that at least *tried* to reason. So, we want to penalize short wrong answers *more* than long wrong answers.\n",
    "\n",
    "Let’s see the code that does this clever scaling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement Cosine Scaled Reward Function\n",
    "def get_cosine_scaled_reward(\n",
    "    min_value_wrong: float = -0.5,\n",
    "    max_value_wrong: float = -0.1,\n",
    "    min_value_correct: float = 0.8,\n",
    "    max_value_correct: float = 1.0,\n",
    "    max_len: int = 1000,\n",
    "):\n",
    "    \"\"\"\n",
    "    Returns a cosine scaled reward function. This function scales the accuracy reward\n",
    "    based on completion length. Shorter correct solutions get higher rewards,\n",
    "    longer incorrect solutions get less penalty.\n",
    "    \"\"\"\n",
    "    def cosine_scaled_reward(completions, solution, accuracy_rewards, **kwargs):\n",
    "        \"\"\"\n",
    "        Cosine scaled reward function that adjusts accuracy rewards based on completion length.\n",
    "        \"\"\"\n",
    "        contents = [completion[0][\"content\"] for completion in completions]\n",
    "        rewards = []\n",
    "\n",
    "        for content, sol, acc_reward in zip(contents, solution, accuracy_rewards):\n",
    "            gen_len = len(content)  # Length of the generated answer\n",
    "            progress = gen_len / max_len # How far we are to max length\n",
    "            cosine = math.cos(progress * math.pi) # Cosine value based on progress\n",
    "\n",
    "            if acc_reward > 0.5: # Assuming accuracy_reward gives ~1.0 for correct answers\n",
    "                min_value = min_value_correct\n",
    "                max_value = max_value_correct\n",
    "            else: # Incorrect answer\n",
    "                min_value = max_value_wrong  # Note the swap!\n",
    "                max_value = min_value_wrong\n",
    "\n",
    "            # Cosine scaling formula!\n",
    "            reward = min_value + 0.5 * (max_value - min_value) * (1.0 + cosine)\n",
    "            rewards.append(float(reward))\n",
    "        return rewards\n",
    "    return cosine_scaled_reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`get_cosine_scaled_reward(...)` generates a reward function for training, customizing scaling with parameters like min_value_wrong/max_value_wrong (penalty range for incorrect answers) and min_value_correct/max_value_correct (reward range for correct ones). max_len sets the maximum length for scaling.\n",
    "\n",
    "Inside, `cosine_scaled_reward(...)` we calculate rewards based on completions, solution, and accuracy_rewards.\n",
    "\n",
    "It computes gen_len, normalizes it as progress `= gen_len / max_len`, and derives a cosine value that starts at 1 (short answers) and decreases to -1 (long answers).\n",
    "\n",
    "If `acc_reward > 0.5`, it uses the correct reward range, otherwise it applies the incorrect range but swaps min/max values to penalize longer wrong answers less."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Repetition Penalty Reward\n",
    "\n",
    "Repetition Penalty Reward is all about discouraging our model from getting stuck in loops and repeating itself. We want it to generate fresh, varied reasoning and answers, not just copy-paste the same phrases over and over!\n",
    "\n",
    "![Repetition Penalty Idea (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/8608/1*9jBhiz-rI_fRGa77g9RZtQ.png)\n",
    "\n",
    "This reward function penalizes the model if it uses the same sequences of words (n-grams) too many times. We’ll use n-grams of size 3 (trigrams) in our example, but you can adjust this.\n",
    "\n",
    "If the model repeats itself a lot, it gets a negative reward (penalty). If it’s more diverse and avoids repetition, the penalty is less.\n",
    "\n",
    "Let’s implement the code to penalize repetition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_repetition_penalty_reward(ngram_size: int = 3, max_penalty: float = -0.1):\n",
    "    \"\"\"\n",
    "    Returns a repetition penalty reward function. Penalizes repetitions of n-grams\n",
    "    in the generated text.\n",
    "    \"\"\"\n",
    "    if max_penalty > 0:\n",
    "        raise ValueError(f\"max_penalty {max_penalty} should not be positive\")\n",
    "\n",
    "    def zipngram(text: str, ngram_size: int):\n",
    "        \"\"\"Helper function to generate n-grams from text.\"\"\"\n",
    "        words = text.lower().split() # Lowercase and split into words\n",
    "        return zip(*[words[i:] for i in range(ngram_size)]) # Create n-grams\n",
    "\n",
    "    def repetition_penalty_reward(completions, **kwargs) -> float:\n",
    "        \"\"\"\n",
    "        Repetition penalty reward function.\n",
    "        \"\"\"\n",
    "        contents = [completion[0][\"content\"] for completion in completions]\n",
    "        rewards = []\n",
    "        for completion in contents:\n",
    "            if completion == \"\": # No penalty for empty completions\n",
    "                rewards.append(0.0)\n",
    "                continue\n",
    "            if len(completion.split()) < ngram_size: # No penalty for short completions\n",
    "                rewards.append(0.0)\n",
    "                continue\n",
    "\n",
    "            ngrams = set() # Use a set to store unique n-grams\n",
    "            total = 0\n",
    "            for ng in zipngram(completion, ngram_size): # Generate n-grams\n",
    "                ngrams.add(ng) # Add n-gram to the set (duplicates are ignored)\n",
    "                total += 1 # Count total n-grams\n",
    "\n",
    "            # Calculate scaling factor: more repetition -> higher scaling\n",
    "            scaling = 1 - len(ngrams) / total\n",
    "            reward = scaling * max_penalty # Apply penalty based on scaling\n",
    "            rewards.append(reward)\n",
    "        return rewards\n",
    "    return repetition_penalty_reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our `get_repetition_penalty_reward(...)` creates a reward function to penalize repetition, with parameters like ngram_size (default 3, for trigrams) and max_penalty (a negative value, e.g., -0.1).\n",
    "\n",
    "A helper function, `zipngram(text, ngram_size)`, generates n-grams by converting text to lowercase, splitting it into words, and using `zip(*[words[i:] for i in range(ngram_size)])` for efficient extraction.\n",
    "\n",
    "Inside, `repetition_penalty_reward(...)` computes the penalty for each completion. If it's empty or too short, it gets a reward of 0.0.\n",
    "\n",
    "The penalty scales as scaling `= 1 - len(ngrams) / total`, where total is the number of n-grams and len(ngrams) is the unique count. More repetition makes scaling approach 1, increasing the penalty.\n",
    "\n",
    "The final reward is scaling * max_penalty, meaning less repetition results in a smaller penalty, while high repetition leads to a stronger negative reward. \n",
    "\n",
    ">We have implemented all five reward functions, Let’s move on to next stage where we define our training args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Configurations for R1 Zero\n",
    "\n",
    "Now we to code a configuration where we can fine-tune how our *reward functions* actually work. So, Let’s define that configuration class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define GRPOScriptArguments for reward function parameters\n",
    "@dataclass\n",
    "class GRPOScriptArguments:\n",
    "    \"\"\"\n",
    "    Script arguments for GRPO training, specifically related to reward functions.\n",
    "    \"\"\"\n",
    "\n",
    "    reward_funcs: list[str] = field(\n",
    "        default_factory=lambda: [\"accuracy\", \"format\"],\n",
    "        metadata={\n",
    "            \"help\": \"List of reward functions. Possible values: 'accuracy', 'format', 'reasoning_steps', 'cosine', 'repetition_penalty'\"\n",
    "        },\n",
    "    )\n",
    "    cosine_min_value_wrong: float = field(\n",
    "        default=-0.5,\n",
    "        metadata={\"help\": \"Minimum reward for cosine scaling for wrong answers\"},\n",
    "    )\n",
    "    cosine_max_value_wrong: float = field(\n",
    "        default=-0.1,\n",
    "        metadata={\"help\": \"Maximum reward for cosine scaling for wrong answers\"},\n",
    "    )\n",
    "    cosine_min_value_correct: float = field(\n",
    "        default=0.8,\n",
    "        metadata={\"help\": \"Minimum reward for cosine scaling for correct answers\"},\n",
    "    )\n",
    "    cosine_max_value_correct: float = field(\n",
    "        default=1.0,\n",
    "        metadata={\"help\": \"Maximum reward for cosine scaling for correct answers\"},\n",
    "    )\n",
    "    cosine_max_len: int = field(\n",
    "        default=1000,\n",
    "        metadata={\"help\": \"Maximum length for cosine scaling\"},\n",
    "    )\n",
    "\n",
    "    repetition_n_grams: int = field(\n",
    "        default=3,\n",
    "        metadata={\"help\": \"Number of n-grams for repetition penalty reward\"},\n",
    "    )\n",
    "    repetition_max_penalty: float = field(\n",
    "        default=-0.1,\n",
    "        metadata={\"help\": \"Maximum (negative) penalty for for repetition penalty reward\"},\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our `@dataclass` decorator makes it easy to create a class for storing data. WhileGRPOScriptArguments class holds reward settings.\n",
    "\n",
    "The reward_funcs list decides which rewards to use, starting with [\"accuracy\", \"format\"], but you can add more like \"reasoning_steps\", \"cosine\", \"repetition_penalty\".\n",
    "\n",
    "Some settings control how the cosine_scaled_reward and repetition_penalty_reward work, letting you adjust how rewards are given.\n",
    "\n",
    "Next up, we have TrainingArguments from the transformers library. This is the **main** configuration object that controls almost **everything** about the training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define TrainingArguments from transformers\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=OUTPUT_DIR,          # Output directory for checkpoints and logs\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=1,             # Total number of training epochs\n",
    "    per_device_train_batch_size=8,  # Batch size per device during training\n",
    "    per_device_eval_batch_size=16,   # Batch size for evaluation\n",
    "    gradient_accumulation_steps=2,  # Accumulate gradients to simulate larger batch size\n",
    "    learning_rate=5e-5,            # Initial learning rate for AdamW optimizer\n",
    "    warmup_ratio=0.1,              # Linear warmup over warmup_ratio fraction of training steps\n",
    "    weight_decay=0.01,             # Apply weight decay to all layers except bias and LayerNorm weights\n",
    "    logging_steps=10,              # Log every X updates steps\n",
    "    evaluation_strategy=\"steps\",    # Evaluate every `eval_steps`\n",
    "    eval_steps=50,                 # Evaluation and logging steps\n",
    "    save_strategy=\"steps\",         # Save checkpoint every `save_steps`\n",
    "    save_steps=50,                 # Save checkpoint every X updates steps\n",
    "    save_total_limit=2,            # Limit the total amount of checkpoints. Deletes the older checkpoints.\n",
    "    dataloader_num_workers=2,      # Number of subprocesses to use for data loading\n",
    "    seed=42,                       # Random seed for reproducibility\n",
    "    bf16=True,                     # Use mixed precision BFP16 training\n",
    "    push_to_hub=False,             # Whether to push the final model to Hugging Face Hub\n",
    "    gradient_checkpointing=True,   # Enable gradient checkpointing\n",
    "    report_to=\"none\",              # Reporting to no one\n",
    "    remove_unused_columns=False,   # Do not remove unused columns from the dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we need to have a ModelConfig. This is where we put settings that are specific to the **model itself**, like which pre-trained model to use, what data type to use (like bfloat16), and whether to trust remote code or not and so.\n",
    "\n",
    "Let’s define our ModelConfig:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelConfig:\n",
    "    \"\"\"\n",
    "    Configuration for the model.\n",
    "    \"\"\"\n",
    "    model_name_or_path: str = field(\n",
    "        default=MODEL_NAME, metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
    "    )\n",
    "    model_revision: Optional[str] = field(\n",
    "        default=\"main\", metadata={\"help\": \"The specific model version to use (can be a branch name, tag name or commit id).\"}\n",
    "    )\n",
    "    torch_dtype: Optional[str] = field(\n",
    "        default=\"bfloat16\", metadata={\"help\": \"Override the default `torch_dtype` and load the model under this dtype.\"}\n",
    "    )\n",
    "    trust_remote_code: bool = field(\n",
    "        default=True, metadata={\"help\": \"Trust remote code when loading model and tokenizer.\"}\n",
    "    )\n",
    "    attn_implementation: Optional[str] = field(\n",
    "        default=\"flash_attention_2\", metadata={\"help\": \"Attention implementation to use. 'flash_attention_2' or None\"}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our **ModelConfig** class holds key settings, including model_name_or_path, which defaults to **Qwen 0.5B Instruct**. We use torch_dtype=\"bfloat16\" for efficiency and set trust_remote_code=True for safe remote loading. Additionally, attn_implementation=\"flash_attention_2\" is enabled for potentially faster training if supported.\n",
    "\n",
    "Now we need to actually **create** instances of these configuration classes so we can use them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate configuration objects\n",
    "script_args = GRPOScriptArguments()\n",
    "model_args = ModelConfig()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to get our list of reward functions and any “callbacks” we want to use during training.\n",
    "\n",
    "Callbacks are like little helpers that can do things at different points in the training process (like logging progress, saving models, etc.). For now, we’ll just use a simple logging callback.\n",
    "\n",
    "Getting our reward functions in one place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function to get reward functions based on script arguments\n",
    "def get_reward_functions(script_args):\n",
    "    \"\"\"\n",
    "    Returns a list of reward functions based on the script arguments.\n",
    "    \"\"\"\n",
    "    reward_funcs_list = []\n",
    "    reward_funcs_registry = {\n",
    "        \"accuracy\": accuracy_reward,  # Assuming accuracy_reward is defined in previous steps\n",
    "        \"format\": format_reward,      # Assuming format_reward is defined in previous steps\n",
    "        \"reasoning_steps\": reasoning_steps_reward, # Assuming reasoning_steps_reward is defined\n",
    "        \"cosine\": get_cosine_scaled_reward( # Assuming get_cosine_scaled_reward is defined\n",
    "            min_value_wrong=script_args.cosine_min_value_wrong,\n",
    "            max_value_wrong=script_args.cosine_max_value_wrong,\n",
    "            min_value_correct=script_args.cosine_min_value_correct,\n",
    "            max_value_correct=script_args.cosine_max_value_correct,\n",
    "            max_len=script_args.cosine_max_len,\n",
    "        ),\n",
    "        \"repetition_penalty\": get_repetition_penalty_reward( # Assuming get_repetition_penalty_reward is defined\n",
    "            ngram_size=script_args.repetition_n_grams,\n",
    "            max_penalty=script_args.repetition_max_penalty,\n",
    "        ),\n",
    "    }\n",
    "\n",
    "    for func_name in script_args.reward_funcs:\n",
    "        if func_name not in reward_funcs_registry:\n",
    "            raise ValueError(f\"Reward function '{func_name}' not found in registry.\")\n",
    "        reward_funcs_list.append(reward_funcs_registry[func_name])\n",
    "\n",
    "    return reward_funcs_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our callback function which will track loss and other important info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "class LoggingCallback(TrainerCallback):\n",
    "    \"\"\"\n",
    "    A simple callback for logging training information at specific steps.\n",
    "    \"\"\"\n",
    "    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):\n",
    "        if state.global_step % args.logging_steps == 0:\n",
    "            if state.log_history and len(state.log_history) > 0:\n",
    "                logger.info(f\"Step {state.global_step}: Loss = {state.log_history[-1].get('loss', None)}, Learning Rate = {state.log_history[-1].get('learning_rate', None)}\")\n",
    "            else:\n",
    "                logger.info(f\"Step {state.global_step}: No logging information available yet\")\n",
    "\n",
    "def get_callbacks(training_args, model_args, script_args):\n",
    "    \"\"\"\n",
    "    Returns a list of callbacks to be used during training.\n",
    "    For now, it includes only the LoggingCallback. You can extend this to add more callbacks.\n",
    "    \"\"\"\n",
    "    callbacks = [LoggingCallback()] # Instantiate our LoggingCallback\n",
    "    return callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, initializing these function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get reward functions and callbacks\n",
    "reward_functions = get_reward_functions(script_args)\n",
    "callbacks = get_callbacks(training_args, model_args, script_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GRPO Training Loop\n",
    "\n",
    "This is the engine that will actually drive our GRPO training. We need to initialize it, giving it all the pieces we’ve prepared: our model, reward functions, training arguments, dataset, and callbacks!\n",
    "\n",
    "Let’s initialize the GRPOTrainer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create GRPOConfig from TrainingArguments\n",
    "grpo_config = GRPOConfig(\n",
    "    **training_args.to_dict(), # Convert TrainingArguments to dictionary and unpack\n",
    "    **{ \n",
    "       # REMOVED model_init_kwargs here \n",
    "       # We are passing the instantiated 'model' object, so GRPOTrainer doesn't need model_init_kwargs\n",
    "    }\n",
    ")\n",
    "\n",
    "grpo_trainer = GRPOTrainer(\n",
    "    model=model,                      # Our initialized Qwen model\n",
    "    reward_funcs=reward_functions,    # List of reward functions from previous step\n",
    "    args=grpo_config,                # GRPOConfig (created from TrainingArguments)\n",
    "    train_dataset=dataset['train'],   # Training dataset\n",
    "    eval_dataset=dataset['test'],    # Evaluation dataset\n",
    "    callbacks=callbacks              # List of callbacks\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now start the **Training Loop**! This is as simple as calling the train() method on our grpo_trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the GRPO Training Loop\n",
    "train_result = grpo_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you run this cell, you should see the training process begin.\n",
    "\n",
    "Training will take some time but we set **num_train_epochs = 1** and are using a small model, it shouldn’t take *too* long for this example.\n",
    "\n",
    "But for real-world GRPO DeepSeek R1 Zero training, you’d likely train for many more epochs and steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving Tiny R1 Zero LLM\n",
    "\n",
    "Once the training completed, we can save our trained model which can be used for inferencing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the path to your trained model (same as OUTPUT_DIR)\n",
    "TRAINED_MODEL_PATH = \"data/Qwen-GRPO-training\"\n",
    "\n",
    "# Save the tokenizer\n",
    "tokenizer.save_pretrained(TRAINED_MODEL_PATH)\n",
    "\n",
    "# Save the trained model\n",
    "grpo_trainer.save_model(TRAINED_MODEL_PATH)\n",
    "\n",
    "print(f\"GRPO Trained model saved to {TRAINED_MODEL_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can simply load the trained model using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the tokenizer - make sure to use trust_remote_code=True if needed\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    TRAINED_MODEL_PATH,\n",
    "    trust_remote_code=True, # If your model config requires it\n",
    "    padding_side=\"right\" # Ensure consistent padding side\n",
    ")\n",
    "\n",
    "# Set pad token if it wasn't saved or loaded correctly\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Load the trained model itself\n",
    "trained_model = AutoModelForCausalLM.from_pretrained(\n",
    "    TRAINED_MODEL_PATH,\n",
    "    trust_remote_code=True, # If your model architecture requires it\n",
    "    torch_dtype=torch.bfloat16 # Keep the same dtype as training for consistency\n",
    ")\n",
    "\n",
    "# Move the loaded model to your device (GPU if available)\n",
    "trained_model.to(device) # 'device' is still our CUDA device from before"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use it for inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing Inference with the Trained Model\n",
    "def test_trained_model_inference(user_input: str):\n",
    "    \"\"\"Test inference with the loaded trained model and tokenizer.\"\"\"\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT}, # Re-use our system prompt\n",
    "        {\"role\": \"user\", \"content\": user_input}\n",
    "    ]\n",
    "\n",
    "    # Apply chat template using our tokenizer\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True\n",
    "    )\n",
    "\n",
    "    # Tokenize the input text\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\").to(device)\n",
    "\n",
    "    # Generate output using our *trained_model*\n",
    "    outputs = trained_model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=200, # Maybe generate a bit longer now\n",
    "        do_sample=True,\n",
    "        temperature=0.7\n",
    "    )\n",
    "\n",
    "    # Decode the generated tokens back to text\n",
    "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    return response\n",
    "\n",
    "# Test the model\n",
    "test_input = \"how are you?\"\n",
    "response = test_trained_model_inference(test_input)\n",
    "print(f\"Test Input: {test_input}\")\n",
    "print(f\"Trained Model Response: {response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two main problems with R1 Zero\n",
    "\n",
    "Now that we have completed our R1 zero training approach using our base model Qwen2–0.5B instead of their DeepSeek V3 (original base model).\n",
    "\n",
    "We cannot identify our trained model problems but researches of DeepSeek saw the R1 Zero model performed really well on reasoning tests, even scoring similarly to more advanced models like **OpenAI-01–0912** on tasks like **AIME 2024**.\n",
    "\n",
    "This showed that using reinforcement learning (RL) to encourage reasoning in language models is a promising approach.\n",
    "\n",
    "But they also noticed DeepSeek-R1-Zero had some key issues that needed fixing for real world use and wider research.\n",
    "\n",
    "![Problem with R1 Zero (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/6378/1*_NdVhpb9cgT3-8o3Qn7mMA.png)\n",
    "\n",
    "Researchers of DeepSeek states that the template is *intentionally simple and structurally focused*. It *avoids* imposing any *content-specific* constraints on the *reasoning process itself*. For example, it doesn’t say:\n",
    "\n",
    "* “You *must* use step-by-step reasoning” (It just says “reasoning process” leaving it open to the model to define what that means).\n",
    "\n",
    "* “You *must* use reflective reasoning”\n",
    "\n",
    "* “You *must* use a specific problem-solving strategy”\n",
    "\n",
    "The main problem was that the reasoning processes inside the `<think>` tags were hard to read, making it tough for humans to follow and analyze.\n",
    "\n",
    "Another issue was language mixing, when asked multi-lingual questions, the model sometimes mixed languages in the same response, leading to inconsistent and confusing outputs.\n",
    "\n",
    "If you asked it questions in, say, Spanish. Suddenly, its “thinking” would be a jumbled mix of **English and Spanish, **not exactly polished! These problems, messy reasoning and language confusion, were the clear roadblocks.\n",
    "> These are the two main reasons they transformed their initial R1 Zero Model into the R1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparing Cold Start Data for SFT\n",
    "\n",
    "So to fix R1 Zero issues and really get DeepSeek reasoning properly, researchers performed a **Cold Start Data Collection and included Supervised Fine Tuning**.\n",
    "\n",
    "You can think of it as giving the model a good foundation in reasoning before the really intense RL training. Basically, they wanted to teach **DeepSeek-V3 Base** what good reasoning looks like and how to present it clearly.\n",
    "\n",
    "One of the example of cold start data is [Bespoke-Stratos-17k](https://huggingface.co/datasets/bespokelabs/Bespoke-Stratos-17k) that we see earlier and will be using for creating R1, but **we need to understand how cold dataset is created so we wont skip any part from the actual training**.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Few-shot Prompting with Long CoT\n",
    "\n",
    "One technique is **Few-shot Prompting with Long Chain-of-Thought (CoT),** in which we try to show DeepSeek-V3 Base (or in our case, Qwen2–0.5B) few examples of questions paired with super detailed, step-by-step solutions. This is Chain-of-Thought (CoT).\n",
    "\n",
    "![Long CoT (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/4068/1*SAhvB0JqaK4d45IiIcj1Ow.png)\n",
    "\n",
    "Goal of this approach is to make the model learn by example and start mimicking this thorough reasoning style.\n",
    "\n",
    "For our example problem “What is 2 + 3 * 4?”, we can create prompts that include a few solved problems as examples. Let’s see how this looks in Python:\n",
    "```python\n",
    "# Loading Model and Tokenizer\n",
    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True, padding_side=\"right\")\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16).to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Generate Long COT Response\n",
    "def generate_response(prompt_text):\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant that provides step-by-step solutions.\"},\n",
    "        {\"role\": \"user\", \"content\": prompt_text}\n",
    "    ]\n",
    "    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)\n",
    "    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False) # Keep it deterministic for example\n",
    "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    return response.split(\"<|im_start|>assistant\\n\")[-1].strip() # Extract assistant's response\n",
    "```\n",
    "Let’s define the few shot examples accordingly for our asked question:\n",
    "```python\n",
    "# Example problems with solutions (using | special_token | as delimiter)\n",
    "few_shot_prompt = \"\"\"\n",
    "Problem: What's the square root of 9 plus 5?\n",
    "Solution: <|special_token|> First, find the square root of 9, which is 3. Then, add 5 to 3.  3 + 5 equals 8. <|special_token|> Summary: The answer is 8.\n",
    "\n",
    "Problem: Train travels at 60 mph for 2 hours, how far?\n",
    "Solution: <|special_token|> Use the formula: Distance = Speed times Time. Speed is 60 mph, Time is 2 hours. Distance = 60 * 2 = 120 miles. <|special_token|> Summary: Train travels 120 miles.\n",
    "\n",
    "Problem: What is 2 + 3 * 4?\n",
    "Solution:\n",
    "\"\"\"\n",
    "```\n",
    "\n",
    "Now using our base model our sample generations looks like this:\n",
    "```python\n",
    "# Generate response for the target problem using few-shot examples\n",
    "target_problem_prompt = few_shot_prompt + \"What is 2 + 3 * 4?\"\n",
    "model_response_few_shot = generate_response(target_problem_prompt)\n",
    "\n",
    "print(\"Few-shot Prompt:\")\n",
    "print(target_problem_prompt)\n",
    "print(\"\\nModel Response (Few-shot CoT):\")\n",
    "print(model_response_few_shot)\n",
    "```\n",
    "\n",
    "It output this structured data\n",
    "\n",
    "```\n",
    "Few-shot Prompt:\n",
    "Problem: What's the square root of 9 plus 5?\n",
    "Solution: <|special_token|> First, find the square root of 9, \n",
    "which is 3. Then, add 5 to 3.  3 + 5 equals 8. \n",
    "<|special_token|> Summary: The answer is 8.\n",
    "\n",
    "Problem: Train travels at 60 mph for 2 hours, how far?\n",
    "Solution: <|special_token|> Use the formula: Distance = Speed times Time. \n",
    "Speed is 60 mph, Time is 2 hours. Distance = 60 * 2 = 120 miles. \n",
    "<|special_token|> Summary: Train travels 120 miles.\n",
    "\n",
    "Problem: What is 2 + 3 * 4?\n",
    "Solution: \n",
    "\n",
    "Model Response (Few-shot CoT):\n",
    "<|special_token|> To solve 2 + 3 * 4, we need to follow the order \n",
    "of operations (PEMDAS/BODMAS). Multiplication should be performed \n",
    "before addition.\n",
    "Step 1: Multiply 3 by 4, which equals 12.\n",
    "Step 2: Add 2 to the result from Step 1: 2 + 12 = 14.\n",
    "<|special_token|> Summary: The answer is 14.\n",
    "```\n",
    "\n",
    "See how the model, after seeing examples, starts to structure its answer with <|special_token|> delimiters and provides step-by-step reasoning leading to the summary and final answer!\n",
    "\n",
    "This is the power of few-shot learning guiding the model towards the desired output format.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct Prompting\n",
    "\n",
    "Another method is **Direct Prompting**. Here, we directly instruct the model to not just solve the problem, but also to explicitly show its reasoning step-by-step and then verify its answer. \n",
    "\n",
    "This is about encouraging a more deliberate and thoughtful problem-solving approach.\n",
    "\n",
    "![Example based learning (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/4656/1*IYyk7UWgDNADFe_djWcXow.png)\n",
    "\n",
    "Let’s craft a prompt for “What is 2 + 3 * 4?” that explicitly asks for reasoning and verification. Here’s the Python code to see it in action:\n",
    "```python\n",
    "# Direct prompting example\n",
    "direct_prompt_text = \"\"\"\n",
    "Problem: Solve this, show reasoning step-by-step, and verify:\n",
    "What is 2 + 3 * 4?\n",
    "\"\"\"\n",
    "\n",
    "model_response_direct = generate_response(direct_prompt_text)\n",
    "\n",
    "print(\"Direct Prompt:\")\n",
    "print(direct_prompt_text)\n",
    "print(\"\\nModel Response (Direct Prompting):\")\n",
    "print(model_response_direct)\n",
    "```\n",
    "The direct prompting output is very easy to understand and this is what it looks like:\n",
    "```\n",
    "Direct Prompt:\n",
    "Problem: Solve this, show reasoning step-by-step, and verify:\n",
    "What is 2 + 3 * 4?\n",
    "\n",
    "Model Response (Direct Prompting):\n",
    "<|special_token|> Reasoning: To solve 2 + 3 * 4, I need to follow \n",
    "the order of operations, which states that multiplication should \n",
    "be done before addition.\n",
    "Step 1: Multiply 3 by 4, which equals 12.\n",
    "Step 2: Add 2 to the result from Step 1: 2 + 12 = 14.\n",
    "Verification: To verify the answer, I can double-check the \n",
    "order of operations and the calculations. Multiplication is \n",
    "indeed performed before addition, and the calculations are correct.\n",
    "<|special_token|> Summary: The answer is 14.\n",
    "```\n",
    "As you can see, by directly asking for reasoning and verification, the model provides a more comprehensive output, including a “Verification” section.\n",
    "\n",
    "This method directly guides the model to produce the kind of detailed reasoning we are looking for."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post Processing Refinement\n",
    "\n",
    "The final technique involves **Post-Processing Refinement**. Interestingly, they even used the outputs from the already trained R1 Zero model for this!\n",
    "\n",
    "Even with its issues, R1 Zero could reason somewhat. So, they took R1 Zero outputs and had human annotators refine them, making them cleaner, more structured, and correcting any mistakes.\n",
    "\n",
    "![Processing Refnement (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/4388/1*-GR29EAnTOVBarQ2JrF5sA.png)\n",
    "\n",
    "Imagine a messy R1 Zero output like this:\n",
    "```\n",
    "<think>  ummm... multiply 3 and 4... get 12... then add 2...</think>\n",
    "<answer> 14 </answer>\n",
    "```\n",
    "\n",
    "Human annotators would then refine it to something much clearer and better formatted:\n",
    "```\n",
    "<|special_token|> Reasoning: To solve this, we use order of operations, doing multiplication before addition.\n",
    "Step 1: Multiply 3 by 4, which is 12.\n",
    "Step 2: Add 2 to the result: 2 + 12 = 14.\n",
    "<|special_token|> Summary: The answer is 14.\n",
    "```\n",
    "\n",
    "While we can’t perfectly simulate human refinement in code, we can demonstrate a basic idea of how you might programmatically reformat and structure a potentially messy output. \n",
    "\n",
    "Let’s take a simulated “messy” output and show how we could refine it:\n",
    "```python\n",
    "# Simulated messy R1 Zero output\n",
    "messy_output = \"<think>  ummm... multiply 3 and 4... get 12... then add 2...</think>\\n<answer> 14 </answer>\"\n",
    "\n",
    "def refine_output(messy_text):\n",
    "    think_content = messy_text.split(\"<think>\")[1].split(\"</think>\")[0].strip()\n",
    "    answer_content = messy_text.split(\"<answer>\")[1].split(\"</answer>\")[0].strip()\n",
    "\n",
    "    refined_text = f\"\"\"<|special_token|> Reasoning: {think_content.replace('umm...', '').strip().capitalize()}.\n",
    "<|special_token|> Summary: The answer is {answer_content}.\"\"\"\n",
    "    return refined_text\n",
    "\n",
    "refined_output_text = refine_output(messy_output)\n",
    "\n",
    "print(\"Messy Output (Simulated R1 Zero):\")\n",
    "print(messy_output)\n",
    "print(\"\\nRefined Output:\")\n",
    "print(refined_output_text)\n",
    "```\n",
    "\n",
    "This will output:\n",
    "```\n",
    "Messy Output (Simulated R1 Zero):\n",
    "<think>  ummm... multiply 3 and 4... get 12... then add 2...</think>\n",
    "<answer> 14 </answer>\n",
    "\n",
    "Refined Output:\n",
    "<|special_token|> Reasoning: Multiply 3 and 4... get 12... then add 2.\n",
    "<|special_token|> Summary: The answer is 14.\n",
    "```\n",
    "\n",
    "This simple refine_output function is just a basic example. Real refinement by humans involves much more nuanced understanding and correction of reasoning steps.\n",
    "\n",
    "However, it shows the core idea: taking initial model outputs and improving their quality and structure to create better training data.\n",
    "> After generating this Cold Start Data, the next crucial step was **Supervised Fine-Tuning (SFT)**, which we’ll explore in the next section!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SFT Stage 1 With Cold Start Data\n",
    "\n",
    "To generate proper cold start data to build R1 using Supervised fine-tuning, we obviously need a proper team along with an excessive amount of code, but thankfully, we already have data ([Bespoke-Stratos-17k](https://huggingface.co/datasets/bespokelabs/Bespoke-Stratos-17k)) that is similar to the cold start form.\n",
    "> We need to know what and how training happens inside the SFT Trainer as it processes our training data?\n",
    "\n",
    "SFT is a form of supervised learning. This means we’re giving the model pairs of inputs and *desired* outputs.\n",
    "\n",
    "In our case, the input might be a problem prompt, and the desired output is the well-reasoned, step-by-step solution from our training dataset. **I hope this point gives a clear view of why there is a need of cold data.**\n",
    "\n",
    "It takes our tokenized training data and feeds it to the model in batches. For each batch, a important set of operations happens, Let’s visualize this internal process:\n",
    "\n",
    "![SFT WorkFlow (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/6838/1*EsEgATw1aSYPjfGtpId2mQ.png)\n",
    "\n",
    "First, the model takes an input, a problem prompt, for instance. It processes this input and generates its best guess for the solution, token by token. These are the *predicted tokens*.\n",
    "\n",
    "Next, the SFT Trainer needs to know how good (or bad) these predictions are. It uses a *loss function*, typically Cross-Entropy Loss. This function mathematically compares the model’s predicted tokens to the *correct* tokens from our training data. Think of it as calculating the “error” of the model’s answer.\n",
    "\n",
    "This “error” isn’t just discarded. It’s the crucial signal for learning. Through a process called *backpropagation*, this error is used to calculate *gradients*. Gradients are like guides, pointing in the direction of parameter adjustments that would reduce the error.\n",
    "\n",
    "Finally, an *optimizer*, like **AdamW** uses these gradients to subtly tweak the model’s internal settings — its parameters. These tweaks are designed to make the model’s next prediction a little bit closer to the correct answer.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1 SFT Trainer Configs for R1\n",
    "\n",
    "Remember those problems we had with R1 Zero messy reasoning and language mixing? SFT is designed to fix exactly that. By training on high-quality, refined data, we’re teaching the model:\n",
    "\n",
    "* **Clear Reasoning Style**: To structure its “thinking” in a way that’s easy to read and follow.\n",
    "\n",
    "* **Consistent Language**: To stick to one language within a response, avoiding confusing mixes.\n",
    "\n",
    "We’re using the Bespoke-Stratos-17k dataset for SFT. As we saw earlier, it’s got 17,000 problems focused on math and code, with a format that looks pretty good for our needs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s quickly remind ourselves of a sample from Bespoke-Stratos-17k:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the \"Bespoke-Stratos-17k\" dataset from bespokelabs\n",
    "bespoke_rl = load_dataset(\"bespokelabs/Bespoke-Stratos-17k\", \"default\")\n",
    "\n",
    "# Access the first sample in the training set\n",
    "bespoke_rl['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dataset, with its system prompts and user-assistant conversations, is perfect for showing our model how conversations with reasoning should look.\n",
    "\n",
    "We’ll use the trl library again, which makes SFT training super easy.\n",
    "\n",
    "First, we need to set up our configurations, similar to what we did for GRPO, but this time for SFT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model and Output Configuration (same as before, or adjust as needed)\n",
    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
    "OUTPUT_DIR = \"data/Qwen-SFT-training\" # New output directory for SFT model\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# Training Arguments - similar to GRPO, but adjust for SFT\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=1,         # Adjust epochs as needed\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=16,\n",
    "    gradient_accumulation_steps=2,\n",
    "    learning_rate=2e-5,        # Adjust learning rate for SFT\n",
    "    warmup_ratio=0.1,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps=10,\n",
    "    evaluation_strategy=\"no\",\n",
    "    eval_steps=50,\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=50,\n",
    "    save_total_limit=2,\n",
    "    dataloader_num_workers=2,\n",
    "    seed=42,\n",
    "    bf16=True,\n",
    "    push_to_hub=False,\n",
    "    gradient_checkpointing=True,\n",
    "    report_to=\"none\",\n",
    ")\n",
    "\n",
    "# Model Configuration - same as before\n",
    "model_args = ModelConfig(\n",
    "    model_name_or_path=MODEL_NAME,\n",
    "    model_revision=\"main\",\n",
    "    torch_dtype=\"bfloat16\",\n",
    "    trust_remote_code=True,\n",
    "    attn_implementation=\"flash_attention_2\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These TrainingArguments and ModelConfig are quite similar to what we used for GRPO, but with a few tweaks that are more suitable for SFT (like a slightly different learning rate, and importantly, packing=True and max_seq_length=4096 for efficient training on longer sequences)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1 STF Training Loop\n",
    "\n",
    "Now, let’s load our dataset and tokenizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Bespoke-Stratos-17k dataset\n",
    "dataset_sft = load_dataset(\"HuggingFaceH4/Bespoke-Stratos-17k\", split='train') # Only using train split for simplicity\n",
    "\n",
    "# Initialize tokenizer - same as before\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    trust_remote_code=True,\n",
    "    padding_side=\"right\"\n",
    ")\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, we initialize the SFTTrainer and start training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize base model for SFT - same as before\n",
    "model_sft = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    trust_remote_code=True,\n",
    "    torch_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# Initialize the SFT Trainer\n",
    "sft_trainer = SFTTrainer(\n",
    "    model=model_sft,                     # Our initialized Qwen model\n",
    "    train_dataset=dataset_sft,           # Bespoke-Stratos-17k dataset\n",
    "    tokenizer=tokenizer,                 # Tokenizer\n",
    "    args=training_args,                  # Training arguments\n",
    ")\n",
    "\n",
    "# Start the SFT Training Loop\n",
    "sft_train_result = sft_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you run this code, you’ll see the SFT training process start. It will look similar to the GRPO training output, showing loss and learning rate at each logging step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like with GRPO, training time will depend on your hardware and chosen epochs. Since we’re still using a small model and only 1 epoch for this example, it should be reasonably quick."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving Tiny R1 LLM\n",
    "\n",
    "After SFT is done, we save our newly fine-tuned model (R1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Saving the Trained SFT Model\n",
    "TRAINED_SFT_MODEL_PATH = \"data/Qwen-SFT-training\" # Same as OUTPUT_DIR\n",
    "\n",
    "# Save the tokenizer\n",
    "tokenizer.save_pretrained(TRAINED_SFT_MODEL_PATH)\n",
    "\n",
    "# Save the trained model\n",
    "sft_trainer.save_model(TRAINED_SFT_MODEL_PATH)\n",
    "\n",
    "print(f\"SFT Trained model saved to {TRAINED_SFT_MODEL_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that’s it for the SFT part! We’ve now taken our base model, shown it lots of examples of good reasoning, and fine-tuned it to be better at producing clear, structured responses.\n",
    "> This finetuned model using SFT is what we called R1 after SFT stage 1\n",
    "\n",
    "The steps after SFT, especially the RL stages and rejection sampling, are complex to implement from scratch in Python. Focusing on the theoretical understanding is key to understand the overall process.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reasoning-Oriented RL\n",
    "\n",
    "After SFT, the model can reason better, but we want to *really* focus on reasoning quality and fix language mixing. This stage uses RL again, but with a smarter reward system.\n",
    "\n",
    "This new reward checks if the model reasoning and answer are in the same language as the question. If you ask in English, the *whole* response should be in English. This fixes language mixing issues.\n",
    "\n",
    "![Reasoning Oriented RL (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/7468/1*Z2oHDdkWb7RnO5uVHPSvMg.png)\n",
    "\n",
    "It adds a **Language Consistency Reward** alongside accuracy to ensure the SFT model reasons and answers in the same language as the input.\n",
    "\n",
    "The GRPO algorithm and training loop from R1 Zero are reused, but the reward signals are improved to specifically target better reasoning and consistent language output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rejection Sampling\n",
    "\n",
    "To get super high-quality reasoning data, DeepSeek uses **Rejection Sampling**. Think of it as a filter to keep only the *best* examples.\n",
    "\n",
    "![Rejection Sampling (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/8520/1*obG-BrhwtIuOv7YBZIpSwg.png)\n",
    "\n",
    "The model generates many reasoning examples. These are then evaluated for correctness and reasoning quality (often using a generative reward model and human checks).\n",
    "\n",
    "Only the *best*, high-quality reasoning examples are kept. Combined with non-reasoning data, this refined dataset is used for a second **SFT Stage 2**, further improving reasoning and general abilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SFT Stage 2 Training\n",
    "\n",
    "The final RL stage focuses on making the model a helpful and safe AI assistant for *all* situations, not just reasoning problems. This is about alignment with human values.\n",
    "\n",
    "**Key Focus: Helpfulness & Harmlessness Rewards**\n",
    "\n",
    "Not just accuracy, the reward system now includes:\n",
    "\n",
    "* **Helpfulness:** Is the response useful and informative?\n",
    "\n",
    "* **Harmlessness:** Is the response safe, unbiased, and ethical?\n",
    "\n",
    "![SFT Stage 2 (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/7086/1*_u5ALx4VYQpsSgT_0s10HQ.png)\n",
    "\n",
    "The training data becomes diverse, including reasoning tasks and human preference data (which output is better — more helpful, less harmful?).\n",
    "\n",
    "The reward system now balances accuracy with **helpfulness and harmlessness**. Iterative RL training (likely GRPO again) optimizes the model to be not just good at reasoning, but also a safe and helpful AI assistant for general use, resulting in DeepSeek R1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distillation\n",
    "\n",
    "To make DeepSeek R1 accessible, they **distilled** its knowledge into smaller models.\n",
    "\n",
    "![Distillation Process (Created by [Fareed Khan](undefined))](https://cdn-images-1.medium.com/max/2500/0*QdOxtvuKaEASreK0.png)\n",
    "\n",
    "Distillation takes the knowledge of a large, powerful “teacher” model (DeepSeek R1) and transfers it to smaller “student” models. Using a large dataset of reasoning examples, the outputs of DeepSeek R1 are used as the *target* answers.\n",
    "\n",
    "Smaller models are then trained (SFT) to mimic these outputs. This results in smaller, faster models that retain a significant portion of DeepSeek R1’s reasoning abilities, making them more practical for wider use."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv-r10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}