Jellywibble's picture
Update README.md
69f4da9
---
license: cc-by-nc-4.0
language:
- en
pipeline_tag: text-classification
tags:
- pytorch
- reward_model
- transformers
- RLHF
library_name: transformers
---
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](https://huggingface.co/datasets/ChaiML/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](https://www.chai-research.com/)
- Model type: Transformer-based Classification Model
- Language: English
- License: cc-by-nc-4.0
- Contact: for general correspondence, please email [hello@chai-research.com](mailto:hello@chai-research.com?subject=Huggingface%20Model%20Inquiry)
## 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](https://huggingface.co/EleutherAI/gpt-j-6B), 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.
```python
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForSequenceClassification.from_pretrained("ChaiML/gpt2_base_retry_and_continue_5m_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 5 million rows out of the [ChaiML/retry_and_continue_50m_reward_model](https://huggingface.co/datasets/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](https://huggingface.co/EleutherAI/gpt-j-6B), with the following sampling parameters:
<figure style="width:30em">
| 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 |
</figure>
### Training procedure
The `gpt2_base_retry_and_continue_5m_reward_model` was trained using a [gpt2](https://huggingface.co/gpt2) 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 156,240 steps. Tensor parallelism and pipeline parallelism were used to distribute the model across GPUs. For evaluation metrics used during training, please see our [Weights and Biases Log](https://wandb.ai/jellywibble/reward).
### BibTeX entry
To cite this model:
```bibtex
@misc{
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](mailto:thomas@chai-research.com?subject=Chai%20Research%20Paper%20Enquiry) or Discord.
### Acknowledgements
This project would not have been possible without the support from members of [Seamless Capital](https://www.seamless-capital.com/)
We thank the following authors from the [Machine Intelligence Laboratory](https://mi.eng.cam.ac.uk/) for their collaboration:
- [Vysas Raina](https://www.linkedin.com/in/vyas-raina-71b226152/)
- [Adian Liusie](https://www.linkedin.com/in/adian-liusie-00b60511a/)