Edit model card

This is part of the Chai reward-model series, using the GPT2 architecture with a classification head, optimising for a user accepting the completion generated by the base model.

Its training dataset consists of purely user-generated content retry_and_continue_50m_reward_model, where a user has the option to decline the generated response via the retry button or end the conversation.

Model details

  • Developed by Chai Research
  • Model type: Transformer-based Classification Model
  • Language: English
  • License: cc-by-nc-4.0
  • Contact: for general correspondence, please email hello@chai-research.com

Uses and limitations

Intended use

This reward model was developed primarily for commercial purposes. It learns an inner representation of response quality rated by humans that can be used to conduct best-of-N sampling and Reinforcement Leanring with the PPO framework.

In addition to scientific uses, you may also further fine-tune and adapt this reward model for deployment, as long as your use is in accordance with the Creative Commons Attribution Non Commercial 4.0 (cc-by-nc-4.0) license, i.e. non-commercial use. This model works with the Transformers Library. If you decide to this pre-trained reward model as a basis for your fine-tuned model, please note that you need to conduct your own risk and bias assessment.

Out-of-scope use

This reward model is not intended for deployment as-is. It is not a product and cannot be used for human-facing interactions without supervision.

This model has not been optimised for common reward-model objectives such as harmfulness, truthfulness and helpfulness, it is only trained based on user actions present on the Chai mobile app platform. Therefore, this model will not rank responses appropriately when evaluating on common open-sourced datasets. All base model responses within the training data were generated using an in-house variant of GPT-J, therefore the model performance may degrade when the input is generated using other language models.

How to use

This reward model can be loaded using the AutoModelForSequenceClassification functionality, with a GPT2 tokenizer where the pad_token_id is set to the EOS token id, padding sides need to be set according to the configurations used during model training.

from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForSequenceClassification.from_pretrained("ChaiML/gpt2_medium_retry_and_continue_12m_reward_model")
tokenizer.pad_token_id = 50256
tokenizer.truncation_side = ‘left’
tokenizer.padding_side = ‘right’
tokens = self.eval_tokenizer(candidates, return_tensors='pt', return_attention_mask=True, padding='longest', truncation=True, max_length=256)
reward = model(**tokens).logits

Model training

Training dataset

This model was trained by randomly sampling 12 million rows out of the ChaiML/retry_and_continue_50m_reward_model dataset. The original dataset contains over 50 million rows of completions (chatbot responses), along with number of remaining user messages within their corresponding conversations and whether the user pressed the "retry" button (where the completion is rejected and resampled). The model which was used to generate these completions is a in-house variant of GPT-J, with the following sampling parameters:

Parameters Value
temperature 0.72
repetition_penalty 1.13125
max_new_tokens 64
top_p 0.725
top_k 0
eos_token_id 198
do_sample True

Training procedure

The gpt2_medium_retry_and_continue_12m_reward_model was trained using a gpt2-medium base model and a classification head with single output. Binary Cross Entropy loss was used. The model was trained on 4xA40 GPUs, 16 per device batch size and gradient accumulation of 1 (therefore the effective batch size is 64), with 1e-5 learning rate for 2 epochs for a total of 375,000 steps. Tensor parallelism and pipeline parallelism were used to distribute the model across GPUs.

BibTeX entry

To cite this model:

  author = {Chai Research, Irvine, Boubert, Raina, Liusie, Mudupalli, Korshuk, Liu, Cremer, Assassi, C. Beauchamp, Lu, Rialan, W. Beauchamp},
  title = {{Rewarding chatbots for real-world engagement with millions of users}},
  howpublished = {\url{https://arxiv.org/abs/2303.06135}},
  year = 2023,
  month = Mar

If you use this model, we would love to hear about it! Reach out on correspondence email or Discord.


This project would not have been possible without the support from members of Seamless Capital

We thank the following authors from the Machine Intelligence Laboratory for their collaboration:

Downloads last month
This model does not have enough activity to be deployed to Inference API (serverless) yet. Increase its social visibility and check back later, or deploy to Inference Endpoints (dedicated) instead.