GRPO (LLM)

.
└── dataset_dir/
    ├── train.jsonl
    └── reward_models.py

train_dataset = [
  {
    "prompt": "...",
    "answer": "..."
  },
]

{"prompt": "...","answer": "..."}
{"prompt": "...","answer": "..."}
{"prompt": "...","answer": "..."}
{"prompt": "...","answer": "..."}

[
  {
    "role": "user",
    "content":"When a spring does work on an object, we cannot find the work by simply multiplying the spring force by the object's displacement. … Provide your reasoning between <think> and </think> and then your final answer between <answer> and (put a float here) </answer>."  
 }
]

1.2

train_dataset = [
  {
    "prompt": [
      {
        "role": "user",
        "content": "When a spring does work on an object, … provide your reasoning between <think> and </think> and then your final answer between <answer> and (put a float here) </answer>"
      }],

    "answer": "1.2",
  },
]

from your_tags import reasoning_start, reasoning_end, solution_start, solution_end
import re
from typing import Callable, List

def formatting_reward_func(completions, **kwargs) -> List[float]:
    """
    Rewards the presence of both reasoning and answer tags in the model's output.
    """
    thinking_pattern = f'{reasoning_start}(.*?){reasoning_end}'
    answer_pattern   = f'{solution_start}(.*?){solution_end}'
    scores = []

    for completion in completions:
        text    = completion[0]['content']
        score   = 0.0
        if len(re.findall(thinking_pattern, text, re.DOTALL)) == 1:
            score += 1.0
        if len(re.findall(answer_pattern, text, re.DOTALL)) == 1:
            score += 1.0
        scores.append(score)

    return scores


def correctness_reward_func(prompts, completions, answer, **kwargs) -> List[float]:
    """
    Rewards exact match of the numeric answer within the <answer> tags.
    """
    answer_pattern = f'{solution_start}(.*?){solution_end}'
    responses = [
        re.findall(answer_pattern, comp[0]['content'], re.DOTALL)
        for comp in completions
    ]

    # Example debug print
    q     = prompts[0][-1]['content']
    a     = answer[0]
    resp0 = completions[0][0]['content']
    print("-"*20, f"Q: {q}\nA: {a}\nR: {resp0}")

    return [
        2.0 if len(r)==1 and r[0].strip()==str(a) else 0.0
        for r in responses
    ]

reward_functions: List[Callable] = [formatting_reward_func, correctness_reward_func]