https://github.com/huggingface/trl/blob/main/examples/notebooks/gpt2-sentiment.ipynb