# Rubric-based Grading for LLMs

- [`data.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/recipes/rubric/data.py) contains the definition for the datapoint class. Each datapoint consists of a conversation prefix and a list of rubric items.
- [`generate_data.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/recipes/rubric/generate_data.py) generates some example datapoints if you want to run our demo on addition.
- [`env.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/recipes/rubric/env.py) determines what each rollout will do. It will let the policy read the prefix, generate a response, ask a grader LLM to grade based on a list of rubric items, and finally provide a reward by summing the response of each grader.
- [`train.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/recipes/rubric/train.py) allows you to train LLMs on any dataset saved in our format (specified in `data.py`). The default script will train on the addition task, whose data is generated by `generate_data.py`.
- [`prometheus_experimental.py`](https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/recipes/rubric/prometheus_experimental.py) contains a script to train the LLMs based on the rubrics from the [`prometheus-eval/Feedback-Collection`](https://huggingface.co/datasets/prometheus-eval/Feedback-Collection/viewer/default/train?row=0&views%5B%5D=train) dataset. It is experimental though -- even though the reward goes up, there is no guarantee that the model is actually better. We hope our script serves as a starting point, and more research is needed.

## A simple example of using a grader LLM with rubrics

We show how to use a rubric-based LLM to provide a reward for an addition task. E.g.

```text
**User**: What's 233 + 100?
**Assistant**: 333
```

Usually, this could be graded by matching the number to the ground truth 333 without needing an LLM. However, for pedagogical purposes, we will grade the response using a language model with a rubric. That is, we will ask a language model "Does the assistant answer 333?"

### Generate an example dataset

To run this, first generate a dataset:

```text
python -m tinker_cookbook.recipes.rubric.generate_data
```

Then you will see two `jsonl` files generated, one for training, one for testing. For example, if you look into `tinker_cookbook/example_data/example_rubric_train.jsonl`, each datapoint consists of

- a convo (the conversation prefix that the policy sees)
- rubric_items: a list of rubric items that specify what is a good response, how the grader should format the response, and how the grading result should be extracted.

```text
{
  "convo": [
    {
      "role": "user",
      "content": "What is 4 + 5?"
    },
    {
      "role": "assistant",
      "content": "9"
    },
    {
      "role": "user",
      "content": "What is 122 + 12?"
    }
  ],
  "rubric_items": [
    {
      "rubric_str": "Does the chatbot correctly get the answer 134?",
      "extraction_regex": "<score>(.*)</score>",
      "grader_output_format_instruction": "Please output your score between 0 and 1 wrapped in <score> ... </score>"
    }
  ]
}
```

### Debugging and Printing What Happens During Rollouts

Run

```text
python -m tinker_cookbook.recipes.rubric.debug_env
```

You can see the message that the policy sees, its response, the grader input, and the grader output.

![Debug output showing the conversation context, policy response, grader prompt, and extracted score](https://github.com/user-attachments/assets/9f4e3c89-f21e-49b0-96d6-e2f27bd21b43)

### An example training run

To train the LLM to add with a rubric-based LLM, run

```text
python -m tinker_cookbook.recipes.rubric.train
```

You can see the reward quickly goes up. In this example, `test/env/all/reward/total` improves from 0.354 at step 0 to 0.994 by step 60, while the final training batch reaches `env/all/rubric_score=1.0`.

![Test reward increasing over training steps for the addition task](https://raw.githubusercontent.com/thinking-machines-lab/tinker-cookbook/main/tinker_cookbook/recipes/rubric/assets/test-reward.png)

### A more realistic dataset

We take the `prometheus-eval/Feedback-Collection` dataset from [Hugging Face](https://huggingface.co/datasets/prometheus-eval/Feedback-Collection/), which contains rubrics to grade general chat responses. Run the following to kick off training:

```text
python -m tinker_cookbook.recipes.rubric.prometheus_experimental
```

We can see that the reward climbs up steadily.

![Prometheus reward climbing steadily over training steps](https://raw.githubusercontent.com/thinking-machines-lab/tinker-cookbook/main/tinker_cookbook/recipes/rubric/assets/prometheus-reward.png)

Note that this training recipe is experimental -- to make the performance better we may need to fine-tune the grader LLM as well. We hope our code serves as a starting point for you to improve rubric-based grading for training LLMs!