self-rewarding-lm-pytorch

The SelfRewardingTrainer implements the training framework for Self-Rewarding Language Models. It requires a transformer model, datasets for SFT and self-rewarding prompts, and tokenizer encoding/decoding functions. Checkpoints are automatically saved to ./checkpoints after each finetuning stage.

Key arguments:

• transformer: The model to train.
• finetune_configs: A dictionary or list of configuration objects defining the training stages.
• tokenizer_decode: A function to convert tokens to strings.
• tokenizer_encode: A function to convert strings to tensors.
• accelerate_kwargs: Arguments passed to the accelerate backend.

import torch
from torch import Tensor
from self_rewarding_lm_pytorch import (
    SelfRewardingTrainer,
    create_mock_dataset
)
from x_transformers import TransformerWrapper, Decoder

transformer = TransformerWrapper(
    num_tokens = 256,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 1,
        heads = 8
    )
)

sft_dataset = create_mock_dataset(100, lambda: (torch.randint(0, 256, (256,)), torch.tensor(1)))
prompt_dataset = create_mock_dataset(100, lambda: 'mock prompt')

def decode_tokens(tokens: Tensor) -> str:
    decode_token = lambda token: str(chr(max(32, token)))
    return ''.join(list(map(decode_token, tokens)))

def encode_str(seq_str: str) -> Tensor:
    return Tensor(list(map(ord, seq_str)))

trainer = SelfRewardingTrainer(
    transformer,
    finetune_configs = dict(
        train_sft_dataset = sft_dataset,
        self_reward_prompt_dataset = prompt_dataset,
        dpo_num_train_steps = 1000
    ),
    tokenizer_decode = decode_tokens,
    tokenizer_encode = encode_str,
    accelerate_kwargs = dict(
        cpu = True
    )
)

trainer(overwrite_checkpoints = True)

The SPINTrainer implements Self-Play Fine-Tuning (SPIN). It can be used as a standalone trainer or integrated into a larger fine-tuning pipeline.

Key arguments:

• transformer: The model to train.
• max_seq_len: Maximum sequence length.
• train_sft_dataset: The SFT dataset.
• checkpoint_every: Frequency of saving checkpoints.
• spin_kwargs: Dictionary of arguments for the SPIN process (e.g., λ).

import torch
from self_rewarding_lm_pytorch import (
    SPINTrainer,
    create_mock_dataset
)
from x_transformers import TransformerWrapper, Decoder

transformer = TransformerWrapper(
    num_tokens = 256,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

sft_dataset = create_mock_dataset(100, lambda: (torch.randint(0, 256, (256,)), torch.tensor(1)))

spin_trainer = SPINTrainer(
    transformer,
    max_seq_len = 16,
    train_sft_dataset = sft_dataset,
    checkpoint_every = 100,
    spin_kwargs = dict(
        λ = 0.1,
    ),
)

spin_trainer()

To use a custom reward prompt instead of the default LLM-as-Judge, pass a RewardConfig instance to the self_reward_prompt_config argument of SelfRewardingTrainer.

RewardConfig requires:

• prompt_template: A string template for the reward prompt. Use {{ prompt }} and {{ response }} as placeholders for the input data.
• reward_regex_template: A regex template used to parse the reward from the LLM's response. Use {{ reward }} as the placeholder for the extracted numeric value.

from self_rewarding_lm_pytorch import RewardConfig

trainer = SelfRewardingTrainer(
    transformer,
    ...,
    self_reward_prompt_config = RewardConfig(
        prompt_template = """
        Pretty please rate the following user prompt and response
        User: {{ prompt }}
        Response: {{ response }}

        Format your score as follows:
        Rating: <rating as integer from 0 - 10>
        """,
        reward_regex_template = """
        Rating: {{ reward }}
        """
    )
)