Rubric-based Grading for LLMs

• 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 generates some example datapoints if you want to run our demo on addition.
• 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 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 contains a script to train the LLMs based on the rubrics from the prometheus-eval/Feedback-Collection 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.

**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:

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.

{
  "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

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.

An example training run

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

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.

A more realistic dataset

We take the prometheus-eval/Feedback-Collection dataset from Hugging Face, which contains rubrics to grade general chat responses. Run the following to kick off training:

python -m tinker_cookbook.recipes.rubric.prometheus_experimental

We can see that the reward climbs up steadily.

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!