# LongBench v2 / 66ed5be2821e116aacb1fb57

task_id: 588a6997-9931-5318-9176-8e639f0ecb9f
task_key: train--66ed5be2821e116aacb1fb57
task_revision_id: 3

{"choice_A":"In the train method in ppo_trainer.py, manually save the model's gradients and reload them at the start of each epoch.","choice_B":"Modify the compute_policy_loss method in ppo_trainer.py to move the model's backward operation from the end of each epoch to the beginning of training for unified processing.","choice_C":"Modify the split_between_epochs method in accelerator.py to ensure that the optimizer state is not reset when splitting data, and maintain gradient accumulation across epochs.","choice_D":"Modify the training_step method in ppo_trainer.py to call accelerator.backward() at the end of each epoch, so that accumulated gradients are preserved into the next epoch instead of being reset within each epoch.","context":"<div style=\"text-align: center\">\n<img src=\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl_banner_dark.png\">\n</div>\n\n# TRL - Transformer Reinforcement Learning\n> Full stack library to fine-tune and align large language models.\n\n<p align=\"center\">\n    <a href=\"https://github.com/huggingface/trl/blob/main/LICENSE\">\n        <img alt=\"License\" src=\"https://img.shields.io/github/license/huggingface/trl.svg?color=blue\">\n    </a>\n    <a href=\"https://huggingface.co/docs/trl/index\">\n        <img alt=\"Documentation\" src=\"https://img.shields.io/website/http/huggingface.co/docs/trl/index.svg?down_color=red&down_message=offline&up_message=online\">\n    </a>\n    <a href=\"https://github.com/huggingface/trl/releases\">\n        <img alt=\"GitHub release\" src=\"https://img.shields.io/github/release/huggingface/trl.svg\">\n    </a>\n</p>\n\n\n## What is it?\n\nThe `trl` library is a full stack tool to fine-tune and align transformer language and diffusion models using methods such as Supervised Fine-tuning step (SFT), Reward Modeling (RM) and the Proximal Policy Optimization (PPO) as well as Direct Preference Optimization (DPO). \n\nThe library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library and thus allows to use any model architecture available there.\n\n\n## Highlights\n\n- **`Efficient and scalable`**: \n    - [`accelerate`](https://github.com/huggingface/accelerate) is the backbone of `trl` which allows to scale model training from a single GPU to a large scale multi-node cluster with methods such as DDP and DeepSpeed.\n    - [`PEFT`](https://github.com/huggingface/peft) is fully integrated and allows to train even the largest models on modest hardware with quantisation and methods such as LoRA or QLoRA.\n    - [`unsloth`](https://github.com/unslothai/unsloth) is also integrated and allows to significantly speed up training with dedicated kernels.\n- **`CLI`**: With the [CLI](https://huggingface.co/docs/trl/clis) you can fine-tune and chat with LLMs without writing any code using a single command and a flexible config system.\n- **`Trainers`**: The Trainer classes are an abstraction to apply many fine-tuning methods with ease such as the [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer), [`DPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.DPOTrainer), [`RewardTrainer`](https://huggingface.co/docs/trl/reward_trainer), [`PPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.PPOTrainer), [`CPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.CPOTrainer), and [`ORPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.ORPOTrainer).\n- **`AutoModels`**: The [`AutoModelForCausalLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForCausalLMWithValueHead) & [`AutoModelForSeq2SeqLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForSeq2SeqLMWithValueHead) classes add an additional value head to the model which allows to train them with RL algorithms such as PPO.\n- **`Examples`**: Train GPT2 to generate positive movie reviews with a BERT sentiment classifier, full RLHF using adapters only, train GPT-j to be less toxic, [StackLlama example](https://huggingface.co/blog/stackllama), etc. following the [examples](https://github.com/huggingface/trl/tree/main/examples).\n\n## Installation\n\n### Python package\nInstall the library with `pip`:\n```bash\npip install trl\n```\n\n### From source\nIf you want to use the latest features before an official release you can install from source:\n```bash\npip install git+https://github.com/huggingface/trl.git\n```\n\n### Repository\nIf you want to use the examples you can clone the repository with the following command:\n```bash\ngit clone https://github.com/huggingface/trl.git\n```\n\n## Command Line Interface (CLI)\n\nYou can use TRL Command Line Interface (CLI) to quickly get started with Supervised Fine-tuning (SFT), Direct Preference Optimization (DPO) and test your aligned model with the chat CLI: \n\n**SFT:**\n\n```bash\ntrl sft --model_name_or_path facebook/opt-125m --dataset_name stanfordnlp/imdb --output_dir opt-sft-imdb\n```\n\n**DPO:**\n\n```bash\ntrl dpo --model_name_or_path facebook/opt-125m --dataset_name trl-internal-testing/hh-rlhf-helpful-base-trl-style --output_dir opt-sft-hh-rlhf \n```\n\n**Chat:**\n\n```bash\ntrl chat --model_name_or_path Qwen/Qwen1.5-0.5B-Chat\n```\n\nRead more about CLI in the [relevant documentation section](https://huggingface.co/docs/trl/main/en/clis) or use `--help` for more details.\n\n## How to use\n\nFor more flexibility and control over the training, you can use the dedicated trainer classes to fine-tune the model in Python.\n\n### `SFTTrainer`\n\nThis is a basic example of how to use the `SFTTrainer` from the library. The `SFTTrainer` is a light wrapper around the `transformers` Trainer to easily fine-tune language models or adapters on a custom dataset.\n\n```python\n# imports\nfrom datasets import load_dataset\nfrom trl import SFTTrainer\n\n# get dataset\ndataset = load_dataset(\"stanfordnlp/imdb\", split=\"train\")\n\n# get trainer\ntrainer = SFTTrainer(\n    \"facebook/opt-350m\",\n    train_dataset=dataset,\n    dataset_text_field=\"text\",\n    max_seq_length=512,\n)\n\n# train\ntrainer.train()\n```\n\n### `RewardTrainer`\n\nThis is a basic example of how to use the `RewardTrainer` from the library. The `RewardTrainer` is a wrapper around the `transformers` Trainer to easily fine-tune reward models or adapters on a custom preference dataset.\n\n```python\n# imports\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer\nfrom trl import RewardTrainer\n\n# load model and dataset - dataset needs to be in a specific format\nmodel = AutoModelForSequenceClassification.from_pretrained(\"gpt2\", num_labels=1)\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n...\n\n# load trainer\ntrainer = RewardTrainer(\n    model=model,\n    tokenizer=tokenizer,\n    train_dataset=dataset,\n)\n\n# train\ntrainer.train()\n```\n\n### `PPOTrainer`\n\nThis is a basic example of how to use the `PPOTrainer` from the library. Based on a query the language model creates a response which is then evaluated. The evaluation could be a human in the loop or another model's output.\n\n```python\n# imports\nimport torch\nfrom transformers import AutoTokenizer\nfrom trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model\nfrom trl.core import respond_to_batch\n\n# get models\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')\nref_model = create_reference_model(model)\n\ntokenizer = AutoTokenizer.from_pretrained('gpt2')\ntokenizer.pad_token = tokenizer.eos_token\n\n# initialize trainer\nppo_config = PPOConfig(batch_size=1, mini_batch_size=1)\n\n# encode a query\nquery_txt = \"This morning I went to the \"\nquery_tensor = tokenizer.encode(query_txt, return_tensors=\"pt\")\n\n# get model response\nresponse_tensor  = respond_to_batch(model, query_tensor)\n\n# create a ppo trainer\nppo_trainer = PPOTrainer(ppo_config, model, ref_model, tokenizer)\n\n# define a reward for response\n# (this could be any reward such as human feedback or output from another model)\nreward = [torch.tensor(1.0)]\n\n# train model for one step with ppo\ntrain_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)\n```\n\n### `DPOTrainer`\n\n`DPOTrainer` is a trainer that uses [Direct Preference Optimization algorithm](https://huggingface.co/papers/2305.18290). This is a basic example of how to use the `DPOTrainer` from the library. The `DPOTrainer` is a wrapper around the `transformers` Trainer to easily fine-tune reward models or adapters on a custom preference dataset.\n\n```python\n# imports\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nfrom trl import DPOTrainer\n\n# load model and dataset - dataset needs to be in a specific format\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2\")\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n...\n\n# load trainer\ntrainer = DPOTrainer(\n    model=model,\n    tokenizer=tokenizer,\n    train_dataset=dataset,\n)\n\n# train\ntrainer.train()\n```\n\n## Development\n\nIf you want to contribute to `trl` or customizing it to your needs make sure to read the [contribution guide](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md) and make sure you make a dev install:\n\n```bash\ngit clone https://github.com/huggingface/trl.git\ncd trl/\nmake dev\n```\n\n## References\n\n### Proximal Policy Optimisation\nThe PPO implementation largely follows the structure introduced in the paper **\"Fine-Tuning Language Models from Human Preferences\"** by D. Ziegler et al. \\[[paper](https://huggingface.co/papers/1909.08593), [code](https://github.com/openai/lm-human-preferences)].\n\n### Direct Preference Optimization\nDPO is based on the original implementation of **\"Direct Preference Optimization: Your Language Model is Secretly a Reward Model\"** by E. Mitchell et al. \\[[paper](https://huggingface.co/papers/2305.18290), [code](https://github.com/eric-mitchell/direct-preference-optimization)]\n\n\n## Citation\n\n```bibtex\n@misc{vonwerra2022trl,\n  author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang},\n  title = {TRL: Transformer Reinforcement Learning},\n  year = {2020},\n  publisher = {GitHub},\n  journal = {GitHub repository},\n  howpublished = {\\url{https://github.com/huggingface/trl}}\n}\n```\n\n\n# How to contribute to TRL?\n\nEveryone is welcome to contribute, and we value everybody's contribution. Code\ncontributions are not the only way to help the community. Answering questions, helping\nothers, and improving the documentation are also immensely valuable.\n\nIt also helps us if you spread the word! Reference the library in blog posts\nabout the awesome projects it made possible, shout out on Twitter every time it has\nhelped you, or simply ⭐️ the repository to say thank you.\n\nHowever you choose to contribute, please be mindful and respect our\n[code of conduct](https://github.com/huggingface/trl/blob/main/CODE_OF_CONDUCT.md).\n\n**This guide was heavily inspired by the awesome [scikit-learn guide to contributing](https://github.com/scikit-learn/scikit-learn/blob/main/CONTRIBUTING.md).**\n\n## Ways to contribute\n\nThere are several ways you can contribute to TRL:\n\n* Fix outstanding issues with the existing code.\n* Submit issues related to bugs or desired new features.\n* Implement trainers for new post-training algorithms.\n* Contribute to the examples or to the documentation.\n\nIf you don't know where to start, there is a special [Good First\nIssue](https://github.com/huggingface/trl/contribute) listing. It will give you a list of\nopen issues that are beginner-friendly and help you start contributing to open-source. The best way to do that is to open a Pull Request and link it to the issue that you'd like to work on. We try to give priority to opened PRs as we can easily track the progress of the fix, and if the contributor does not have time anymore, someone else can take the PR over.\n\nFor something slightly more challenging, you can also take a look at the [Good Second Issue](https://github.com/huggingface/trl/labels/Good%20Second%20Issue) list. In general though, if you feel like you know what you're doing, go for it and we'll help you get there! 🚀\n\n> All contributions are equally valuable to the community. 🥰\n\nBefore you start contributing make sure you have installed all the dev tools:\n\n```bash\nmake dev\n```\n\n## Fixing outstanding issues\n\nIf you notice an issue with the existing code and have a fix in mind, feel free to [start contributing](#create-a-pull-request) and open a Pull Request!\n\n## Submitting a bug-related issue or feature request\n\nDo your best to follow these guidelines when submitting a bug-related issue or a feature request. It will make it easier for us to come back to you quickly and with good feedback.\n\n### Did you find a bug?\n\nThe TRL library is robust and reliable thanks to users who report the problems they encounter.\n\nBefore you report an issue, we would really appreciate it if you could **make sure the bug was not\nalready reported** (use the search bar on GitHub under Issues). Your issue should also be related to bugs in the library itself, and not your code.\n\nOnce you've confirmed the bug hasn't already been reported, please include the following information in your issue so we can quickly resolve it:\n\n* Your **OS type and version**, **Python**, **PyTorch**, **TRL** and **Transformers** versions.\n* A short, self-contained, code snippet that allows us to reproduce the bug in\n  less than 30s.\n* The *full* traceback if an exception is raised.\n* Attach any other additional information, like screenshots, you think may help.\n\nTo get the OS and software versions automatically, run the following command:\n\n```bash\ntransformers-cli env\n```\n\n### Do you want a new feature?\n\nIf there is a new feature you'd like to see in TRL, please open an issue and describe:\n\n1. What is the *motivation* behind this feature? Is it related to a problem or frustration with the library? Is it a feature related to something you need for a project? Is it something you worked on and think it could benefit the community?\n\n   Whatever it is, we'd love to hear about it!\n\n2. Describe your requested feature in as much detail as possible. The more you can tell us about it, the better we'll be able to help you.\n3. Provide a *code snippet* that demonstrates the features usage.\n4. If the feature is related to a paper, please include a link.\n\nIf your issue is well written we're already 80% of the way there by the time you create it.\n\n## Do you want to implement a new trainer?\n\nNew post-training methods are published on a frequent basis and those which satisfy the following criteria are good candidates to be integrated in TRL:\n\n* **Simplicity:** does the new method achieve similar performance as prior methods, but with less complexity? A good example is [Direct Preference Optimization](https://arxiv.org/abs/2305.18290) (DPO), which provided a simpler and compelling alternative to RLHF methods.\n* **Efficiency:** does the new method provide a significant improvement in training efficiency? A good example is [Odds Ratio Preference Optimization](https://arxiv.org/abs/2403.07691v2), which utilises a similar objective as DPO, but requires half the GPU VRAM.\n\nMethods which only provide incremental improvements at the expense of added complexity or compute costs are unlikely to be included in TRL.\n\nIf you want to implement a trainer for a new post-training method, first open an issue and provide the following information:\n\n* A short description of the method and a link to the paper.\n* Link to the implementation if it is open-sourced.\n* Link to model weights trained with the method if they are available.\n\nBased on the community and maintainer feedback, the next step will be to implement the trainer and config classes. See the following examples for inspiration:\n\n* Paired preference optimisation: [`dpo_trainer.py`](./trl/trainer/dpo_trainer.py) and [`dpo_config.py`](./trl/trainer/dpo_config.py)\n* RL-based optimisation: [`rloo_trainer.py](./trl/trainer/rloo_trainer.py) and [`rloo_config.py](./trl/trainer/rloo_config.py)\n* Online optimisation: [`online_dpo_trainer.py`](./trl/trainer/online_dpo_trainer.py) and [`online_dpo_config.py`](./trl/trainer/online_dpo_config.py)\n\n## Do you want to add documentation?\n\nWe're always looking for improvements to the documentation that make it more clear and accurate. Please let us know how the documentation can be improved, such as typos, dead links and any missing, unclear or inaccurate content.. We'll be happy to make the changes or help you make a contribution if you're interested!\n\n## Submitting a pull request (PR)\n\nBefore writing code, we strongly advise you to search through the existing PRs or\nissues to make sure that nobody is already working on the same thing. If you are\nunsure, it is always a good idea to open an issue to get some feedback.\n\nYou will need basic `git` proficiency to be able to contribute to\nTRL. `git` is not the easiest tool to use but it has the greatest\nmanual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro\nGit](https://git-scm.com/book/en/v2) is a very good reference.\n\nFollow these steps to start contributing:\n\n1. Fork the [repository](https://github.com/huggingface/trl) by\n   clicking on the 'Fork' button on the repository's page. This creates a copy of the code\n   under your GitHub user account.\n\n2. Clone your fork to your local disk, and add the base repository as a remote. The following command\n   assumes you have your public SSH key uploaded to GitHub. See the following guide for more\n   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).\n\n   ```bash\n   $ git clone git@github.com:<your Github handle>/trl.git\n   $ cd trl\n   $ git remote add upstream https://github.com/huggingface/trl.git\n   ```\n\n3. Create a new branch to hold your development changes, and do this for every new PR you work on.\n\n   Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):\n\n   ```bash\n   $ git checkout main\n   $ git fetch upstream\n   $ git merge upstream/main\n   ```\n\n   Once your `main` branch is synchronized, create a new branch from it:\n\n   ```bash\n   $ git checkout -b a-descriptive-name-for-my-changes\n   ```\n\n   **Do not** work on the `main` branch.\n\n4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:\n\n   ```bash\n   $ make dev\n   ```\n\n   (If TRL was already installed in the virtual environment, remove\n   it with `pip uninstall trl` before reinstalling it.)\n\n   Alternatively, if you are using [Visual Studio Code](https://code.visualstudio.com/Download), the fastest way to get set up is by using\n   the provided Dev Container. Documentation on how to get started with dev containers is available [here](https://code.visualstudio.com/docs/remote/containers).\n\n5. Develop the features on your branch.\n\n   As you work on the features, you should make sure that the test suite\n   passes. You should run the tests impacted by your changes like this (see \n   below an explanation regarding the environment variable):\n\n   ```bash\n   $ pytest tests/<TEST_TO_RUN>.py\n   ```\n   \n   > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on\n   > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).\n\n   You can also run the full suite with the following command.\n\n   ```bash\n   $ make test\n   ```\n\n   TRL relies on `ruff` to format its source code\n   consistently. After you make changes, apply automatic style corrections and code verifications\n   that can't be automated in one go with:\n\n   This target is also optimized to only work with files modified by the PR you're working on.\n\n   If you prefer to run the checks one after the other, the following command apply the\n   style corrections:\n\n   ```bash\n   $ make precommit\n   ```\n\n   Once you're happy with your changes, add changed files using `git add` and\n   make a commit with `git commit` to record your changes locally:\n\n   ```bash\n   $ git add modified_file.py\n   $ git commit\n   ```\n\n   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).\n\n   It is a good idea to sync your copy of the code with the original\n   repository regularly. This way you can quickly account for changes:\n\n   ```bash\n   $ git fetch upstream\n   $ git rebase upstream/main\n   ```\n\n   Push the changes to your account using:\n\n   ```bash\n   $ git push -u origin a-descriptive-name-for-my-changes\n   ```\n\n6. Once you are satisfied (**and the checklist below is happy too**), go to the\n   webpage of your fork on GitHub. Click on 'Pull request' to send your changes\n   to the project maintainers for review.\n\n7. It's ok if maintainers ask you for changes. It happens to core contributors\n   too! So everyone can see the changes in the Pull request, work in your local\n   branch and push the changes to your fork. They will automatically appear in\n   the pull request.\n\n\n### Checklist\n\n1. The title of your pull request should be a summary of its contribution;\n2. If your pull request addresses an issue, please mention the issue number in\n   the pull request description to make sure they are linked (and people\n   consulting the issue know you are working on it);\n3. To indicate a work in progress please prefix the title with `[WIP]`, or mark\n   the PR as a draft PR. These are useful to avoid duplicated work, and to differentiate\n   it from PRs ready to be merged;\n4. Make sure existing tests pass;\n5. Add high-coverage tests. No quality testing = no merge.\n\n\n### Tests\n\nAn extensive test suite is included to test the library behavior and several examples. Library tests can be found in\nthe [tests folder](https://github.com/huggingface/trl/tree/main/tests).\n\nWe use `pytest` in order to run the tests. From the root of the\nrepository, here's how to run tests with `pytest` for the library:\n\n```bash\n$ python -m pytest -sv ./tests\n```\n\nIn fact, that's how `make test` is implemented (sans the `pip install` line)!\n\nYou can specify a smaller set of tests in order to test only the feature\nyou're working on.\n\n\ndatasets>=1.17.0\ntorch>=1.4.0\ntqdm\ntransformers>=4.40.0\naccelerate\npeft>=0.3.0\ntyro>=0.5.7\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"trl is an open library for RL with transformer models.\n\nNote:\n\n   VERSION needs to be formatted following the MAJOR.MINOR.PATCH convention\n   (we need to follow this convention to be able to retrieve versioned scripts)\n\nSimple check list for release from AllenNLP repo: https://github.com/allenai/allennlp/blob/master/setup.py\n\nTo create the package for pypi.\n\n0. Prerequisites:\n   - Dependencies:\n     - twine: \"pip install twine\"\n   - Create an account in (and join the 'trl' project):\n     - PyPI: https://pypi.org/\n     - Test PyPI: https://test.pypi.org/\n\n1. Change the version in:\n   - __init__.py\n   - setup.py\n\n2. Commit these changes: \"git commit -m 'Release: VERSION'\"\n\n3. Add a tag in git to mark the release: \"git tag VERSION -m 'Add tag VERSION for pypi'\"\n   Push the tag to remote: git push --tags origin main\n\n4. Build both the sources and the wheel. Do not change anything in setup.py between\n   creating the wheel and the source distribution (obviously).\n\n   First, delete any \"build\" directory that may exist from previous builds.\n\n   For the wheel, run: \"python setup.py bdist_wheel\" in the top level directory.\n   (this will build a wheel for the python version you use to build it).\n\n   For the sources, run: \"python setup.py sdist\"\n   You should now have a /dist directory with both .whl and .tar.gz source versions.\n\n5. Check that everything looks correct by uploading the package to the pypi test server:\n\n   twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/\n\n   Check that you can install it in a virtualenv/notebook by running:\n   pip install huggingface_hub fsspec aiohttp\n   pip install -U tqdm\n   pip install -i https://testpypi.python.org/pypi evaluate\n\n6. Upload the final version to actual pypi:\n   twine upload dist/* -r pypi\n\n7. Fill release notes in the tag in github once everything is looking hunky-dory.\n\n8. Change the version in __init__.py and setup.py to X.X.X+1.dev0 (e.g. VERSION=1.18.3 -> 1.18.4.dev0).\n   Then push the change with a message 'set dev version'\n\"\"\"\n\nimport os\n\nfrom setuptools import find_packages, setup\n\n\n__version__ = \"0.12.0.dev0\"  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)\n\nREQUIRED_PKGS = [\n    \"torch>=1.4.0\",\n    \"transformers>=4.40.0\",\n    \"numpy>=1.18.2;platform_system!='Windows'\",\n    \"numpy<2;platform_system=='Windows'\",\n    \"accelerate\",\n    \"datasets\",\n    \"tyro>=0.5.11\",\n]\nEXTRAS = {\n    \"test\": [\n        \"parameterized\",\n        \"peft>=0.8.0\",\n        \"pytest\",\n        \"pytest-xdist\",\n        \"pytest-cov\",\n        \"pytest-xdist\",\n        \"scikit-learn\",\n        \"Pillow\",\n        \"pytest-rerunfailures\",\n        \"llm-blender>=0.0.2\",\n    ],\n    \"peft\": [\"peft>=0.8.0\"],\n    \"liger\": [\"liger-kernel>=0.2.1\"],\n    \"diffusers\": [\"diffusers>=0.18.0\"],\n    \"deepspeed\": [\"deepspeed>=0.14.4\"],\n    \"benchmark\": [\"wandb\", \"ghapi\", \"openrlbenchmark==0.2.1a5\", \"requests\", \"deepspeed\"],\n    \"quantization\": [\"bitsandbytes<=0.41.1\"],\n    \"llm_judge\": [\"openai>=1.23.2\", \"huggingface_hub>=0.22.2\", \"llm-blender>=0.0.2\"],\n}\nEXTRAS[\"dev\"] = []\nfor reqs in EXTRAS.values():\n    EXTRAS[\"dev\"].extend(reqs)\n\ntry:\n    file_path = os.path.dirname(os.path.abspath(__file__))\n    os.symlink(os.path.join(file_path, \"examples/scripts\"), os.path.join(file_path, \"trl/commands/scripts\"))\n\n    setup(\n        name=\"trl\",\n        license=\"Apache 2.0\",\n        classifiers=[\n            \"Development Status :: 2 - Pre-Alpha\",\n            \"Intended Audience :: Developers\",\n            \"Intended Audience :: Science/Research\",\n            \"License :: OSI Approved :: Apache Software License\",\n            \"Natural Language :: English\",\n            \"Operating System :: OS Independent\",\n            \"Programming Language :: Python :: 3\",\n            \"Programming Language :: Python :: 3.9\",\n            \"Programming Language :: Python :: 3.10\",\n            \"Programming Language :: Python :: 3.11\",\n        ],\n        url=\"https://github.com/huggingface/trl\",\n        entry_points={\n            \"console_scripts\": [\"trl=trl.commands.cli:main\"],\n        },\n        include_package_data=True,\n        package_data={\"trl\": [\"commands/scripts/config/*\", \"commands/scripts/*\"]},\n        packages=find_packages(exclude={\"tests\"}),\n        install_requires=REQUIRED_PKGS,\n        extras_require=EXTRAS,\n        python_requires=\">=3.7\",\n        long_description=open(\"README.md\", encoding=\"utf-8\").read(),\n        long_description_content_type=\"text/markdown\",\n        zip_safe=False,\n        version=__version__,\n        description=\"Train transformer language models with reinforcement learning.\",\n        keywords=\"ppo, transformers, huggingface, gpt2, language modeling, rlhf\",\n        author=\"Leandro von Werra\",\n        author_email=\"leandro.vonwerra@gmail.com\",\n    )\nfinally:\n    os.unlink(os.path.join(file_path, \"trl/commands/scripts\"))\n\n\n\n# Contributor Covenant Code of Conduct\n\n## Our Pledge\n\nWe as members, contributors, and leaders pledge to make participation in our\ncommunity a harassment-free experience for everyone, regardless of age, body\nsize, visible or invisible disability, ethnicity, sex characteristics, gender\nidentity and expression, level of experience, education, socio-economic status,\nnationality, personal appearance, race, caste, color, religion, or sexual\nidentity and orientation.\n\nWe pledge to act and interact in ways that contribute to an open, welcoming,\ndiverse, inclusive, and healthy community.\n\n## Our Standards\n\nExamples of behavior that contributes to a positive environment for our\ncommunity include:\n\n* Demonstrating empathy and kindness toward other people\n* Being respectful of differing opinions, viewpoints, and experiences\n* Giving and gracefully accepting constructive feedback\n* Accepting responsibility and apologizing to those affected by our mistakes,\n  and learning from the experience\n* Focusing on what is best not just for us as individuals, but for the overall\n  community\n\nExamples of unacceptable behavior include:\n\n* The use of sexualized language or imagery, and sexual attention or advances of\n  any kind\n* Trolling, insulting or derogatory comments, and personal or political attacks\n* Public or private harassment\n* Publishing others' private information, such as a physical or email address,\n  without their explicit permission\n* Other conduct which could reasonably be considered inappropriate in a\n  professional setting\n\n## Enforcement Responsibilities\n\nCommunity leaders are responsible for clarifying and enforcing our standards of\nacceptable behavior and will take appropriate and fair corrective action in\nresponse to any behavior that they deem inappropriate, threatening, offensive,\nor harmful.\n\nCommunity leaders have the right and responsibility to remove, edit, or reject\ncomments, commits, code, wiki edits, issues, and other contributions that are\nnot aligned to this Code of Conduct, and will communicate reasons for moderation\ndecisions when appropriate.\n\n## Scope\n\nThis Code of Conduct applies within all community spaces, and also applies when\nan individual is officially representing the community in public spaces.\nExamples of representing our community include using an official e-mail address,\nposting via an official social media account, or acting as an appointed\nrepresentative at an online or offline event.\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be\nreported to the community leaders responsible for enforcement at\nfeedback@huggingface.co.\nAll complaints will be reviewed and investigated promptly and fairly.\n\nAll community leaders are obligated to respect the privacy and security of the\nreporter of any incident.\n\n## Enforcement Guidelines\n\nCommunity leaders will follow these Community Impact Guidelines in determining\nthe consequences for any action they deem in violation of this Code of Conduct:\n\n### 1. Correction\n\n**Community Impact**: Use of inappropriate language or other behavior deemed\nunprofessional or unwelcome in the community.\n\n**Consequence**: A private, written warning from community leaders, providing\nclarity around the nature of the violation and an explanation of why the\nbehavior was inappropriate. A public apology may be requested.\n\n### 2. Warning\n\n**Community Impact**: A violation through a single incident or series of\nactions.\n\n**Consequence**: A warning with consequences for continued behavior. No\ninteraction with the people involved, including unsolicited interaction with\nthose enforcing the Code of Conduct, for a specified period of time. This\nincludes avoiding interactions in community spaces as well as external channels\nlike social media. Violating these terms may lead to a temporary or permanent\nban.\n\n### 3. Temporary Ban\n\n**Community Impact**: A serious violation of community standards, including\nsustained inappropriate behavior.\n\n**Consequence**: A temporary ban from any sort of interaction or public\ncommunication with the community for a specified period of time. No public or\nprivate interaction with the people involved, including unsolicited interaction\nwith those enforcing the Code of Conduct, is allowed during this period.\nViolating these terms may lead to a permanent ban.\n\n### 4. Permanent Ban\n\n**Community Impact**: Demonstrating a pattern of violation of community\nstandards, including sustained inappropriate behavior, harassment of an\nindividual, or aggression toward or disparagement of classes of individuals.\n\n**Consequence**: A permanent ban from any sort of public interaction within the\ncommunity.\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage],\nversion 2.1, available at\n[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].\n\nCommunity Impact Guidelines were inspired by\n[Mozilla's code of conduct enforcement ladder][Mozilla CoC].\n\nFor answers to common questions about this code of conduct, see the FAQ at\n[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at\n[https://www.contributor-covenant.org/translations][translations].\n\n[homepage]: https://www.contributor-covenant.org\n[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html\n[Mozilla CoC]: https://github.com/mozilla/diversity\n[FAQ]: https://www.contributor-covenant.org/faq\n[translations]: https://www.contributor-covenant.org/translations\n\n# What does this PR do?\n\n<!--\nCongratulations! You've made it this far! You're not quite done yet though.\n\nOnce merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.\n\nThen, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.\n\nOnce you're done, someone will review your PR shortly. They may suggest changes to make the code even better.\n-->\n\n<!-- Remove if not applicable -->\n\nFixes # (issue)\n\n\n## Before submitting\n- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).\n- [ ] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),\n      Pull Request section?\n- [ ] Was this discussed/approved via a GitHub issue? Please add a link\n      to it if that's the case.\n- [ ] Did you make sure to update the documentation with your changes? Here are the\n      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).\n- [ ] Did you write any new necessary tests?\n\n\n## Who can review?\n\nAnyone in the community is free to review the PR once the tests have passed. Feel free to tag\nmembers/contributors who may be interested in your PR.\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom typing import Callable\n\nfrom datasets import Dataset, load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom trl.extras.dataset_formatting import get_formatting_func_from_dataset\nfrom trl.models.utils import ChatMlSpecialTokens, setup_chat_format\n\n\nclass DatasetFormattingTestCase(unittest.TestCase):\n    def setUp(self):\n        self.llama_tokenizer = AutoTokenizer.from_pretrained(\"hf-internal-testing/llama-tokenizer\")\n        self.chatml_tokenizer = AutoTokenizer.from_pretrained(\"philschmid/gpt2-chatml-tokenizer\")\n\n    def test_get_formatting_func_from_dataset_with_chatml_messages(self):\n        dataset = Dataset.from_dict(\n            {\n                \"messages\": [\n                    [\n                        {\"role\": \"system\", \"content\": \"You are helpful\"},\n                        {\"role\": \"user\", \"content\": \"Hello\"},\n                        {\"role\": \"assistant\", \"content\": \"Hi, how can I help you?\"},\n                    ]\n                ]\n            }\n        )\n\n        # Llama tokenizer\n        formatting_func = get_formatting_func_from_dataset(dataset, self.llama_tokenizer)\n        assert isinstance(formatting_func, Callable)\n        formatted_text = formatting_func(dataset[0])\n        expected = \"<s>[INST] <<SYS>>\\nYou are helpful\\n<</SYS>>\\n\\nHello [/INST] Hi, how can I help you? </s>\"\n        assert formatted_text == expected\n        formatted_text = formatting_func(dataset[0:1])\n        assert formatted_text == [expected]\n\n        # ChatML tokenizer\n        formatting_func = get_formatting_func_from_dataset(dataset, self.chatml_tokenizer)\n        formatted_text = formatting_func(dataset[0])\n        expected = \"<|im_start|>system\\nYou are helpful<|im_end|>\\n<|im_start|>user\\nHello<|im_end|>\\n<|im_start|>assistant\\nHi, how can I help you?<|im_end|>\\n\"\n        assert formatted_text == expected\n        formatted_text = formatting_func(dataset[0:1])\n        assert formatted_text == [expected]\n\n    def test_get_formatting_func_from_dataset_with_chatml_conversations(self):\n        dataset = Dataset.from_dict(\n            {\n                \"conversations\": [\n                    [\n                        {\"role\": \"system\", \"content\": \"You are helpful\"},\n                        {\"role\": \"user\", \"content\": \"Hello\"},\n                        {\"role\": \"assistant\", \"content\": \"Hi, how can I help you?\"},\n                    ]\n                ]\n            }\n        )\n        # Llama tokenizer\n        formatting_func = get_formatting_func_from_dataset(dataset, self.llama_tokenizer)\n        assert isinstance(formatting_func, Callable)\n        formatted_text = formatting_func(dataset[0])\n        expected = \"<s>[INST] <<SYS>>\\nYou are helpful\\n<</SYS>>\\n\\nHello [/INST] Hi, how can I help you? </s>\"\n        assert formatted_text == expected\n        formatted_text = formatting_func(dataset[0:1])\n        assert formatted_text == [expected]\n\n        # ChatML tokenizer\n        formatting_func = get_formatting_func_from_dataset(dataset, self.chatml_tokenizer)\n        formatted_text = formatting_func(dataset[0])\n        expected = \"<|im_start|>system\\nYou are helpful<|im_end|>\\n<|im_start|>user\\nHello<|im_end|>\\n<|im_start|>assistant\\nHi, how can I help you?<|im_end|>\\n\"\n        assert formatted_text == expected\n        formatted_text = formatting_func(dataset[0:1])\n        assert formatted_text == [expected]\n\n    def test_get_formatting_func_from_dataset_with_instruction(self):\n        dataset = Dataset.from_list(\n            [{\"prompt\": \"What is 2+2?\", \"completion\": \"4\"}, {\"prompt\": \"What is 3+3?\", \"completion\": \"6\"}]\n        )\n        formatting_func = get_formatting_func_from_dataset(dataset, self.llama_tokenizer)\n        assert formatting_func is not None\n        assert isinstance(formatting_func, Callable)\n        formatted_text = formatting_func(dataset[0])\n        assert formatted_text == \"<s>[INST] What is 2+2? [/INST] 4 </s>\"\n        formatted_text = formatting_func(dataset[0:1])\n        assert formatted_text == [\"<s>[INST] What is 2+2? [/INST] 4 </s>\"]\n\n    def test_get_formatting_func_from_dataset_from_hub(self):\n        ds_1 = load_dataset(\"philschmid/trl-test-instruction\", split=\"train\")\n        ds_2 = load_dataset(\"philschmid/dolly-15k-oai-style\", split=\"train\")\n        for ds in [ds_1, ds_2]:\n            formatting_func = get_formatting_func_from_dataset(ds, self.llama_tokenizer)\n            assert formatting_func is not None\n            assert isinstance(formatting_func, Callable)\n        ds_3 = load_dataset(\"philschmid/guanaco-sharegpt-style\", split=\"train\")\n        formatting_func = get_formatting_func_from_dataset(ds_3, self.llama_tokenizer)\n        assert formatting_func is None\n\n    def test_get_formatting_func_from_dataset_with_unknown_format(self):\n        dataset = Dataset.from_dict({\"text\": \"test\"})\n        formatting_func = get_formatting_func_from_dataset(dataset, self.llama_tokenizer)\n        assert formatting_func is None\n\n\nclass SetupChatFormatTestCase(unittest.TestCase):\n    def setUp(self):\n        self.tokenizer = AutoTokenizer.from_pretrained(\"hf-internal-testing/llama-tokenizer\")\n        self.model = AutoModelForCausalLM.from_pretrained(\"hf-internal-testing/tiny-random-MistralForCausalLM\")\n\n    def test_setup_chat_format(self):\n        original_tokenizer_len = len(self.tokenizer)\n        modified_model, modified_tokenizer = setup_chat_format(\n            self.model, self.tokenizer, format=\"chatml\", resize_to_multiple_of=64\n        )\n\n        _chatml = ChatMlSpecialTokens()\n        # Check if special tokens are correctly set\n        assert modified_tokenizer.eos_token == \"<|im_end|>\"\n        assert modified_tokenizer.pad_token == \"<|im_end|>\"\n        assert modified_tokenizer.bos_token == \"<|im_start|>\"\n        assert modified_tokenizer.eos_token == _chatml.eos_token\n        assert modified_tokenizer.pad_token == _chatml.pad_token\n        assert modified_tokenizer.bos_token == _chatml.bos_token\n        assert len(modified_tokenizer) == (original_tokenizer_len + 2)\n        assert (self.model.get_input_embeddings().weight.shape[0] % 64) == 0\n        assert self.model.get_input_embeddings().weight.shape[0] == (original_tokenizer_len + 64)\n\n    def test_example_with_setup_model(self):\n        modified_model, modified_tokenizer = setup_chat_format(\n            self.model,\n            self.tokenizer,\n        )\n        messages = [\n            {\"role\": \"system\", \"content\": \"You are helpful\"},\n            {\"role\": \"user\", \"content\": \"Hello\"},\n            {\"role\": \"assistant\", \"content\": \"Hi, how can I help you?\"},\n        ]\n        prompt = modified_tokenizer.apply_chat_template(messages, tokenize=False)\n\n        assert (\n            prompt\n            == \"<|im_start|>system\\nYou are helpful<|im_end|>\\n<|im_start|>user\\nHello<|im_end|>\\n<|im_start|>assistant\\nHi, how can I help you?<|im_end|>\\n\"\n        )\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tempfile\nimport unittest\n\nimport numpy as np\nimport pytest\nimport torch\nfrom datasets import Dataset, features, load_dataset\nfrom parameterized import parameterized\nfrom PIL import Image\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoModelForSeq2SeqLM,\n    AutoModelForVision2Seq,\n    AutoProcessor,\n    AutoTokenizer,\n)\nfrom transformers.testing_utils import require_bitsandbytes, require_peft\n\nfrom trl import DPOConfig, DPOTrainer, FDivergenceType\nfrom trl.trainer.dpo_trainer import _build_tokenized_answer, _truncate_tokens\n\nfrom .testing_utils import require_no_wandb\n\n\nclass TestBuildTokenizedAnswer(unittest.TestCase):\n    def setUp(self):\n        self.tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n    def test_basic_functionality(self):\n        prompt = \"Hello, how are you?\"\n        answer = \"I'm doing well, thank you!\"\n\n        result = _build_tokenized_answer(prompt, answer, tokenizer=self.tokenizer)\n\n        self.assertIn(\"prompt_input_ids\", result)\n        self.assertIn(\"prompt_attention_mask\", result)\n        self.assertIn(\"input_ids\", result)\n        self.assertIn(\"attention_mask\", result)\n\n        self.assertEqual(len(result[\"prompt_input_ids\"]), len(result[\"prompt_attention_mask\"]))\n        self.assertEqual(len(result[\"input_ids\"]), len(result[\"attention_mask\"]))\n\n        decoded_prompt = self.tokenizer.decode(result[\"prompt_input_ids\"])\n        self.assertTrue(prompt in decoded_prompt)\n\n        decoded_answer = self.tokenizer.decode(result[\"input_ids\"])\n        self.assertTrue(answer in decoded_answer)\n\n    def test_with_processor(self):\n        def mock_processor(text, images=None, add_special_tokens=True):\n            return {\"input_ids\": torch.tensor([[1, 2, 3]]), \"attention_mask\": torch.tensor([[1, 1, 1]])}\n\n        prompt = \"Describe this image:\"\n        answer = \"A beautiful sunset over the ocean.\"\n\n        result = _build_tokenized_answer(prompt, answer, processor=mock_processor)\n\n        self.assertIn(\"prompt_input_ids\", result)\n        self.assertIn(\"prompt_attention_mask\", result)\n        self.assertIn(\"input_ids\", result)\n        self.assertIn(\"attention_mask\", result)\n\n        self.assertEqual(result[\"prompt_input_ids\"], [1, 2, 3])\n        self.assertEqual(result[\"prompt_attention_mask\"], [1, 1, 1])\n\n    def test_token_merging(self):\n        prompt = \"The quick brown\"\n        answer = \" fox jumps over the lazy dog.\"\n\n        result = _build_tokenized_answer(prompt, answer, tokenizer=self.tokenizer)\n\n        full_text = prompt + answer\n        full_tokenized = self.tokenizer(full_text, add_special_tokens=False)\n\n        self.assertEqual(result[\"prompt_input_ids\"] + result[\"input_ids\"], full_tokenized[\"input_ids\"])\n\n    def test_vision_model(self):\n        def mock_vision_processor(text, images=None, add_special_tokens=True):\n            return {\n                \"input_ids\": torch.tensor([[1, 2, 3]]),\n                \"attention_mask\": torch.tensor([[1, 1, 1]]),\n                \"pixel_values\": torch.rand(1, 3, 224, 224),\n                \"pixel_attention_mask\": torch.ones(1, 224, 224),\n            }\n\n        prompt = \"Describe this image:\"\n        answer = \"A cat sitting on a windowsill.\"\n\n        result = _build_tokenized_answer(prompt, answer, processor=mock_vision_processor)\n\n        self.assertIn(\"prompt_pixel_values\", result)\n        self.assertIn(\"prompt_pixel_attention_mask\", result)\n        self.assertTrue(torch.is_tensor(result[\"prompt_pixel_values\"]))\n        self.assertTrue(torch.is_tensor(result[\"prompt_pixel_attention_mask\"]))\n\n\nclass TestTruncateTokens(unittest.TestCase):\n    def setUp(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            self.training_args = DPOConfig(\n                max_length=20, max_prompt_length=10, truncation_mode=\"keep_start\", output_dir=tmp_dir\n            )\n\n    def test_truncate_tokens(self):\n        chosen_tokens = [\n            {\n                \"prompt_input_ids\": list(range(15)),\n                \"prompt_attention_mask\": [1] * 15,\n                \"input_ids\": list(range(10)),\n                \"attention_mask\": [1] * 10,\n            }\n        ]\n        rejected_tokens = [\n            {\n                \"prompt_input_ids\": list(range(15)),\n                \"prompt_attention_mask\": [1] * 15,\n                \"input_ids\": list(range(12)),\n                \"attention_mask\": [1] * 12,\n            }\n        ]\n        prompt_tokens = [{\"prompt_input_ids\": list(range(15)), \"prompt_attention_mask\": [1] * 15}]\n\n        _truncate_tokens(chosen_tokens, rejected_tokens, prompt_tokens, self.training_args)\n\n        # Check if prompt is truncated correctly\n        self.assertEqual(len(chosen_tokens[0][\"prompt_input_ids\"]), 10)\n        self.assertEqual(len(chosen_tokens[0][\"prompt_attention_mask\"]), 10)\n        self.assertEqual(len(rejected_tokens[0][\"prompt_input_ids\"]), 10)\n        self.assertEqual(len(rejected_tokens[0][\"prompt_attention_mask\"]), 10)\n        self.assertEqual(len(prompt_tokens[0][\"prompt_input_ids\"]), 10)\n        self.assertEqual(len(prompt_tokens[0][\"prompt_attention_mask\"]), 10)\n\n        # Check if responses are truncated correctly\n        self.assertEqual(len(chosen_tokens[0][\"input_ids\"]), 10)\n        self.assertEqual(len(chosen_tokens[0][\"attention_mask\"]), 10)\n        self.assertEqual(len(rejected_tokens[0][\"input_ids\"]), 10)\n        self.assertEqual(len(rejected_tokens[0][\"attention_mask\"]), 10)\n\n    def test_truncation_mode_keep_end(self):\n        self.training_args.truncation_mode = \"keep_end\"\n        chosen_tokens = [\n            {\n                \"prompt_input_ids\": list(range(15)),\n                \"prompt_attention_mask\": [1] * 15,\n                \"input_ids\": list(range(15, 25)),\n                \"attention_mask\": [1] * 10,\n            }\n        ]\n        rejected_tokens = [\n            {\n                \"prompt_input_ids\": list(range(15)),\n                \"prompt_attention_mask\": [1] * 15,\n                \"input_ids\": list(range(15, 28)),\n                \"attention_mask\": [1] * 13,\n            }\n        ]\n        prompt_tokens = [{\"prompt_input_ids\": list(range(15)), \"prompt_attention_mask\": [1] * 15}]\n\n        _truncate_tokens(chosen_tokens, rejected_tokens, prompt_tokens, self.training_args)\n\n        # Check if prompt is truncated correctly from the end\n        self.assertEqual(prompt_tokens[0][\"prompt_input_ids\"], list(range(5, 15)))\n        self.assertEqual(prompt_tokens[0][\"prompt_attention_mask\"], [1] * 10)\n\n        # Check if chosen tokens are truncated correctly\n        self.assertEqual(chosen_tokens[0][\"prompt_input_ids\"], list(range(5, 15)))\n        self.assertEqual(chosen_tokens[0][\"prompt_attention_mask\"], [1] * 10)\n        self.assertEqual(chosen_tokens[0][\"input_ids\"], list(range(15, 25)))\n        self.assertEqual(chosen_tokens[0][\"attention_mask\"], [1] * 10)\n\n        # Check if rejected tokens are truncated correctly\n        self.assertEqual(rejected_tokens[0][\"prompt_input_ids\"], list(range(5, 15)))\n        self.assertEqual(rejected_tokens[0][\"prompt_attention_mask\"], [1] * 10)\n        self.assertEqual(rejected_tokens[0][\"input_ids\"], list(range(15, 25)))\n        self.assertEqual(rejected_tokens[0][\"attention_mask\"], [1] * 10)\n\n    def test_invalid_truncation_mode(self):\n        self.training_args.truncation_mode = \"invalid_mode\"\n        with self.assertRaises(ValueError):\n            _truncate_tokens([], [], [], self.training_args)\n\n\nclass DPOTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/T5ForConditionalGeneration-correct-vocab-calibrated\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_ref_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", \"sigmoid\", True],\n            [\"t5\", \"hinge\", False],\n            [\"gpt2\", \"ipo\", False],\n            [\"t5\", \"ipo\", True],\n            [\"gpt2\", \"aot_pair\", True],\n            [\"t5\", \"aot_pair\", False],\n            [\"gpt2\", \"aot\", True],\n            [\"t5\", \"aot\", False],\n            [\"gpt2\", \"bco_pair\", False],\n            [\"t5\", \"bco_pair\", True],\n            [\"gpt2\", \"sppo_hard\", False],\n            [\"t5\", \"sppo_hard\", True],\n            [\"gpt2\", \"nca_pair\", False],\n            [\"t5\", \"nca_pair\", True],\n            [\"gpt2\", \"robust\", True],\n            [\"gpt2\", \"exo_pair\", False],\n            [\"t5\", \"exo_pair\", True],\n            [\"gpt2\", \"apo_zero\", True],\n            [\"t5\", \"apo_down\", False],\n        ]\n    )\n    def test_dpo_trainer(self, name, loss_type, pre_compute):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                loss_type=loss_type,\n                precompute_ref_log_probs=pre_compute,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            if name == \"gpt2\":\n                model = self.model\n                ref_model = self.ref_model\n                tokenizer = self.tokenizer\n            elif name == \"t5\":\n                model = self.t5_model\n                ref_model = self.t5_ref_model\n                tokenizer = self.t5_tokenizer\n\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=ref_model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.allclose(param, new_param, rtol=1e-12, atol=1e-12)\n\n    @parameterized.expand(\n        [\n            [None, \"Test when rpo_alpha is set to None\"],\n            [0.5, \"Test when rpo_alpha is set to 0.5\"],\n        ]\n    )\n    def test_dpo_trainer_without_providing_ref_model(self, rpo_alpha, _):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                precompute_ref_log_probs=True,\n                rpo_alpha=rpo_alpha,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            trainer = DPOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.equal(param, new_param)\n\n    def test_dpo_trainer_with_ref_model_is_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            with self.assertRaises(ValueError):\n                DPOTrainer(\n                    model=self.model,\n                    ref_model=self.model,  # ref_model can't be the same as model\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n    @require_peft\n    def test_dpo_trainer_without_providing_ref_model_with_lora(self):\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                precompute_ref_log_probs=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            trainer = DPOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                if \"lora\" in n:\n                    new_param = trainer.model.get_parameter(n)\n                    # check the params have changed - ignore 0 biases\n                    if param.sum() != 0:\n                        assert not torch.equal(param, new_param)\n\n    def test_dpo_trainer_padding_token_is_none(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n            tokenizer.pad_token = None\n\n            with self.assertRaisesRegex(\n                ValueError,\n                expected_regex=r\"Padding is enabled, but the tokenizer is not configured with a padding token.\"\n                r\" Explicitly set `tokenizer.pad_token` \\(e.g. `tokenizer.pad_token = tokenizer.eos_token`\\)\"\n                r\" before calling the trainer.\",\n            ):\n                trainer = DPOTrainer(\n                    model=self.model,\n                    ref_model=None,\n                    args=training_args,\n                    tokenizer=tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                )\n\n                trainer.train()\n\n    def test_dpo_trainer_w_dataset_num_proc(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                dataset_num_proc=5,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n            tokenizer.pad_token = None\n\n            with self.assertRaisesRegex(\n                ValueError,\n                expected_regex=r\"Padding is enabled, but the tokenizer is not configured with a padding token.\"\n                r\" Explicitly set `tokenizer.pad_token` \\(e.g. `tokenizer.pad_token = tokenizer.eos_token`\\)\"\n                r\" before calling the trainer.\",\n            ):\n                trainer = DPOTrainer(\n                    model=self.model,\n                    ref_model=None,\n                    args=training_args,\n                    tokenizer=tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                )\n\n                trainer.train()\n\n    def test_tr_dpo_trainer(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                precompute_ref_log_probs=False,\n                sync_ref_model=True,\n                ref_model_mixup_alpha=0.5,\n                ref_model_sync_steps=1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            trainer = DPOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                beta=0.1,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            # params of the ref model as its the same as the model\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.ref_model.get_parameter(n)\n                # check the ref model's params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.equal(param, new_param)\n\n    @require_no_wandb\n    def test_dpo_trainer_generate_during_eval_no_wandb(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                generate_during_eval=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            with self.assertRaisesRegex(\n                ValueError,\n                expected_regex=\"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install `wandb` to resolve.\",\n            ):\n                DPOTrainer(\n                    model=self.model,\n                    ref_model=None,\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                )\n\n    @require_peft\n    def test_dpo_lora_save(self):\n        from peft import LoraConfig, get_peft_model\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        model_peft = get_peft_model(model, lora_config)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                precompute_ref_log_probs=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model_peft,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save peft adapter\n            trainer.save_model()\n\n            # assert that the model is loaded without giving OSError\n            try:\n                AutoModelForCausalLM.from_pretrained(tmp_dir)\n            except OSError:\n                self.fail(\"Loading the saved peft adapter failed\")\n\n    @require_peft\n    @require_bitsandbytes\n    def test_dpo_lora_bf16_autocast_llama(self):\n        # Note this test only works on compute capability > 7 GPU devices\n        from peft import LoraConfig\n\n        model_id = \"trl-internal-testing/tiny-random-LlamaForCausalLM\"\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                bf16=True,\n                beta=0.1,\n                generate_during_eval=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save peft adapter\n            trainer.save_model()\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", \"sigmoid\", False, False],\n            [\"gpt2\", \"sigmoid\", False, True],\n            [\"gpt2\", \"sigmoid\", True, False],\n            [\"gpt2\", \"sigmoid\", True, True],\n            [\"gpt2\", \"ipo\", False, False],\n            [\"gpt2\", \"ipo\", False, True],\n            [\"gpt2\", \"ipo\", True, False],\n            [\"gpt2\", \"ipo\", True, True],\n            [\"gpt2\", \"aot_pair\", False, False],\n            [\"gpt2\", \"aot_pair\", False, True],\n            [\"gpt2\", \"aot_pair\", True, False],\n            [\"gpt2\", \"aot_pair\", True, True],\n            [\"gpt2\", \"aot\", False, False],\n            [\"gpt2\", \"aot\", False, True],\n            [\"gpt2\", \"aot\", True, False],\n            [\"gpt2\", \"aot\", True, True],\n            [\"gpt2\", \"bco_pair\", False, False],\n            [\"gpt2\", \"bco_pair\", False, True],\n            [\"gpt2\", \"bco_pair\", True, False],\n            [\"gpt2\", \"bco_pair\", True, True],\n            [\"gpt2\", \"robust\", False, False],\n            [\"gpt2\", \"robust\", False, True],\n            [\"gpt2\", \"robust\", True, False],\n            [\"gpt2\", \"robust\", True, True],\n        ]\n    )\n    @require_bitsandbytes\n    @require_peft\n    @unittest.skip(\"You need a GPU with bf16 support in order to run these tests\")\n    def test_dpo_lora_bf16_autocast(self, name, loss_type, pre_compute, gen_during_eval):\n        # Note this test only works on compute capability > 7 GPU devices\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(self.model_id, load_in_4bit=True)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                bf16=True,\n                beta=0.1,\n                generate_during_eval=gen_during_eval,\n                loss_type=loss_type,\n                precompute_ref_log_probs=pre_compute,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save peft adapter\n            trainer.save_model()\n\n    @require_peft\n    def test_dpo_lora_tags(self):\n        from peft import LoraConfig\n\n        model_id = \"trl-internal-testing/tiny-random-LlamaForCausalLM\"\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            assert trainer.model.model_tags == trainer._tag_names\n\n    @require_peft\n    def test_dpo_tags(self):\n        model_id = \"HuggingFaceM4/tiny-random-LlamaForCausalLM\"\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            assert trainer.model.model_tags == trainer._tag_names\n\n    @require_peft\n    def test_dpo_lora_force_use_ref(self):\n        from peft import LoraConfig, get_peft_model\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        model_peft = get_peft_model(model, lora_config)\n\n        ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            with self.assertRaises(ValueError):\n                # passing a peft_model as model and ref_model should error out,\n                # unless you pass `force_use_ref_model`\n                trainer = DPOTrainer(\n                    model=model_peft,\n                    ref_model=ref_model,\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                    peft_config=lora_config,\n                )\n\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                force_use_ref_model=True,\n                report_to=\"none\",\n            )\n\n            trainer = DPOTrainer(\n                model=model_peft,\n                ref_model=ref_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n    def test_dpo_trainer_torch_dtype(self):\n        # See https://github.com/huggingface/trl/issues/1751\n        dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=1,\n                model_init_kwargs={\"torch_dtype\": \"float16\"},\n                ref_model_init_kwargs={\"torch_dtype\": \"float16\"},\n                report_to=\"none\",\n            )\n\n            trainer = DPOTrainer(\n                model=self.model_id,\n                ref_model=self.model_id,\n                tokenizer=self.tokenizer,\n                args=training_args,\n                train_dataset=dummy_dataset[\"train\"],\n            )\n            assert trainer.model.config.torch_dtype == torch.float16\n            assert trainer.ref_model.config.torch_dtype == torch.float16\n\n        # Now test when `torch_dtype` is provided but is wrong to either the model or the ref_model\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=1,\n                model_init_kwargs={\"torch_dtype\": -1},\n                report_to=\"none\",\n            )\n\n            with pytest.raises(\n                ValueError,\n                match=\"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got -1.\",\n            ):\n                _ = DPOTrainer(\n                    model=self.model_id,\n                    tokenizer=self.tokenizer,\n                    args=training_args,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=1,\n                ref_model_init_kwargs={\"torch_dtype\": -1},\n                report_to=\"none\",\n            )\n\n            with pytest.raises(\n                ValueError,\n                match=\"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got -1.\",\n            ):\n                _ = DPOTrainer(\n                    model=self.model_id,\n                    ref_model=self.model_id,\n                    tokenizer=self.tokenizer,\n                    args=training_args,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n    def test_dpo_loss_alpha_div_f(self):\n        model_id = \"trl-internal-testing/tiny-random-LlamaForCausalLM\"\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                f_divergence_type=FDivergenceType.ALPHA_DIVERGENCE.value,\n                f_alpha_divergence_coef=0.5,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            # Fake chosen and rejected log probs\n            policy_chosen_logps = torch.FloatTensor([410.0, 0.1])\n            policy_rejected_logps = torch.FloatTensor([810.5, 0.2])\n            reference_chosen_logps = torch.FloatTensor([-610.0, -0.1])\n            reference_rejected_logps = torch.FloatTensor([110.6, 0.5])\n            losses, _, _ = trainer.dpo_loss(\n                policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps\n            )\n            assert torch.isfinite(losses).cpu().numpy().all()\n\n    def test_dpo_loss_js_div_f(self):\n        model_id = \"trl-internal-testing/tiny-random-LlamaForCausalLM\"\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                f_divergence_type=FDivergenceType.JS_DIVERGENCE.value,\n                f_alpha_divergence_coef=0.5,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            # dpo train lora model with a lora config\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            # Fake chosen and rejected log probs\n            policy_chosen_logps = torch.FloatTensor([410.0, 0.1])\n            policy_rejected_logps = torch.FloatTensor([95.5, 0.2])\n            reference_chosen_logps = torch.FloatTensor([-610.0, -0.1])\n            reference_rejected_logps = torch.FloatTensor([5.5, 0.5])\n            losses, _, _ = trainer.dpo_loss(\n                policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps\n            )\n            assert torch.isfinite(losses).cpu().numpy().all()\n\n\nclass DPOVisionTrainerTester(unittest.TestCase):\n    @parameterized.expand(\n        [\n            [\"trl-internal-testing/tiny-random-idefics2\"],\n            [\"trl-internal-testing/tiny-random-paligemma\"],\n            [\"trl-internal-testing/tiny-random-llava-1.5\"],\n        ]\n    )\n    def test_vdpo_trainer(self, model_id):\n        # fmt: off\n        dataset_dict = {\n            \"prompt\": [\n                [{\"role\": \"user\", \"content\": [{\"type\": \"image\"}, {\"type\": \"text\", \"text\": \"Describe the image in great detail.\"}]}],\n                [{\"role\": \"user\", \"content\": [{\"type\": \"image\"}, {\"type\": \"text\", \"text\": \"Is this bus in the USA?\"}]}],\n                [{\"role\": \"user\", \"content\": [{\"type\": \"image\"}, {\"type\": \"text\", \"text\": \"Give a thorough description of the image.\"}]}],\n                [{\"role\": \"user\", \"content\": [{\"type\": \"image\"}, {\"type\": \"text\", \"text\": \"Who are the people in the image?\"}]}],\n                [{\"role\": \"user\", \"content\": [{\"type\": \"image\"}, {\"type\": \"text\", \"text\": \"What ise written?\"}]}],\n            ],\n            \"chosen\": [\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"The image features a modern, multi-colored train.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"Yes, it can be assumed that this bus is in the USA.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"The image features a forest path.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"There are two individuals, possibly girls or women.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": '\"ccpb\".'}]}],\n            ],\n            \"rejected\": [\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"The image features a modern, colorful train.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"No, it's not in the USA.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"The image features a forest path surrounded by trees.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"In the image, there are two individuals.\"}]}],\n                [{\"role\": \"assistant\", \"content\": [{\"text\": '\"ccpb\".', \"type\": \"text\"}]}],\n            ],\n            \"images\": [\n                [Image.fromarray(np.random.randint(0, 255, (92, 33, 3), dtype=np.uint8))],\n                [Image.fromarray(np.random.randint(0, 255, (64, 48, 3), dtype=np.uint8))],\n                [Image.fromarray(np.random.randint(0, 255, (80, 152, 3), dtype=np.uint8))],\n                [Image.fromarray(np.random.randint(0, 255, (57, 24, 3), dtype=np.uint8))],\n                [Image.fromarray(np.random.randint(0, 255, (102, 48, 3), dtype=np.uint8))],\n            ],\n        }\n        # fmt: on\n        dataset = Dataset.from_dict(dataset_dict)\n        dataset = dataset.cast_column(\"images\", features.Sequence(features.Image()))\n\n        # Instantiate the model and processor\n        model = AutoModelForVision2Seq.from_pretrained(model_id)\n        ref_model = AutoModelForVision2Seq.from_pretrained(model_id)\n        processor = AutoProcessor.from_pretrained(model_id)\n\n        # Apply chat template to the dataset\n        def apply_chat_template(example):\n            example[\"prompt\"] = processor.apply_chat_template(example[\"prompt\"])\n            example[\"chosen\"] = processor.apply_chat_template(example[\"chosen\"])\n            example[\"rejected\"] = processor.apply_chat_template(example[\"rejected\"])\n            return example\n\n        dataset = dataset.map(apply_chat_template)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_length=512,\n                max_prompt_length=128,\n                remove_unused_columns=False,\n                report_to=\"none\",\n            )\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=ref_model,\n                args=training_args,\n                tokenizer=processor,\n                train_dataset=dataset,\n                eval_dataset=dataset,\n            )\n\n            # Save the initial weights, so we can check if they have changed after training\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # Check that the trainable params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                if param.sum() != 0:  # ignore 0 biases\n                    if model_id == \"trl-internal-testing/tiny-random-llava-1.5\" and (\n                        n.startswith(\"vision_tower.vision_model.encoder.layers.3\")\n                        or n == \"vision_tower.vision_model.post_layernorm.weight\"\n                    ):\n                        # For some reason, these params are not updated. This is probably not related to TRL, but to\n                        # the model itself. We should investigate this further, but for now we just skip these params.\n                        continue\n                    assert not torch.allclose(param, new_param, rtol=1e-12, atol=1e-12)\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport sys\nimport tempfile\nimport unittest\n\nimport pytest\nimport torch\nfrom transformers import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM\n\nfrom trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, create_reference_model\n\n\nALL_CAUSAL_LM_MODELS = [\n    \"trl-internal-testing/tiny-random-CodeGenForCausalLM\",\n    \"trl-internal-testing/tiny-random-GPTJForCausalLM\",\n    \"trl-internal-testing/tiny-random-GPTNeoForCausalLM\",\n    \"trl-internal-testing/tiny-random-GPTNeoXForCausalLM\",\n    \"trl-internal-testing/tiny-random-OPTForCausalLM\",\n    \"trl-internal-testing/tiny-random-BloomForCausalLM\",\n    \"trl-internal-testing/tiny-random-GPT2LMHeadModel\",\n    \"trl-internal-testing/tiny-random-CodeGenForCausalLM-sharded\",\n    \"trl-internal-testing/tiny-random-GPTNeoXForCausalLM-safetensors-sharded\",\n    \"trl-internal-testing/tiny-random-GPTNeoXForCausalLM-safetensors\",\n    \"trl-internal-testing/tiny-random-LlamaForCausalLM\",\n]\n\nALL_SEQ2SEQ_MODELS = [\n    \"trl-internal-testing/tiny-random-BartForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-BigBirdPegasusForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-BlenderbotForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-BlenderbotSmallForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-FSMTForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-LEDForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-LongT5ForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-M2M100ForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-MarianMTModel\",\n    \"trl-internal-testing/tiny-random-MBartForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-MT5ForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-MvpForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-PegasusForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-PegasusXForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-PLBartForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-ProphetNetForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-SwitchTransformersForConditionalGeneration\",\n    \"trl-internal-testing/tiny-random-T5ForConditionalGeneration\",\n]\n\n\nclass VHeadModelTester:\n    all_model_names = None\n    trl_model_class = None\n    transformers_model_class = None\n\n    def test_value_head(self):\n        r\"\"\"\n        Test if the v-head is added to the model successfully\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            assert hasattr(model, \"v_head\")\n\n    def test_value_head_shape(self):\n        r\"\"\"\n        Test if the v-head has the correct shape\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            assert model.v_head.summary.weight.shape[0] == 1\n\n    def test_value_head_init_random(self):\n        r\"\"\"\n        Test if the v-head has been randomly initialized.\n        We can check that by making sure the bias is different\n        than zeros by default.\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            assert not torch.allclose(model.v_head.summary.bias, torch.zeros_like(model.v_head.summary.bias))\n\n    def test_value_head_not_str(self):\n        r\"\"\"\n        Test if the v-head is added to the model successfully, by passing a non `PretrainedModel`\n        as an argument to `from_pretrained`.\n        \"\"\"\n        for model_name in self.all_model_names:\n            pretrained_model = self.transformers_model_class.from_pretrained(model_name)\n            model = self.trl_model_class.from_pretrained(pretrained_model)\n            assert hasattr(model, \"v_head\")\n\n    @unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\n    def test_from_save_trl(self):\n        \"\"\"\n        Test if the model can be saved and loaded from a directory and get the same weights\n        Including the additional modules (e.g. v_head)\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n\n            with tempfile.TemporaryDirectory() as tmp_dir:\n                model.save_pretrained(tmp_dir)\n\n                model_from_save = self.trl_model_class.from_pretrained(tmp_dir)\n\n            # Check if the weights are the same\n            for key in model_from_save.state_dict():\n                assert torch.allclose(model_from_save.state_dict()[key], model.state_dict()[key])\n\n    @unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\n    def test_from_save_trl_sharded(self):\n        \"\"\"\n        Test if the model can be saved and loaded from a directory and get the same weights - sharded case\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n\n            with tempfile.TemporaryDirectory() as tmp_dir:\n                model.save_pretrained(tmp_dir)\n\n                model_from_save = self.trl_model_class.from_pretrained(tmp_dir)\n\n            # Check if the weights are the same\n            for key in model_from_save.state_dict():\n                assert torch.allclose(model_from_save.state_dict()[key], model.state_dict()[key])\n\n    @unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\n    def test_from_save_transformers_sharded(self):\n        \"\"\"\n        Test if the model can be saved and loaded using transformers and get the same weights - sharded case\n        \"\"\"\n        for model_name in self.all_model_names:\n            transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)\n\n            trl_model = self.trl_model_class.from_pretrained(model_name)\n\n            with tempfile.TemporaryDirectory() as tmp_dir:\n                trl_model.save_pretrained(tmp_dir, max_shard_size=\"1MB\")\n                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(tmp_dir)\n\n            # Check if the weights are the same\n            for key in transformers_model.state_dict():\n                assert torch.allclose(\n                    transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]\n                )\n\n    @unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\n    def test_from_save_transformers(self):\n        \"\"\"\n        Test if the model can be saved and loaded using transformers and get the same weights.\n        We override the test of the super class to check if the weights are the same.\n        \"\"\"\n        for model_name in self.all_model_names:\n            transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)\n\n            trl_model = self.trl_model_class.from_pretrained(model_name)\n\n            with tempfile.TemporaryDirectory() as tmp_dir:\n                trl_model.save_pretrained(tmp_dir)\n                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(tmp_dir)\n\n            # Check if the weights are the same\n            for key in transformers_model.state_dict():\n                assert torch.allclose(\n                    transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]\n                )\n\n            # Check if the trl model has the same keys as the transformers model\n            # except the v_head\n            for key in trl_model.state_dict():\n                if \"v_head\" not in key:\n                    assert key in transformers_model.state_dict()\n                    # check if the weights are the same\n                    assert torch.allclose(trl_model.state_dict()[key], transformers_model.state_dict()[key])\n\n            # check if they have the same modules\n            assert set(transformers_model_from_save.state_dict().keys()) == set(transformers_model.state_dict().keys())\n\n\nclass CausalLMValueHeadModelTester(VHeadModelTester, unittest.TestCase):\n    \"\"\"\n    Testing suite for v-head models.\n    \"\"\"\n\n    all_model_names = ALL_CAUSAL_LM_MODELS\n    trl_model_class = AutoModelForCausalLMWithValueHead\n    transformers_model_class = AutoModelForCausalLM\n\n    def tearDown(self):\n        # free memory\n        gc.collect()\n\n    def test_inference(self):\n        r\"\"\"\n        Test if the model can be used for inference and outputs 3 values\n        - logits, loss, and value states\n        \"\"\"\n        EXPECTED_OUTPUT_SIZE = 3\n\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n            outputs = model(input_ids)\n\n            # Check if the outputs are of the right size - here\n            # we always output 3 values - logits, loss, and value states\n            assert len(outputs) == EXPECTED_OUTPUT_SIZE\n\n    def test_dropout_config(self):\n        r\"\"\"\n        Test if we instantiate a model by adding `summary_drop_prob` to the config\n        it will be added to the v_head\n        \"\"\"\n        for model_name in self.all_model_names:\n            pretrained_model = self.transformers_model_class.from_pretrained(model_name)\n            pretrained_model.config.summary_dropout_prob = 0.5\n            model = self.trl_model_class.from_pretrained(pretrained_model)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob\n\n    def test_dropout_kwargs(self):\n        r\"\"\"\n        Test if we instantiate a model by adding `summary_drop_prob` to the config\n        it will be added to the v_head\n        \"\"\"\n        for model_name in self.all_model_names:\n            v_head_kwargs = {\"summary_dropout_prob\": 0.5}\n\n            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == 0.5\n\n            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == 0.5\n\n    def test_generate(self):\n        r\"\"\"\n        Test if `generate` works for every model\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n            # Just check if the generation works\n            _ = model.generate(input_ids)\n\n    def test_raise_error_not_causallm(self):\n        # Test with a model without a LM head\n        model_id = \"trl-internal-testing/tiny-random-GPT2Model\"\n        # This should raise a ValueError\n        with pytest.raises(ValueError):\n            pretrained_model = AutoModelForCausalLM.from_pretrained(model_id)\n            _ = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model.transformer)\n\n    def test_transformers_bf16_kwargs(self):\n        r\"\"\"\n        Test if the transformers kwargs are correctly passed\n        Here we check that loading a model in half precision works as expected, i.e. the weights of\n        the `pretrained_model` attribute is loaded in half precision and you can run a dummy\n        forward pass without any issue.\n        \"\"\"\n        for model_name in self.all_model_names:\n            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16)\n\n            lm_head_namings = self.trl_model_class.lm_head_namings\n\n            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings)\n\n            for lm_head_naming in lm_head_namings:\n                if hasattr(trl_model.pretrained_model, lm_head_naming):\n                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16\n\n            dummy_input = torch.LongTensor([[0, 1, 0, 1]])\n\n            # check dummy forward pass works in half precision\n            _ = trl_model(dummy_input)\n\n    @unittest.skip(\"This test needs to be run manually due to HF token issue.\")\n    def test_push_to_hub(self):\n        for model_name in self.all_model_names:\n            model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)\n            if \"sharded\" in model_name:\n                model.push_to_hub(model_name + \"-ppo\", use_auth_token=True, max_shard_size=\"1MB\")\n            else:\n                model.push_to_hub(model_name + \"-ppo\", use_auth_token=True)\n\n            model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(model_name + \"-ppo\")\n            # check all keys\n            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()\n\n            for name, param in model.state_dict().items():\n                assert torch.allclose(\n                    param, model_from_pretrained.state_dict()[name]\n                ), f\"Parameter {name} is not the same after push_to_hub and from_pretrained\"\n\n\nclass Seq2SeqValueHeadModelTester(VHeadModelTester, unittest.TestCase):\n    \"\"\"\n    Testing suite for v-head models.\n    \"\"\"\n\n    all_model_names = ALL_SEQ2SEQ_MODELS\n    trl_model_class = AutoModelForSeq2SeqLMWithValueHead\n    transformers_model_class = AutoModelForSeq2SeqLM\n\n    def tearDown(self):\n        # free memory\n        gc.collect()\n\n    def test_inference(self):\n        r\"\"\"\n        Test if the model can be used for inference and outputs 3 values\n        - logits, loss, and value states\n        \"\"\"\n        EXPECTED_OUTPUT_SIZE = 3\n\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n            outputs = model(input_ids, decoder_input_ids=decoder_input_ids)\n\n            # Check if the outputs are of the right size - here\n            # we always output 3 values - logits, loss, and value states\n            assert len(outputs) == EXPECTED_OUTPUT_SIZE\n\n    def test_dropout_config(self):\n        r\"\"\"\n        Test if we instantiate a model by adding `summary_drop_prob` to the config\n        it will be added to the v_head\n        \"\"\"\n        for model_name in self.all_model_names:\n            pretrained_model = self.transformers_model_class.from_pretrained(model_name)\n            pretrained_model.config.summary_dropout_prob = 0.5\n            model = self.trl_model_class.from_pretrained(pretrained_model)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob\n\n    def test_dropout_kwargs(self):\n        r\"\"\"\n        Test if we instantiate a model by adding `summary_drop_prob` to the config\n        it will be added to the v_head\n        \"\"\"\n        for model_name in self.all_model_names:\n            v_head_kwargs = {\"summary_dropout_prob\": 0.5}\n\n            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == 0.5\n\n            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)\n\n            # Check if v head of the model has the same dropout as the config\n            assert model.v_head.dropout.p == 0.5\n\n    def test_generate(self):\n        r\"\"\"\n        Test if `generate` works for every model\n        \"\"\"\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n            # Just check if the generation works\n            _ = model.generate(input_ids, decoder_input_ids=decoder_input_ids)\n\n    def test_raise_error_not_causallm(self):\n        # Test with a model without a LM head\n        model_id = \"trl-internal-testing/tiny-random-T5Model\"\n        # This should raise a ValueError\n        with pytest.raises(ValueError):\n            pretrained_model = AutoModel.from_pretrained(model_id)\n            _ = self.trl_model_class.from_pretrained(pretrained_model)\n\n    @unittest.skip(\"This test needs to be run manually due to HF token issue.\")\n    def test_push_to_hub(self):\n        for model_name in self.all_model_names:\n            model = self.trl_model_class.from_pretrained(model_name)\n            if \"sharded\" in model_name:\n                model.push_to_hub(model_name + \"-ppo\", use_auth_token=True, max_shard_size=\"1MB\")\n            else:\n                model.push_to_hub(model_name + \"-ppo\", use_auth_token=True)\n\n            model_from_pretrained = self.trl_model_class.from_pretrained(model_name + \"-ppo\")\n            # check all keys\n            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()\n\n            for name, param in model.state_dict().items():\n                assert torch.allclose(\n                    param, model_from_pretrained.state_dict()[name]\n                ), f\"Parameter {name} is not the same after push_to_hub and from_pretrained\"\n\n    def test_transformers_bf16_kwargs(self):\n        r\"\"\"\n        Test if the transformers kwargs are correctly passed\n        Here we check that loading a model in half precision works as expected, i.e. the weights of\n        the `pretrained_model` attribute is loaded in half precision and you can run a dummy\n        forward pass without any issue.\n        \"\"\"\n        for model_name in self.all_model_names:\n            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16)\n\n            lm_head_namings = self.trl_model_class.lm_head_namings\n\n            if model_name == \"trl-internal-testing/tiny-random-FSMTForConditionalGeneration\":\n                # skip the test for FSMT as it does not support mixed-prec\n                continue\n\n            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings)\n\n            for lm_head_naming in lm_head_namings:\n                if hasattr(trl_model.pretrained_model, lm_head_naming):\n                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16\n\n            dummy_input = torch.LongTensor([[0, 1, 0, 1]])\n\n            # check dummy forward pass works in half precision\n            _ = trl_model(input_ids=dummy_input, decoder_input_ids=dummy_input)\n\n\nclass ReferenceModelTest(unittest.TestCase):\n    def setUp(self):\n        self.model = AutoModelForCausalLMWithValueHead.from_pretrained(\n            \"trl-internal-testing/tiny-random-GPT2LMHeadModel\"\n        )\n        self.test_input = torch.tensor([[0, 1, 2, 3]])\n        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1)\n        self.layer_format = \"pretrained_model.transformer.h.{layer}.attn.c_attn.weight\"\n\n    def test_independent_reference(self):\n        layer_0 = self.layer_format.format(layer=0)\n        layer_5 = self.layer_format.format(layer=4)\n\n        ref_model = create_reference_model(self.model)\n\n        first_layer_before = self.model.get_parameter(layer_0).data.clone()\n        last_layer_before = self.model.get_parameter(layer_5).data.clone()\n\n        first_ref_layer_before = ref_model.get_parameter(layer_0).data.clone()\n        last_ref_layer_before = ref_model.get_parameter(layer_5).data.clone()\n\n        output = self.model(input_ids=self.test_input, labels=self.test_input)\n        output[1].backward()\n        self.optimizer.step()\n\n        first_layer_after = self.model.get_parameter(layer_0).data.clone()\n        last_layer_after = self.model.get_parameter(layer_5).data.clone()\n\n        first_ref_layer_after = ref_model.get_parameter(layer_0).data.clone()\n        last_ref_layer_after = ref_model.get_parameter(layer_5).data.clone()\n\n        # before optimization ref and model are identical\n        assert (first_layer_before == first_ref_layer_before).all()\n        assert (last_layer_before == last_ref_layer_before).all()\n        # ref model stays identical after optimization\n        assert (first_ref_layer_before == first_ref_layer_after).all()\n        assert (last_ref_layer_before == last_ref_layer_after).all()\n        # optimized model changes\n        assert not (first_layer_before == first_layer_after).all()\n        assert not (last_layer_before == last_layer_after).all()\n\n    def test_shared_layers(self):\n        layer_0 = self.layer_format.format(layer=0)\n        layer_1 = self.layer_format.format(layer=1)\n\n        ref_model = create_reference_model(self.model, num_shared_layers=1)\n\n        first_layer_before = self.model.get_parameter(layer_0).data.clone()\n        second_layer_before = self.model.get_parameter(layer_1).data.clone()\n\n        first_ref_layer_before = ref_model.get_parameter(layer_0).data.clone()\n        second_ref_layer_before = ref_model.get_parameter(layer_1).data.clone()\n\n        output = self.model(input_ids=self.test_input, labels=self.test_input)\n        output[1].backward()\n        self.optimizer.step()\n\n        first_layer_after = self.model.get_parameter(layer_0).data.clone()\n        second_layer_after = self.model.get_parameter(layer_1).data.clone()\n\n        first_ref_layer_after = ref_model.get_parameter(layer_0).data.clone()\n        second_ref_layer_after = ref_model.get_parameter(layer_1).data.clone()\n\n        # before optimization ref and model are identical\n        assert (first_layer_before == first_ref_layer_before).all()\n        assert (second_layer_before == second_ref_layer_before).all()\n        # ref model stays identical after optimization\n        assert (first_ref_layer_before == first_ref_layer_after).all()\n        assert (second_ref_layer_before == second_ref_layer_after).all()\n        # first layer of optimized model stays the same\n        assert (first_layer_before == first_layer_after).all()\n        # other layers in optimized model change\n        assert not (second_layer_before == second_layer_after).all()\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nimport torch\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer\nfrom transformers.testing_utils import require_peft\n\nfrom trl import KTOConfig, KTOTrainer\nfrom trl.trainer.kto_trainer import _get_kl_dataset, _process_tokens, _tokenize\n\nfrom .testing_utils import require_no_wandb\n\n\nclass KTOTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_ref_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", \"kto\", True, True],\n            [\"gpt2\", \"kto\", True, False],\n            [\"gpt2\", \"kto\", False, True],\n            [\"gpt2\", \"kto\", False, False],\n            [\"gpt2\", \"apo_zero_unpaired\", True, True],\n            [\"gpt2\", \"apo_zero_unpaired\", True, False],\n            [\"gpt2\", \"apo_zero_unpaired\", False, True],\n            [\"gpt2\", \"apo_zero_unpaired\", False, False],\n        ]\n    )\n    def test_kto_trainer(self, name, loss_type, pre_compute, eval_dataset):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                precompute_ref_log_probs=pre_compute,\n                loss_type=loss_type,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            if name == \"gpt2\":\n                model = self.model\n                ref_model = self.ref_model\n                tokenizer = self.tokenizer\n            elif name == \"t5\":\n                model = self.t5_model\n                ref_model = self.t5_ref_model\n                tokenizer = self.t5_tokenizer\n\n            trainer = KTOTrainer(\n                model=model,\n                ref_model=ref_model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"] if eval_dataset else None,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    self.assertFalse(torch.equal(param, new_param))\n\n    def test_kto_trainer_with_ref_model_is_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            with self.assertRaises(ValueError):\n                KTOTrainer(\n                    model=self.model,\n                    ref_model=self.model,  # ref_model can't be the same as model\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n    def test_tokenize_and_process_tokens(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = KTOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            train_dataset = dummy_dataset[\"train\"]\n            tokenized_dataset = train_dataset.map(\n                _tokenize,\n                fn_kwargs={\"tokenizer\": trainer.tokenizer},\n                batched=True,\n                batch_size=2,\n            )\n            self.assertListEqual(tokenized_dataset[\"prompt\"], train_dataset[\"prompt\"])\n            self.assertListEqual(tokenized_dataset[\"completion\"], train_dataset[\"completion\"])\n            self.assertListEqual(tokenized_dataset[\"label\"], train_dataset[\"label\"])\n            self.assertListEqual(tokenized_dataset[\"prompt_input_ids\"][0], [5377, 11141])\n            self.assertListEqual(tokenized_dataset[\"prompt_attention_mask\"][0], [1, 1])\n            self.assertListEqual(tokenized_dataset[\"answer_input_ids\"][0], [318, 1365, 621, 8253, 13])\n            self.assertListEqual(tokenized_dataset[\"answer_attention_mask\"][0], [1, 1, 1, 1, 1])\n\n            # Test corruption of (prompt, completion) pairs for KL dataset\n            for batch_size in [2, 3]:\n                tokenized_kl_dataset = tokenized_dataset.map(_get_kl_dataset, batched=True, batch_size=batch_size)\n\n                # Verify that the \"answer_input_ids\" have been modified, meaning the new \"answer_input_ids\" differ\n                # from the original ones. However, when the length of the dataset modulo batch_size equals 1,\n                # the last batch remains unaltered. This is a rare scenario that does not impact the training\n                # process, so we exclude it from testing by iterating only up to len - 1.\n                for i in range(len(tokenized_kl_dataset[\"answer_input_ids\"]) - 1):\n                    self.assertListEqual(\n                        tokenized_dataset[\"prompt_input_ids\"][i],\n                        tokenized_kl_dataset[\"prompt_input_ids\"][i],\n                    )\n                    self.assertListEqual(\n                        tokenized_dataset[\"prompt_attention_mask\"][i],\n                        tokenized_kl_dataset[\"prompt_attention_mask\"][i],\n                    )\n                    self.assertNotEqual(\n                        tokenized_dataset[\"answer_input_ids\"][i],\n                        tokenized_kl_dataset[\"answer_input_ids\"][i],\n                    )\n\n            fn_kwargs = {\n                \"prefix\": \"\",\n                \"is_encoder_decoder\": trainer.is_encoder_decoder,\n                \"tokenizer\": trainer.tokenizer,\n                \"max_length\": trainer.max_length,\n                \"truncation_mode\": trainer.truncation_mode,\n                \"label_pad_token_id\": trainer.label_pad_token_id,\n                \"max_prompt_length\": trainer.max_prompt_length,\n            }\n            processed_dataset = tokenized_dataset.map(_process_tokens, fn_kwargs=fn_kwargs, num_proc=2)\n            self.assertListEqual(processed_dataset[\"prompt\"], train_dataset[\"prompt\"])\n            self.assertListEqual(processed_dataset[\"completion\"], train_dataset[\"completion\"])\n            self.assertListEqual(processed_dataset[\"label\"], train_dataset[\"label\"])\n            self.assertListEqual(processed_dataset[\"prompt_input_ids\"][0], [50256, 5377, 11141])\n            self.assertListEqual(processed_dataset[\"prompt_attention_mask\"][0], [1, 1, 1])\n            self.assertListEqual(\n                processed_dataset[\"completion_input_ids\"][0], [50256, 5377, 11141, 318, 1365, 621, 8253, 13, 50256]\n            )\n            self.assertListEqual(processed_dataset[\"completion_attention_mask\"][0], [1, 1, 1, 1, 1, 1, 1, 1, 1])\n            self.assertListEqual(\n                processed_dataset[\"completion_labels\"][0], [-100, -100, -100, 318, 1365, 621, 8253, 13, 50256]\n            )\n\n    def test_kto_trainer_without_providing_ref_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = KTOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    self.assertFalse(torch.equal(param, new_param))\n\n    @require_peft\n    def test_kto_trainer_without_providing_ref_model_with_lora(self):\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = KTOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                if \"lora\" in n:\n                    new_param = trainer.model.get_parameter(n)\n                    # check the params have changed - ignore 0 biases\n                    if param.sum() != 0:\n                        self.assertFalse(torch.equal(param, new_param))\n\n    @require_no_wandb\n    def test_kto_trainer_generate_during_eval_no_wandb(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                generate_during_eval=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            with self.assertRaisesRegex(\n                ValueError,\n                expected_regex=\"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install with `pip install wandb` to resolve.\",\n            ):\n                KTOTrainer(\n                    model=self.model,\n                    ref_model=None,\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                )\n\n    @require_peft\n    def test_kto_lora_save(self):\n        from peft import LoraConfig, get_peft_model\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        model_peft = get_peft_model(model, lora_config)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            # kto train lora model with a lora config\n            trainer = KTOTrainer(\n                model=model_peft,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save peft adapter\n            trainer.save_model()\n\n            # assert that the model is loaded without giving OSError\n            try:\n                AutoModelForCausalLM.from_pretrained(tmp_dir)\n            except OSError:\n                self.fail(\"Loading the saved peft adapter failed\")\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport sys\nimport unittest\nfrom functools import partial\nfrom unittest.mock import patch\n\nimport pytest\nimport torch\nfrom transformers import AutoTokenizer\nfrom transformers.utils import import_utils\n\n\nclass DummyDataset(torch.utils.data.Dataset):\n    def __init__(self, query_data, response_data):\n        self.query_data = query_data\n        self.response_data = response_data\n\n    def __len__(self):\n        return len(self.query_data)\n\n    def __getitem__(self, idx):\n        return self.query_data[idx], self.response_data[idx]\n\n\nEXPECTED_STATS = [\n    \"objective/kl\",\n    \"objective/kl_dist\",\n    \"objective/logprobs\",\n    \"objective/ref_logprobs\",\n    \"objective/kl_coef\",\n    \"objective/entropy\",\n    \"ppo/mean_non_score_reward\",\n    \"ppo/loss/policy\",\n    \"ppo/loss/value\",\n    \"ppo/loss/total\",\n    \"ppo/policy/entropy\",\n    \"ppo/policy/approxkl\",\n    \"ppo/policy/policykl\",\n    \"ppo/policy/clipfrac\",\n    \"ppo/policy/advantages\",\n    \"ppo/policy/advantages_mean\",\n    \"ppo/policy/ratio\",\n    \"ppo/returns/mean\",\n    \"ppo/returns/var\",\n    \"ppo/val/vpred\",\n    \"ppo/val/error\",\n    \"ppo/val/clipfrac\",\n    \"ppo/val/mean\",\n    \"ppo/val/var\",\n    \"ppo/val/var_explained\",\n    \"time/ppo/forward_pass\",\n    \"time/ppo/compute_rewards\",\n    \"time/ppo/optimize_step\",\n    \"time/ppo/calc_stats\",\n    \"time/ppo/total\",\n    \"ppo/learning_rate\",\n]\n\n\nclass TestPeftDependancy(unittest.TestCase):\n    def setUp(self):\n        self.causal_lm_model_id = \"trl-internal-testing/tiny-random-GPTNeoXForCausalLM\"\n        self.seq_to_seq_model_id = \"trl-internal-testing/tiny-random-T5ForConditionalGeneration\"\n\n    def test_no_peft(self):\n        _peft_available = import_utils._peft_available\n        import_utils._peft_available = False  # required so that is_peft_available() returns False\n        with patch.dict(sys.modules, {\"peft\": None}):\n            from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead\n\n            # Check that loading a model with `peft` will raise an error\n            with pytest.raises(ModuleNotFoundError):\n                import peft  # noqa: F401\n\n            _trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.causal_lm_model_id)\n            _trl_seq2seq_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(self.seq_to_seq_model_id)\n        import_utils._peft_available = _peft_available\n\n    def test_imports_no_peft(self):\n        _peft_available = import_utils._peft_available\n        import_utils._peft_available = False  # required so that is_peft_available() returns False\n        with patch.dict(sys.modules, {\"peft\": None}):\n            from trl import (  # noqa: F401\n                AutoModelForCausalLMWithValueHead,\n                AutoModelForSeq2SeqLMWithValueHead,\n                PPOConfig,\n                PPOTrainer,\n                PreTrainedModelWrapper,\n            )\n        import_utils._peft_available = _peft_available\n\n    def test_ppo_trainer_no_peft(self):\n        _peft_available = import_utils._peft_available\n        import_utils._peft_available = False  # required so that is_peft_available() returns False\n        with patch.dict(sys.modules, {\"peft\": None}):\n            from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer\n\n            ppo_model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n\n            trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_model_id)\n            tokenizer = AutoTokenizer.from_pretrained(ppo_model_id)\n            tokenizer.pad_token_id = tokenizer.eos_token_id\n\n            ppo_config = PPOConfig(batch_size=2, mini_batch_size=1, log_with=None)\n\n            dummy_dataset = DummyDataset(\n                [torch.LongTensor([0, 1, 0, 1, 0, 1]), torch.LongTensor([0, 1, 0, 1, 0, 1])],\n                [torch.LongTensor([1, 0, 1, 0, 1, 0]), torch.LongTensor([0, 1, 0, 1, 0, 1])],\n            )\n\n            ppo_trainer = PPOTrainer(\n                config=ppo_config,\n                model=trl_model,\n                ref_model=None,\n                tokenizer=tokenizer,\n                dataset=dummy_dataset,\n            )\n            ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n            dummy_dataloader = ppo_trainer.dataloader\n\n            for query_tensor, response_tensor in dummy_dataloader:\n                # define a reward for response\n                # (this could be any reward such as human feedback or output from another model)\n                reward = [torch.tensor(1.0), torch.tensor(0.0)]\n                # train model\n                train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n                break\n\n            # check gradients are not None\n            for _, param in trl_model.named_parameters():\n                if param.requires_grad:\n                    assert param.grad is not None\n\n            # check expected stats\n            for stat in EXPECTED_STATS:\n                assert stat in train_stats\n        import_utils._peft_available = _peft_available\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport copy\nimport os\nimport tempfile\nimport unittest\n\nimport numpy as np\nimport pytest\nimport torch\nfrom datasets import Dataset, Image, Sequence, load_dataset\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoProcessor,\n    AutoTokenizer,\n    LlavaForConditionalGeneration,\n    TrainingArguments,\n    is_vision_available,\n)\nfrom transformers.testing_utils import require_peft, require_vision\nfrom transformers.utils import is_peft_available\n\nfrom trl import SFTConfig, SFTTrainer\nfrom trl.trainer import ConstantLengthDataset, DataCollatorForCompletionOnlyLM\n\n\ndef formatting_prompts_func(example):\n    text = f\"### Question: {example['question']}\\n ### Answer: {example['answer']}\"\n    return text\n\n\ndef formatting_prompts_func_batched(example):\n    output_text = []\n    for i, question in enumerate(example[\"question\"]):\n        text = f\"### Question: {question}\\n ### Answer: {example['answer'][i]}\"\n        output_text.append(text)\n    return output_text\n\n\nif is_peft_available():\n    from peft import LoraConfig, PeftModel\n\nif is_vision_available():\n    from PIL import Image as PILImage\n\n\nclass SFTTrainerTester(unittest.TestCase):\n    r\"\"\" \"\"\"\n\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n        self.dummy_dataset = Dataset.from_dict(\n            {\n                \"question\": [\n                    \"Does llamas know how to code?\",\n                    \"Does llamas know how to fly?\",\n                    \"Does llamas know how to talk?\",\n                    \"Does llamas know how to code?\",\n                    \"Does llamas know how to fly?\",\n                    \"Does llamas know how to talk?\",\n                    \"Does llamas know how to swim?\",\n                ],\n                \"answer\": [\n                    \"Yes, llamas are very good at coding.\",\n                    \"No, llamas can't fly.\",\n                    \"Yes, llamas are very good at talking.\",\n                    \"Yes, llamas are very good at coding.\",\n                    \"No, llamas can't fly.\",\n                    \"Yes, llamas are very good at talking.\",\n                    \"No, llamas can't swim.\",\n                ],\n                \"text\": [\n                    \"### Question: Does llamas know how to code?\\n ### Answer: Yes, llamas are very good at coding.\",\n                    \"### Question: Does llamas know how to fly?\\n ### Answer: No, llamas can't fly.\",\n                    \"### Question: Does llamas know how to talk?\\n ### Answer: Yes, llamas are very good at talking.\",\n                    \"### Question: Does llamas know how to code?\\n ### Answer: Yes, llamas are very good at coding.\",\n                    \"### Question: Does llamas know how to fly?\\n ### Answer: No, llamas can't fly.\",\n                    \"### Question: Does llamas know how to talk?\\n ### Answer: Yes, llamas are very good at talking.\",\n                    \"### Question: Does llamas know how to swim?\\n ### Answer: No, llamas can't swim.\",\n                ],\n            }\n        )\n        self.conversational_lm_dataset = load_dataset(\"trl-internal-testing/zen\", \"conversational_language_modeling\")\n        self.standard_prompt_completion_dataset = load_dataset(\n            \"trl-internal-testing/zen\", \"standard_prompt_completion\"\n        )\n\n        if is_vision_available():\n            self.dummy_vsft_instruction_dataset = Dataset.from_dict(\n                {\n                    \"messages\": [\n                        [\n                            {\n                                \"role\": \"user\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"What is in this image?\"}, {\"type\": \"image\"}],\n                            },\n                            {\n                                \"role\": \"assistant\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"It is random noise.\"}],\n                            },\n                            {\n                                \"role\": \"user\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"Oh ye, you are right, what is 1+1\"}],\n                            },\n                            {\n                                \"role\": \"assistant\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"2\"}],\n                            },\n                        ],\n                        [\n                            {\n                                \"role\": \"user\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"What is in this image?\"}, {\"type\": \"image\"}],\n                            },\n                            {\n                                \"role\": \"assistant\",\n                                \"content\": [{\"type\": \"text\", \"text\": \"It is random noise.\"}],\n                            },\n                        ],\n                    ],\n                    \"images\": [\n                        [PILImage.fromarray((np.random.rand(40, 50, 3) * 255).astype(\"uint8\")).convert(\"RGBA\")],\n                        [PILImage.fromarray((np.random.rand(50, 60, 3) * 255).astype(\"uint8\")).convert(\"RGBA\")],\n                    ],\n                }\n            )\n            self.dummy_vsft_instruction_dataset.cast_column(\"images\", Sequence(Image()))\n            self.dummy_vsft_instruction_dataset = self.dummy_vsft_instruction_dataset.cast_column(\n                \"images\", Sequence(Image())\n            )\n\n        self.train_dataset = ConstantLengthDataset(\n            self.tokenizer,\n            self.dummy_dataset,\n            dataset_text_field=None,\n            formatting_func=formatting_prompts_func,\n            seq_length=16,\n            num_of_sequences=16,\n        )\n\n        self.eval_dataset = ConstantLengthDataset(\n            self.tokenizer,\n            self.dummy_dataset,\n            dataset_text_field=None,\n            formatting_func=formatting_prompts_func,\n            seq_length=16,\n            num_of_sequences=16,\n        )\n\n    def test_constant_length_dataset(self):\n        formatted_dataset = ConstantLengthDataset(\n            self.tokenizer,\n            self.dummy_dataset,\n            dataset_text_field=None,\n            formatting_func=formatting_prompts_func,\n        )\n\n        assert len(formatted_dataset) == len(self.dummy_dataset)\n        assert len(formatted_dataset) > 0\n\n        for example in formatted_dataset:\n            assert \"input_ids\" in example\n            assert \"labels\" in example\n\n            assert len(example[\"input_ids\"]) == formatted_dataset.seq_length\n            assert len(example[\"labels\"]) == formatted_dataset.seq_length\n\n            decoded_text = self.tokenizer.decode(example[\"input_ids\"])\n            assert (\"Question\" in decoded_text) and (\"Answer\" in decoded_text)\n\n    def test_sft_trainer_backward_compatibility(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                hub_token=\"not_a_real_token\",\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                formatting_func=formatting_prompts_func,\n            )\n\n            assert trainer.args.hub_token == training_args.hub_token\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n    def test_sft_trainer(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n    def test_sft_trainer_uncorrect_data(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            # Shouldn't work as `dataset_text_field` is missing from the arguments\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n            with pytest.raises(ValueError):\n                _ = SFTTrainer(\n                    model=self.model,\n                    args=training_args,\n                    train_dataset=self.dummy_dataset,\n                )\n\n            # Shoud work as SFTTrainer natively supports conversational lm dataset\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=32,  # make sure there is at least 1 packed sequence\n                num_of_sequences=32,\n                packing=True,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.conversational_lm_dataset[\"train\"],\n            )\n\n            # Same, but without packing\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=False,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.conversational_lm_dataset[\"train\"],\n            )\n\n            # Same, but with packing with `max_seq_length`\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=16,  # make sure there is at least 1 packed sequence\n                packing=True,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.standard_prompt_completion_dataset[\"train\"],\n            )\n\n            # Same but with prompt completion dataset\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=False,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.standard_prompt_completion_dataset[\"train\"],\n            )\n\n            # Should work as dummy dataset are supported with a formatting function\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=32,  # make sure there is at least 1 packed sequence\n                packing=True,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n                formatting_func=formatting_prompts_func,\n            )\n\n            # This should not work because not enough data for one sample\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=1024,  # make sure there is NOT at least 1 packed sequence\n                packing=True,\n                report_to=\"none\",\n            )\n            with pytest.raises(ValueError):\n                _ = SFTTrainer(\n                    model=self.model,\n                    args=training_args,\n                    train_dataset=self.dummy_dataset,\n                    formatting_func=formatting_prompts_func,\n                )\n\n            # This should not work as well\n            with pytest.raises(ValueError):\n                training_args = SFTConfig(\n                    output_dir=tmp_dir,\n                    dataloader_drop_last=True,\n                    eval_strategy=\"steps\",\n                    max_steps=2,\n                    eval_steps=1,\n                    save_steps=1,\n                    per_device_train_batch_size=2,\n                    packing=False,\n                    report_to=\"none\",\n                )\n                _ = SFTTrainer(\n                    model=self.model,\n                    args=training_args,\n                    train_dataset=self.dummy_dataset,\n                    formatting_func=formatting_prompts_func,\n                )\n\n            # but this should work\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=False,\n                report_to=\"none\",\n            )\n            _ = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n                formatting_func=formatting_prompts_func_batched,\n            )\n\n    def test_sft_trainer_with_model_num_train_epochs(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                num_train_epochs=2,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                num_train_epochs=2,\n                per_device_train_batch_size=2,\n                dataset_text_field=\"text\",\n                max_seq_length=16,\n                num_of_sequences=16,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                num_train_epochs=2,\n                per_device_train_batch_size=2,\n                dataset_text_field=\"text\",\n                max_seq_length=16,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-1\")\n\n    def test_sft_trainer_with_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                dataset_text_field=\"text\",\n                max_seq_length=16,\n                num_of_sequences=16,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        # with formatting_func + packed\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=16,\n                num_of_sequences=16,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n                formatting_func=formatting_prompts_func,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        # with formatting_func + packed\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                max_seq_length=16,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n                formatting_func=formatting_prompts_func_batched,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                dataset_text_field=\"text\",\n                max_seq_length=16,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-1\")\n\n    def test_sft_trainer_with_multiple_eval_datasets(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=1,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset={\n                    \"data1\": self.eval_dataset,\n                    \"data2\": self.eval_dataset,\n                },\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_data1_loss\"] is not None\n            assert trainer.state.log_history[1][\"eval_data2_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-1\")\n\n    def test_data_collator_completion_lm(self):\n        response_template = \"### Response:\\n\"\n        data_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=self.tokenizer, mlm=False)\n\n        text = \"\"\"\\n\\n### Instructions:\\nHello all this should be masked\\n\\n### Response:\\nI have not been masked correctly.\"\"\"\n        encoded_text = self.tokenizer(text)\n\n        examples = [encoded_text]\n\n        batch = data_collator(examples)\n        labels = batch[\"labels\"]\n        last_pad_idx = np.where(labels == -100)[1][-1]\n        result_text = self.tokenizer.decode(batch[\"input_ids\"][0, last_pad_idx + 1 :])\n        assert result_text == \"I have not been masked correctly.\"\n\n    def test_data_collator_completion_lm_with_multiple_text(self):\n        tokenizer = copy.deepcopy(self.tokenizer)\n        tokenizer.padding_side = \"left\"\n\n        response_template = \"### Response:\\n\"\n        data_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer, mlm=False)\n\n        text1 = \"\"\"\\n\\n### Instructions:\\nHello all this should be masked\\n\\n### Response:\\nI have not been masked correctly.\"\"\"\n        text2 = \"\"\"\\n\\n### Instructions:\\nThis is another longer text that should also be masked. This text is significantly longer than the previous one.\\n\\n### Response:\\nI have not been masked correctly.\"\"\"\n\n        encoded_text1 = tokenizer(text1)\n        encoded_text2 = tokenizer(text2)\n\n        examples = [encoded_text1, encoded_text2]\n\n        batch = data_collator(examples)\n\n        for i in range(2):\n            labels = batch[\"labels\"][i]\n            last_pad_idx = np.where(labels == -100)[0][-1]\n            result_text = tokenizer.decode(batch[\"input_ids\"][i, last_pad_idx + 1 :])\n            assert result_text == \"I have not been masked correctly.\"\n\n    def test_data_collator_chat_completion_lm(self):\n        instruction_template = \"### Human:\"\n        assistant_template = \"### Assistant:\"\n        data_collator = DataCollatorForCompletionOnlyLM(\n            response_template=assistant_template,\n            instruction_template=instruction_template,\n            tokenizer=self.tokenizer,\n            mlm=False,\n        )\n\n        text = \"\"\"### Human: Hello all this should be masked.### Assistant: I should not be masked.### Human: All this should be masked too.### Assistant: I should not be masked too.\"\"\"\n        encoded_text = self.tokenizer(text)\n\n        examples = [encoded_text]\n\n        batch = data_collator(examples)\n        labels = batch[\"labels\"]\n        non_masked_tokens = batch[\"input_ids\"][labels != -100]\n        result_text = self.tokenizer.decode(non_masked_tokens)\n        assert result_text == \" I should not be masked. I should not be masked too.\"\n\n    def test_data_collator_chat_completion_lm_with_multiple_text(self):\n        tokenizer = copy.deepcopy(self.tokenizer)\n        tokenizer.padding_side = \"left\"\n\n        instruction_template = \"### Human:\"\n        assistant_template = \"### Assistant:\"\n        data_collator = DataCollatorForCompletionOnlyLM(\n            response_template=assistant_template,\n            instruction_template=instruction_template,\n            tokenizer=tokenizer,\n            mlm=False,\n        )\n\n        text1 = \"\"\"### Human: Hello all this should be masked.### Assistant: I should not be masked.\"\"\"\n        text2 = \"\"\"### Human: Hello all this should be masked.### Assistant: I should not be masked.### Human: All this should be masked too.### Assistant: I should not be masked too.\"\"\"\n        encoded_text1 = tokenizer(text1)\n        encoded_text2 = tokenizer(text2)\n\n        examples = [encoded_text1, encoded_text2]\n\n        batch = data_collator(examples)\n        labels = batch[\"labels\"]\n        input_ids = batch[\"input_ids\"]\n\n        non_masked_tokens1 = input_ids[0][labels[0] != -100]\n        result_text1 = tokenizer.decode(non_masked_tokens1)\n        assert result_text1 == \" I should not be masked.\"\n\n        non_masked_tokens2 = input_ids[1][labels[1] != -100]\n        result_text2 = tokenizer.decode(non_masked_tokens2)\n        assert result_text2 == \" I should not be masked. I should not be masked too.\"\n\n    def test_sft_trainer_infinite_with_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=5,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                packing=True,\n                max_seq_length=500,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            assert trainer.train_dataset.infinite\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            # make sure the trainer did 5 steps\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-5\")\n\n    def test_sft_trainer_infinite_with_model_epochs(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                num_train_epochs=1,\n                per_device_train_batch_size=2,\n                save_strategy=\"epoch\",\n                packing=True,\n                max_seq_length=500,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            assert not trainer.train_dataset.infinite\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            # make sure the trainer did 5 steps\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-4\")\n\n    def test_sft_trainer_with_model_neftune(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=2,\n                eval_steps=1,\n                save_steps=1,\n                per_device_train_batch_size=2,\n                neftune_noise_alpha=5,\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.model = trainer._activate_neftune(trainer.model)\n\n            device = trainer.model.get_input_embeddings().weight.device\n            trainer.model.train()\n\n            torch.random.manual_seed(42)\n            embeds_neftune = trainer.model.get_input_embeddings()(torch.LongTensor([[1, 0, 1]]).to(device))\n\n            torch.random.manual_seed(24)\n            embeds_neftune_2 = trainer.model.get_input_embeddings()(torch.LongTensor([[1, 0, 1]]).to(device))\n\n            assert not torch.allclose(embeds_neftune, embeds_neftune_2)\n            assert len(trainer.model.get_input_embeddings()._forward_hooks) > 0\n\n            trainer.neftune_hook_handle.remove()\n\n            trainer.train()\n\n            # Make sure forward pass works fine\n            _ = trainer.model(torch.LongTensor([[1, 0, 1]]).to(device))\n            assert len(trainer.model.get_input_embeddings()._forward_hooks) == 0\n\n    @require_peft\n    def test_peft_sft_trainer_str(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n\n            training_args = SFTConfig(\n                packing=True,\n                output_dir=tmp_dir,\n                report_to=\"none\",\n            )\n\n            _ = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=peft_config,\n            )\n\n    @require_peft\n    def test_peft_sft_trainer(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"adapter_model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"adapter_config.json\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"model.safetensors\" not in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n    @require_peft\n    def test_peft_sft_trainer_gc(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"adapter_model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"adapter_config.json\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"model.safetensors\" not in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n    @require_peft\n    def test_peft_sft_trainer_neftune(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                neftune_noise_alpha=5,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=peft_config,\n            )\n\n            trainer.model = trainer._activate_neftune(trainer.model)\n\n            assert isinstance(trainer.model, PeftModel)\n\n            device = trainer.model.get_input_embeddings().weight.device\n            trainer.model.train()\n\n            torch.random.manual_seed(42)\n            embeds_neftune = trainer.model.get_input_embeddings()(torch.LongTensor([[1, 0, 1]]).to(device))\n\n            torch.random.manual_seed(24)\n            embeds_neftune_2 = trainer.model.get_input_embeddings()(torch.LongTensor([[1, 0, 1]]).to(device))\n\n            assert not torch.allclose(embeds_neftune, embeds_neftune_2)\n            assert len(trainer.model.get_input_embeddings()._forward_hooks) > 0\n\n            trainer.neftune_hook_handle.remove()\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"adapter_model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"adapter_config.json\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n            assert \"model.safetensors\" not in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n            # Make sure forward pass works fine to check if embeddings forward is not broken.\n            _ = trainer.model(torch.LongTensor([[1, 0, 1]]).to(device))\n            assert len(trainer.model.get_input_embeddings()._forward_hooks) == 0\n\n    @require_peft\n    def test_peft_sft_trainer_tag(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=peft_config,\n            )\n\n            assert trainer.model.model_tags == trainer._tag_names\n\n    @require_peft\n    def test_sft_trainer_tag(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                packing=True,\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            assert trainer.model.model_tags == trainer._tag_names\n\n    def test_sft_trainer_only_train_packing(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                packing=True,\n                max_seq_length=16,  # make sure there is at least 1 packed sequence\n                eval_packing=False,\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.conversational_lm_dataset[\"train\"],\n                eval_dataset=self.conversational_lm_dataset[\"test\"],\n            )\n\n            assert len(trainer.train_dataset[\"input_ids\"]) == 16  # with the used dataset, we end up with 16 sequences\n            assert len(trainer.eval_dataset[\"input_ids\"]) == len(self.conversational_lm_dataset[\"test\"])\n\n    def test_sft_trainer_eval_packing(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                max_seq_length=16,  # make sure there is at least 1 packed sequence\n                packing=True,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.conversational_lm_dataset[\"train\"],\n                eval_dataset=self.conversational_lm_dataset[\"test\"],\n            )\n\n            assert len(trainer.train_dataset[\"input_ids\"]) == 16  # with the used dataset, we end up with 16 sequences\n            assert len(trainer.eval_dataset[\"input_ids\"]) == 1  # with the used dataset, we end up with 1 sequence\n\n    def test_sft_trainer_no_packing(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                max_seq_length=16,  # make sure there is at least 1 packed sequence\n                packing=False,\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.conversational_lm_dataset[\"train\"],\n                eval_dataset=self.conversational_lm_dataset[\"test\"],\n            )\n\n            assert len(trainer.train_dataset[\"input_ids\"]) == len(self.conversational_lm_dataset[\"train\"])\n            assert len(trainer.eval_dataset[\"input_ids\"]) == len(self.conversational_lm_dataset[\"test\"])\n\n    @require_vision\n    def test_sft_trainer_skip_prepare_dataset(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                remove_unused_columns=False,\n                dataset_text_field=\"text\",  # need a dummy field\n                dataset_kwargs={\"skip_prepare_dataset\": True},\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.dummy_vsft_instruction_dataset,\n                eval_dataset=self.dummy_vsft_instruction_dataset,\n            )\n            assert trainer.train_dataset.features == self.dummy_vsft_instruction_dataset.features\n            assert trainer.eval_dataset.features == self.dummy_vsft_instruction_dataset.features\n\n    def test_sft_trainer_skip_prepare_dataset_with_no_packing(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                gradient_checkpointing=True,\n                remove_unused_columns=False,\n                packing=False,\n                dataset_kwargs={\"skip_prepare_dataset\": True},\n                report_to=\"none\",\n            )\n\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.dummy_dataset,\n            )\n            assert trainer.train_dataset.features == self.dummy_dataset.features\n\n    @require_vision\n    def test_sft_trainer_llava(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                per_device_eval_batch_size=2,\n                remove_unused_columns=False,\n                dataset_text_field=\"\",  # need a dummy field\n                dataset_kwargs={\"skip_prepare_dataset\": True},\n                report_to=\"none\",\n            )\n            tiny_llava = LlavaForConditionalGeneration.from_pretrained(\n                \"trl-internal-testing/tiny-random-LlavaForConditionalGeneration\"\n            )\n            processor = AutoProcessor.from_pretrained(\"trl-internal-testing/tiny-random-LlavaForConditionalGeneration\")\n\n            processor.chat_template = \"\"\"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. {% for message in messages %}{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% elif item['type'] == 'image' %}<image>{% endif %}{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}\"\"\"\n\n            def collate_fn(examples):\n                # Get the texts and images, and apply the chat template\n                texts = [processor.apply_chat_template(example[\"messages\"], tokenize=False) for example in examples]\n                images = [example[\"images\"][0] for example in examples]\n\n                # Tokenize the texts and process the images\n                batch = processor(texts, images, return_tensors=\"pt\", padding=True)\n\n                # The labels are the input_ids, and we mask the padding tokens in the loss computation\n                labels = batch[\"input_ids\"].clone()\n                labels[labels == processor.tokenizer.pad_token_id] = -100\n                batch[\"labels\"] = labels\n\n                return batch\n\n            trainer = SFTTrainer(\n                model=tiny_llava,\n                args=training_args,\n                data_collator=collate_fn,\n                train_dataset=self.dummy_vsft_instruction_dataset,\n                eval_dataset=self.dummy_vsft_instruction_dataset,\n            )\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n            assert trainer.state.log_history[0][\"eval_loss\"] is not None\n\n            assert \"model.safetensors\" in os.listdir(tmp_dir + \"/checkpoint-2\")\n\n    def test_sft_trainer_torch_dtype(self):\n        # See https://github.com/huggingface/trl/issues/1751\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                model_init_kwargs={\"torch_dtype\": torch.float16},\n                report_to=\"none\",\n            )\n            trainer = SFTTrainer(\n                model=self.model_id,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                formatting_func=formatting_prompts_func,\n            )\n            assert trainer.model.config.torch_dtype == torch.float16\n\n        # Now test when `torch_dtype` is provided but is wrong\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                model_init_kwargs={\"torch_dtype\": -1},\n                report_to=\"none\",\n            )\n            with pytest.raises(\n                ValueError,\n                match=\"Invalid `torch_dtype` passed to the SFTConfig. Expected a string with either `torch.dtype` or 'auto', but got -1.\",\n            ):\n                _ = SFTTrainer(\n                    model=self.model_id,\n                    args=training_args,\n                    train_dataset=self.train_dataset,\n                    eval_dataset=self.eval_dataset,\n                )\n\n\n# Copyright 2023 metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport unittest\n\nimport torch\nfrom transformers.utils import is_peft_available\n\nfrom trl import is_diffusers_available\n\nfrom .testing_utils import require_diffusers\n\n\nif is_diffusers_available() and is_peft_available():\n    from trl import DDPOConfig, DDPOTrainer, DefaultDDPOStableDiffusionPipeline\n\n\ndef scorer_function(images, prompts, metadata):\n    return torch.randn(1) * 3.0, {}\n\n\ndef prompt_function():\n    return (\"cabbages\", {})\n\n\n@require_diffusers\nclass DDPOTrainerTester(unittest.TestCase):\n    \"\"\"\n    Test the DDPOTrainer class.\n    \"\"\"\n\n    def setUp(self):\n        self.training_args = DDPOConfig(\n            num_epochs=2,\n            train_gradient_accumulation_steps=1,\n            per_prompt_stat_tracking_buffer_size=32,\n            sample_num_batches_per_epoch=2,\n            sample_batch_size=2,\n            mixed_precision=None,\n            save_freq=1000000,\n        )\n        pretrained_model = \"hf-internal-testing/tiny-stable-diffusion-torch\"\n        pretrained_revision = \"main\"\n\n        pipeline = DefaultDDPOStableDiffusionPipeline(\n            pretrained_model, pretrained_model_revision=pretrained_revision, use_lora=False\n        )\n\n        self.trainer = DDPOTrainer(self.training_args, scorer_function, prompt_function, pipeline)\n\n        return super().setUp()\n\n    def tearDown(self) -> None:\n        gc.collect()\n\n    def test_loss(self):\n        advantage = torch.tensor([-1.0])\n        clip_range = 0.0001\n        ratio = torch.tensor([1.0])\n        loss = self.trainer.loss(advantage, clip_range, ratio)\n        assert loss.item() == 1.0\n\n    def test_generate_samples(self):\n        samples, output_pairs = self.trainer._generate_samples(1, 2)\n        assert len(samples) == 1\n        assert len(output_pairs) == 1\n        assert len(output_pairs[0][0]) == 2\n\n    def test_calculate_loss(self):\n        samples, _ = self.trainer._generate_samples(1, 2)\n        sample = samples[0]\n\n        latents = sample[\"latents\"][0, 0].unsqueeze(0)\n        next_latents = sample[\"next_latents\"][0, 0].unsqueeze(0)\n        log_probs = sample[\"log_probs\"][0, 0].unsqueeze(0)\n        timesteps = sample[\"timesteps\"][0, 0].unsqueeze(0)\n        prompt_embeds = sample[\"prompt_embeds\"]\n        advantage = torch.tensor([1.0], device=prompt_embeds.device)\n\n        assert latents.shape == (1, 4, 64, 64)\n        assert next_latents.shape == (1, 4, 64, 64)\n        assert log_probs.shape == (1,)\n        assert timesteps.shape == (1,)\n        assert prompt_embeds.shape == (2, 77, 32)\n        loss, approx_kl, clipfrac = self.trainer.calculate_loss(\n            latents, timesteps, next_latents, log_probs, advantage, prompt_embeds\n        )\n\n        assert torch.isfinite(loss.cpu())\n\n\n@require_diffusers\nclass DDPOTrainerWithLoRATester(DDPOTrainerTester):\n    \"\"\"\n    Test the DDPOTrainer class.\n    \"\"\"\n\n    def setUp(self):\n        self.training_args = DDPOConfig(\n            num_epochs=2,\n            train_gradient_accumulation_steps=1,\n            per_prompt_stat_tracking_buffer_size=32,\n            sample_num_batches_per_epoch=2,\n            sample_batch_size=2,\n            mixed_precision=None,\n            save_freq=1000000,\n        )\n        pretrained_model = \"hf-internal-testing/tiny-stable-diffusion-torch\"\n        pretrained_revision = \"main\"\n\n        pipeline = DefaultDDPOStableDiffusionPipeline(\n            pretrained_model, pretrained_model_revision=pretrained_revision, use_lora=True\n        )\n\n        self.trainer = DDPOTrainer(self.training_args, scorer_function, prompt_function, pipeline)\n\n        return super().setUp()\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport tempfile\nimport unittest\n\nimport torch\nfrom transformers import AutoModelForCausalLM\nfrom transformers.testing_utils import require_bitsandbytes, require_peft\nfrom transformers.utils import is_peft_available\n\nfrom trl import AutoModelForCausalLMWithValueHead\n\n\nif is_peft_available():\n    from peft import LoraConfig, get_peft_model\n\n\n@require_peft\nclass PeftModelTester(unittest.TestCase):\n    def setUp(self):\n        self.causal_lm_model_id = \"trl-internal-testing/tiny-random-GPTNeoXForCausalLM\"\n        self.lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n    def test_create_peft_model(self):\n        r\"\"\"\n        Simply creates a peft model and checks that it can be loaded.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        _ = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n\n    def test_peft_requires_grad(self):\n        r\"\"\"\n        Check that the value head of the returned model has requires_grad=True.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n\n        # Check that the value head has requires_grad=True\n        assert model.v_head.summary.weight.requires_grad\n\n    def test_check_peft_model_nb_trainable_params(self):\n        r\"\"\"\n        Check that the number of trainable parameters is correct.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n\n        # Check that the number of trainable parameters is correct\n        nb_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 10273\n\n        # Check that the number of trainable param for the non-peft model is correct\n        non_peft_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.causal_lm_model_id)\n        nb_trainable_params = sum(p.numel() for p in non_peft_model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 99578\n\n    def test_create_peft_model_from_config(self):\n        r\"\"\"\n        Simply creates a peft model and checks that it can be loaded.\n        \"\"\"\n        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(\n            self.causal_lm_model_id, peft_config=self.lora_config\n        )\n        # Check that the number of trainable parameters is correct\n        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 10273\n\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(causal_lm_model, peft_config=self.lora_config)\n        # Check that the number of trainable parameters is correct\n        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 10273\n\n    @require_bitsandbytes\n    def test_create_bnb_peft_model_from_config(self):\n        r\"\"\"\n        Simply creates a peft model and checks that it can be loaded.\n        \"\"\"\n        from bitsandbytes.nn import Linear8bitLt\n\n        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(\n            self.causal_lm_model_id, peft_config=self.lora_config, load_in_8bit=True\n        )\n        # Check that the number of trainable parameters is correct\n        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 10273\n        assert trl_model.pretrained_model.model.gpt_neox.layers[0].mlp.dense_h_to_4h.__class__ == Linear8bitLt\n\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(\n            self.causal_lm_model_id, load_in_8bit=True, device_map=\"auto\"\n        )\n        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(causal_lm_model, peft_config=self.lora_config)\n        # Check that the number of trainable parameters is correct\n        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)\n        assert nb_trainable_params == 10273\n        assert trl_model.pretrained_model.model.gpt_neox.layers[0].mlp.dense_h_to_4h.__class__ == Linear8bitLt\n\n    def test_save_pretrained_peft(self):\n        r\"\"\"\n        Check that the model can be saved and loaded properly.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            model.save_pretrained(tmp_dir)\n\n            # check that the files `adapter_model.safetensors` and `adapter_config.json` are in the directory\n            assert os.path.isfile(\n                f\"{tmp_dir}/adapter_model.safetensors\"\n            ), f\"{tmp_dir}/adapter_model.safetensors does not exist\"\n            assert os.path.exists(f\"{tmp_dir}/adapter_config.json\"), f\"{tmp_dir}/adapter_config.json does not exist\"\n            # check also for `pytorch_model.bin` and make sure it only contains `v_head` weights\n            assert os.path.exists(f\"{tmp_dir}/pytorch_model.bin\"), f\"{tmp_dir}/pytorch_model.bin does not exist\"\n            maybe_v_head = torch.load(f\"{tmp_dir}/pytorch_model.bin\", weights_only=True)\n            # check that only keys that starts with `v_head` are in the dict\n            assert all(\n                k.startswith(\"v_head\") for k in maybe_v_head.keys()\n            ), f\"keys in {tmp_dir}/pytorch_model.bin do not start with `v_head`\"\n\n            model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(tmp_dir)\n\n            # check all the weights are the same\n            for p1, p2 in zip(model.named_parameters(), model_from_pretrained.named_parameters()):\n                assert torch.allclose(p1[1], p2[1]), f\"{p1[0]} != {p2[0]}\"\n\n    def test_load_pretrained_peft(self):\n        r\"\"\"\n        Check that the model saved with peft class interface can be loaded properly.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            pretrained_model.save_pretrained(tmp_dir)\n            model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(tmp_dir)\n\n            # check that the files `adapter_model.safetensors` and `adapter_config.json` are in the directory\n            assert os.path.isfile(\n                f\"{tmp_dir}/adapter_model.safetensors\"\n            ), f\"{tmp_dir}/adapter_model.safetensors does not exist\"\n            assert os.path.exists(f\"{tmp_dir}/adapter_config.json\"), f\"{tmp_dir}/adapter_config.json does not exist\"\n\n            # check all the weights are the same\n            for p1, p2 in zip(model.named_parameters(), model_from_pretrained.named_parameters()):\n                if p1[0] not in [\"v_head.summary.weight\", \"v_head.summary.bias\"]:\n                    assert torch.allclose(p1[1], p2[1]), f\"{p1[0]} != {p2[0]}\"\n\n    def test_continue_training_peft_model(self):\n        r\"\"\"\n        Load peft and checks that it can continue training.\n        \"\"\"\n        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)\n        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            pretrained_model.save_pretrained(tmp_dir)\n            # set is_trainable to True\n            model = AutoModelForCausalLMWithValueHead.from_pretrained(tmp_dir, is_trainable=True)\n            # Check that the number of trainable parameters is correct\n            nb_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n            assert nb_trainable_params == 10273\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport unittest\n\nimport torch\nfrom transformers import AutoModelForCausalLM, GenerationConfig\n\nfrom trl.models.modeling_base import GeometricMixtureWrapper, create_reference_model\n\n\nclass TestGeometricMixtureWrapper(unittest.TestCase):\n    def setUp(self):\n        self.model = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n        self.ref_model = create_reference_model(self.model)\n        self.generation_config = GenerationConfig.from_pretrained(\"gpt2\")\n        self.mixture_coef = 0.5\n        self.wrapper = GeometricMixtureWrapper(\n            self.model, self.ref_model, self.generation_config, mixture_coef=self.mixture_coef\n        )\n\n    def test_forward(self):\n        input_ids = torch.tensor([[1, 2, 3, 4, 5]])\n        attention_mask = torch.ones_like(input_ids)\n\n        output = self.wrapper(input_ids=input_ids, attention_mask=attention_mask)\n\n        self.assertIsNotNone(output)\n        self.assertTrue(hasattr(output, \"logits\"))\n        self.assertEqual(output.logits.shape, (1, 5, self.model.config.vocab_size))\n\n    def test_mixture_coefficient(self):\n        input_ids = torch.tensor([[1, 2, 3, 4, 5]])\n        attention_mask = torch.ones_like(input_ids)\n\n        with torch.no_grad():\n            model_output = self.model(input_ids=input_ids, attention_mask=attention_mask)\n            ref_model_output = self.ref_model(input_ids=input_ids, attention_mask=attention_mask)\n            wrapper_output = self.wrapper(input_ids=input_ids, attention_mask=attention_mask)\n\n        expected_logits = torch.nn.functional.log_softmax(\n            self.mixture_coef * ref_model_output.logits + (1 - self.mixture_coef) * model_output.logits, dim=-1\n        )\n\n        self.assertTrue(torch.allclose(wrapper_output.logits, expected_logits, atol=1e-5))\n\n    def test_prepare_inputs_for_generation(self):\n        input_ids = torch.tensor([[1, 2, 3, 4, 5]])\n        attention_mask = torch.ones_like(input_ids)\n\n        inputs = self.wrapper.prepare_inputs_for_generation(input_ids, attention_mask=attention_mask, use_cache=True)\n\n        self.assertIn(\"input_ids\", inputs)\n        self.assertIn(\"attention_mask\", inputs)\n        self.assertFalse(inputs.get(\"use_cache\", False))\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nfrom trl import HfPairwiseJudge, PairRMJudge, RandomPairwiseJudge, RandomRankJudge\n\n\nclass TestJudges(unittest.TestCase):\n    def _get_prompts_and_completions(self):\n        prompts = [\"The capital of France is\", \"The biggest planet in the solar system is\"]\n        completions = [[\"Paris\", \"Marseille\"], [\"Saturn\", \"Jupiter\"]]\n        return prompts, completions\n\n    def test_random_pairwise_judge(self):\n        judge = RandomPairwiseJudge()\n        prompts, completions = self._get_prompts_and_completions()\n        ranks = judge.judge(prompts=prompts, completions=completions)\n        self.assertEqual(len(ranks), 2)\n        self.assertTrue(all(isinstance(rank, int) for rank in ranks))\n\n    def test_random_rank_judge(self):\n        judge = RandomRankJudge()\n        prompts, completions = self._get_prompts_and_completions()\n        ranks = judge.judge(prompts=prompts, completions=completions)\n        self.assertEqual(len(ranks), 2)\n        self.assertTrue(all(isinstance(rank, list) for rank in ranks))\n        self.assertTrue(all(all(isinstance(rank, int) for rank in ranks) for ranks in ranks))\n\n    @unittest.skip(\"This test needs to be run manually since it requires a valid Hugging Face API key.\")\n    def test_hugging_face_judge(self):\n        judge = HfPairwiseJudge()\n        prompts, completions = self._get_prompts_and_completions()\n        ranks = judge.judge(prompts=prompts, completions=completions)\n        self.assertEqual(len(ranks), 2)\n        self.assertTrue(all(isinstance(rank, int) for rank in ranks))\n        self.assertEqual(ranks, [0, 1])\n\n    def test_pair_rm_judge(self):\n        judge = PairRMJudge()\n        prompts, completions = self._get_prompts_and_completions()\n        ranks = judge.judge(prompts=prompts, completions=completions)\n        self.assertEqual(len(ranks), 2)\n        self.assertTrue(all(isinstance(rank, int) for rank in ranks))\n        self.assertEqual(ranks, [0, 1])\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport unittest\n\nimport torch\n\nfrom trl.core import masked_mean, masked_var, masked_whiten, whiten\n\n\nclass CoreTester(unittest.TestCase):\n    \"\"\"\n    A wrapper class for testing core utils functions\n    \"\"\"\n\n    def setUp(self):\n        self.test_input = torch.Tensor([1, 2, 3, 4])\n        self.test_mask = torch.Tensor([0, 1, 1, 0])\n        self.test_input_unmasked = self.test_input[1:3]\n\n    def test_masked_mean(self):\n        assert torch.mean(self.test_input_unmasked) == masked_mean(self.test_input, self.test_mask)\n\n    def test_masked_var(self):\n        assert torch.var(self.test_input_unmasked) == masked_var(self.test_input, self.test_mask)\n\n    def test_masked_whiten(self):\n        whiten_unmasked = whiten(self.test_input_unmasked)\n        whiten_masked = masked_whiten(self.test_input, self.test_mask)[1:3]\n        diffs = (whiten_unmasked - whiten_masked).sum()\n        assert abs(diffs.item()) < 0.00001\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tempfile\nimport unittest\n\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, Trainer, TrainingArguments\n\nfrom trl import BasePairwiseJudge, WinRateCallback\n\n\nclass HalfPairwiseJudge(BasePairwiseJudge):\n    \"\"\"Naive pairwise judge that always returns [1, 0]\"\"\"\n\n    def judge(self, prompts, completions, shuffle_order=True):\n        # just check that the batch size is 2\n        assert len(prompts) == 2\n        return [1, 0]\n\n\nclass TrainerWithRefModel(Trainer):\n    # This is a dummy class to test the callback. Compared to the Trainer class, it only has an additional\n    # ref_model attribute\n    def __init__(self, model, ref_model, args, train_dataset, eval_dataset, tokenizer):\n        super().__init__(\n            model=model, args=args, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=tokenizer\n        )\n        self.ref_model = ref_model\n\n\nclass WinRateCallbackTester(unittest.TestCase):\n    def setUp(self):\n        self.model = AutoModelForCausalLM.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        self.ref_model = AutoModelForCausalLM.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        self.tokenizer = AutoTokenizer.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n        dataset[\"train\"] = dataset[\"train\"].select(range(8))\n        self.expected_winrates = [\n            {\"eval_win_rate\": 0.5, \"epoch\": 0.5, \"step\": 2},\n            {\"eval_win_rate\": 0.5, \"epoch\": 1.0, \"step\": 4},\n            {\"eval_win_rate\": 0.5, \"epoch\": 1.5, \"step\": 6},\n            {\"eval_win_rate\": 0.5, \"epoch\": 2.0, \"step\": 8},\n            {\"eval_win_rate\": 0.5, \"epoch\": 2.5, \"step\": 10},\n            {\"eval_win_rate\": 0.5, \"epoch\": 3.0, \"step\": 12},\n        ]\n\n        def tokenize_function(examples):\n            out = self.tokenizer(examples[\"prompt\"], padding=\"max_length\", max_length=16, truncation=True)\n            out[\"labels\"] = out[\"input_ids\"].copy()\n            return out\n\n        self.dataset = dataset.map(tokenize_function, batched=True)\n\n        self.generation_config = GenerationConfig(max_length=32)\n        self.judge = HalfPairwiseJudge()\n\n    def test_basic(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                eval_steps=2,  # evaluate every 2 steps\n                per_device_train_batch_size=2,  # 8 samples in total so 4 batches of 2 per epoch\n                per_device_eval_batch_size=2,\n                report_to=\"none\",\n            )\n            trainer = TrainerWithRefModel(\n                model=self.model,\n                ref_model=self.ref_model,\n                args=training_args,\n                train_dataset=self.dataset[\"train\"],\n                eval_dataset=self.dataset[\"test\"],\n                tokenizer=self.tokenizer,\n            )\n            win_rate_callback = WinRateCallback(\n                judge=self.judge, trainer=trainer, generation_config=self.generation_config\n            )\n            trainer.add_callback(win_rate_callback)\n            trainer.train()\n            winrate_history = [h for h in trainer.state.log_history if \"eval_win_rate\" in h]\n            self.assertListEqual(winrate_history, self.expected_winrates)\n\n    def test_without_ref_model(self):\n        # Same as before, but without the ref_model attribute. It should use the model attribute instead\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                eval_steps=2,  # evaluate every 2 steps\n                per_device_train_batch_size=2,  # 8 samples in total so 4 batches of 2 per epoch\n                per_device_eval_batch_size=2,\n                report_to=\"none\",\n            )\n            trainer = Trainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dataset[\"train\"],\n                eval_dataset=self.dataset[\"test\"],\n                tokenizer=self.tokenizer,\n            )\n            win_rate_callback = WinRateCallback(\n                judge=self.judge, trainer=trainer, generation_config=self.generation_config\n            )\n            trainer.add_callback(win_rate_callback)\n            trainer.train()\n            winrate_history = [h for h in trainer.state.log_history if \"eval_win_rate\" in h]\n            self.assertListEqual(winrate_history, self.expected_winrates)\n\n    def test_lora(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            peft_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n            self.model.add_adapter(peft_config)\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                eval_strategy=\"steps\",\n                eval_steps=2,  # evaluate every 2 steps\n                per_device_train_batch_size=2,  # 8 samples in total so 4 batches of 2 per epoch\n                per_device_eval_batch_size=2,\n                report_to=\"none\",\n            )\n            trainer = Trainer(\n                model=self.model,\n                args=training_args,\n                train_dataset=self.dataset[\"train\"],\n                eval_dataset=self.dataset[\"test\"],\n                tokenizer=self.tokenizer,\n            )\n            win_rate_callback = WinRateCallback(\n                judge=self.judge, trainer=trainer, generation_config=self.generation_config\n            )\n            trainer.add_callback(win_rate_callback)\n            trainer.train()\n            winrate_history = [h for h in trainer.state.log_history if \"eval_win_rate\" in h]\n            self.assertListEqual(winrate_history, self.expected_winrates)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tempfile\nimport unittest\n\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\n\nfrom trl import (\n    BCOConfig,\n    BCOTrainer,\n    CPOConfig,\n    CPOTrainer,\n    DPOConfig,\n    DPOTrainer,\n    KTOConfig,\n    KTOTrainer,\n    OnlineDPOConfig,\n    OnlineDPOTrainer,\n    ORPOConfig,\n    ORPOTrainer,\n    SFTConfig,\n    SFTTrainer,\n)\n\n\nclass TrainerArgTester(unittest.TestCase):\n    def test_bco(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                tmp_dir,\n                max_length=256,\n                max_prompt_length=64,\n                max_completion_length=64,\n                beta=0.5,\n                label_pad_token_id=-99,\n                padding_value=-99,\n                truncation_mode=\"keep_start\",\n                # generate_during_eval=True, # ignore this one, it requires wandb\n                is_encoder_decoder=True,\n                precompute_ref_log_probs=True,\n                model_init_kwargs={\"trust_remote_code\": True},\n                ref_model_init_kwargs={\"trust_remote_code\": True},\n                dataset_num_proc=4,\n                prompt_sample_size=512,\n                min_density_ratio=0.2,\n                max_density_ratio=20.0,\n            )\n            trainer = BCOTrainer(\n                model=\"gpt2\", ref_model=\"gpt2\", args=training_args, train_dataset=dataset, tokenizer=tokenizer\n            )\n            self.assertEqual(trainer.args.max_length, 256)\n            self.assertEqual(trainer.args.max_prompt_length, 64)\n            self.assertEqual(trainer.args.max_completion_length, 64)\n            self.assertEqual(trainer.args.beta, 0.5)\n            self.assertEqual(trainer.args.label_pad_token_id, -99)\n            self.assertEqual(trainer.args.padding_value, -99)\n            self.assertEqual(trainer.args.truncation_mode, \"keep_start\")\n            # self.assertEqual(trainer.args.generate_during_eval, True)\n            self.assertEqual(trainer.args.is_encoder_decoder, True)\n            self.assertEqual(trainer.args.precompute_ref_log_probs, True)\n            self.assertEqual(trainer.args.model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.ref_model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n            self.assertEqual(trainer.args.prompt_sample_size, 512)\n            self.assertEqual(trainer.args.min_density_ratio, 0.2)\n            self.assertEqual(trainer.args.max_density_ratio, 20.0)\n\n    def test_cpo(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = CPOConfig(\n                tmp_dir,\n                max_length=256,\n                max_prompt_length=64,\n                max_completion_length=64,\n                beta=0.5,\n                label_smoothing=0.5,\n                loss_type=\"hinge\",\n                disable_dropout=False,\n                cpo_alpha=0.5,\n                simpo_gamma=0.2,\n                label_pad_token_id=-99,\n                padding_value=-99,\n                truncation_mode=\"keep_start\",\n                # generate_during_eval=True, # ignore this one, it requires wandb\n                is_encoder_decoder=True,\n                model_init_kwargs={\"trust_remote_code\": True},\n                dataset_num_proc=4,\n            )\n            trainer = CPOTrainer(model=\"gpt2\", args=training_args, train_dataset=dataset, tokenizer=tokenizer)\n            self.assertEqual(trainer.args.max_length, 256)\n            self.assertEqual(trainer.args.max_prompt_length, 64)\n            self.assertEqual(trainer.args.max_completion_length, 64)\n            self.assertEqual(trainer.args.beta, 0.5)\n            self.assertEqual(trainer.args.label_smoothing, 0.5)\n            self.assertEqual(trainer.args.loss_type, \"hinge\")\n            self.assertEqual(trainer.args.disable_dropout, False)\n            self.assertEqual(trainer.args.cpo_alpha, 0.5)\n            self.assertEqual(trainer.args.simpo_gamma, 0.2)\n            self.assertEqual(trainer.args.label_pad_token_id, -99)\n            self.assertEqual(trainer.args.padding_value, -99)\n            self.assertEqual(trainer.args.truncation_mode, \"keep_start\")\n            # self.assertEqual(trainer.args.generate_during_eval, True)\n            self.assertEqual(trainer.args.is_encoder_decoder, True)\n            self.assertEqual(trainer.args.model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n\n    def test_dpo(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                tmp_dir,\n                beta=0.5,\n                label_smoothing=0.5,\n                loss_type=\"hinge\",\n                label_pad_token_id=-99,\n                padding_value=-99,\n                truncation_mode=\"keep_start\",\n                max_length=256,\n                max_prompt_length=64,\n                max_completion_length=64,\n                is_encoder_decoder=True,\n                disable_dropout=False,\n                # generate_during_eval=True, # ignore this one, it requires wandb\n                precompute_ref_log_probs=True,\n                dataset_num_proc=4,\n                model_init_kwargs={\"trust_remote_code\": True},\n                ref_model_init_kwargs={\"trust_remote_code\": True},\n                model_adapter_name=\"dummy_adapter\",\n                ref_adapter_name=\"dummy_adapter\",\n                reference_free=True,\n                force_use_ref_model=True,\n                f_divergence_type=\"js_divergence\",\n                f_alpha_divergence_coef=0.5,\n                sync_ref_model=True,\n                ref_model_mixup_alpha=0.5,\n                ref_model_sync_steps=32,\n                rpo_alpha=0.5,\n            )\n            trainer = DPOTrainer(\n                model=\"gpt2\", ref_model=\"gpt2\", args=training_args, train_dataset=dataset, tokenizer=tokenizer\n            )\n            self.assertEqual(trainer.args.beta, 0.5)\n            self.assertEqual(trainer.args.label_smoothing, 0.5)\n            self.assertEqual(trainer.args.loss_type, \"hinge\")\n            self.assertEqual(trainer.args.label_pad_token_id, -99)\n            self.assertEqual(trainer.args.padding_value, -99)\n            self.assertEqual(trainer.args.truncation_mode, \"keep_start\")\n            self.assertEqual(trainer.args.max_length, 256)\n            self.assertEqual(trainer.args.max_prompt_length, 64)\n            self.assertEqual(trainer.args.max_completion_length, 64)\n            self.assertEqual(trainer.args.is_encoder_decoder, True)\n            self.assertEqual(trainer.args.disable_dropout, False)\n            # self.assertEqual(trainer.args.generate_during_eval, True)\n            self.assertEqual(trainer.args.precompute_ref_log_probs, True)\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n            self.assertEqual(trainer.args.model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.ref_model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.model_adapter_name, \"dummy_adapter\")\n            self.assertEqual(trainer.args.ref_adapter_name, \"dummy_adapter\")\n            self.assertEqual(trainer.args.reference_free, True)\n            self.assertEqual(trainer.args.force_use_ref_model, True)\n            self.assertEqual(trainer.args.f_divergence_type, \"js_divergence\")\n            self.assertEqual(trainer.args.f_alpha_divergence_coef, 0.5)\n            self.assertEqual(trainer.args.sync_ref_model, True)\n            self.assertEqual(trainer.args.ref_model_mixup_alpha, 0.5)\n            self.assertEqual(trainer.args.ref_model_sync_steps, 32)\n            self.assertEqual(trainer.args.rpo_alpha, 0.5)\n\n    def test_kto(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = KTOConfig(\n                tmp_dir,\n                max_length=256,\n                max_prompt_length=64,\n                max_completion_length=64,\n                beta=0.5,\n                desirable_weight=0.5,\n                undesirable_weight=0.5,\n                label_pad_token_id=-99,\n                padding_value=-99,\n                truncation_mode=\"keep_start\",\n                # generate_during_eval=True, # ignore this one, it requires wandb\n                is_encoder_decoder=True,\n                precompute_ref_log_probs=True,\n                model_init_kwargs={\"trust_remote_code\": True},\n                ref_model_init_kwargs={\"trust_remote_code\": True},\n                dataset_num_proc=4,\n            )\n            trainer = KTOTrainer(\n                model=\"gpt2\", ref_model=\"gpt2\", args=training_args, train_dataset=dataset, tokenizer=tokenizer\n            )\n            self.assertEqual(trainer.args.max_length, 256)\n            self.assertEqual(trainer.args.max_prompt_length, 64)\n            self.assertEqual(trainer.args.max_completion_length, 64)\n            self.assertEqual(trainer.args.beta, 0.5)\n            self.assertEqual(trainer.args.desirable_weight, 0.5)\n            self.assertEqual(trainer.args.undesirable_weight, 0.5)\n            self.assertEqual(trainer.args.label_pad_token_id, -99)\n            self.assertEqual(trainer.args.padding_value, -99)\n            self.assertEqual(trainer.args.truncation_mode, \"keep_start\")\n            # self.assertEqual(trainer.args.generate_during_eval, True)\n            self.assertEqual(trainer.args.is_encoder_decoder, True)\n            self.assertEqual(trainer.args.precompute_ref_log_probs, True)\n            self.assertEqual(trainer.args.model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.ref_model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n\n    def test_online_dpo(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                tmp_dir,\n                max_new_tokens=42,\n                temperature=0.5,\n                missing_eos_penalty=0.33,\n                beta=0.6,\n                loss_type=\"hinge\",\n                dataset_num_proc=4,\n            )\n            model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/pythia-14m\")\n            ref_model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/pythia-14m\")\n            reward_model = AutoModelForSequenceClassification.from_pretrained(\"EleutherAI/pythia-14m\", num_labels=1)\n            trainer = OnlineDPOTrainer(\n                args=training_args,\n                tokenizer=tokenizer,\n                model=model,\n                ref_model=ref_model,\n                reward_model=reward_model,\n                train_dataset=dataset,\n            )\n            self.assertEqual(trainer.args.max_new_tokens, 42)\n            self.assertEqual(trainer.args.temperature, 0.5)\n            self.assertEqual(trainer.args.missing_eos_penalty, 0.33)\n            self.assertEqual(trainer.args.beta, 0.6)\n            self.assertEqual(trainer.args.loss_type, \"hinge\")\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n\n    def test_orpo(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = ORPOConfig(\n                tmp_dir,\n                max_length=256,\n                max_prompt_length=64,\n                max_completion_length=64,\n                beta=0.5,\n                disable_dropout=False,\n                label_pad_token_id=-99,\n                padding_value=-99,\n                truncation_mode=\"keep_start\",\n                # generate_during_eval=True, # ignore this one, it requires wandb\n                is_encoder_decoder=True,\n                model_init_kwargs={\"trust_remote_code\": True},\n                dataset_num_proc=4,\n            )\n\n            trainer = ORPOTrainer(model=\"gpt2\", args=training_args, train_dataset=dataset, tokenizer=tokenizer)\n            self.assertEqual(trainer.args.max_length, 256)\n            self.assertEqual(trainer.args.max_prompt_length, 64)\n            self.assertEqual(trainer.args.max_completion_length, 64)\n            self.assertEqual(trainer.args.beta, 0.5)\n            self.assertEqual(trainer.args.disable_dropout, False)\n            self.assertEqual(trainer.args.label_pad_token_id, -99)\n\n    def test_sft(self):\n        dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_language_modeling\", split=\"train\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                tmp_dir,\n                dataset_text_field=\"dummy_text_field\",\n                packing=True,\n                max_seq_length=256,\n                dataset_num_proc=4,\n                dataset_batch_size=512,\n                neftune_noise_alpha=0.1,\n                model_init_kwargs={\"trust_remote_code\": True},\n                dataset_kwargs={\"append_concat_token\": True, \"skip_prepare_dataset\": True},\n                eval_packing=True,\n                num_of_sequences=32,\n                chars_per_token=4.2,\n            )\n            trainer = SFTTrainer(\"gpt2\", args=training_args, train_dataset=dataset)\n            self.assertEqual(trainer.args.dataset_text_field, \"dummy_text_field\")\n            self.assertEqual(trainer.args.packing, True)\n            self.assertEqual(trainer.args.max_seq_length, 256)\n            self.assertEqual(trainer.args.dataset_num_proc, 4)\n            self.assertEqual(trainer.args.dataset_batch_size, 512)\n            self.assertEqual(trainer.args.neftune_noise_alpha, 0.1)\n            self.assertEqual(trainer.args.model_init_kwargs, {\"trust_remote_code\": True})\n            self.assertIn(\"append_concat_token\", trainer.args.dataset_kwargs)\n            self.assertEqual(trainer.args.dataset_kwargs[\"append_concat_token\"], True)\n            self.assertEqual(trainer.args.eval_packing, True)\n            self.assertEqual(trainer.args.num_of_sequences, 32)\n            self.assertEqual(trainer.args.chars_per_token, 4.2)\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport copy\nimport fnmatch\nimport gc\nimport re\nimport tempfile\nimport unittest\nfrom functools import partial\n\nimport pytest\nimport torch\nfrom huggingface_hub import HfApi\nfrom parameterized import parameterized\nfrom requests.exceptions import HTTPError\nfrom transformers import AutoTokenizer\nfrom transformers.testing_utils import require_peft, require_torch_multi_accelerator\n\nfrom trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed\nfrom trl.core import respond_to_batch\n\nfrom .testing_constants import CI_HUB_ENDPOINT, CI_HUB_USER\n\n\nEXPECTED_STATS = [\n    \"objective/kl\",\n    \"objective/kl_dist\",\n    \"objective/logprobs\",\n    \"objective/ref_logprobs\",\n    \"objective/kl_coef\",\n    \"objective/entropy\",\n    \"ppo/mean_non_score_reward\",\n    \"ppo/loss/policy\",\n    \"ppo/loss/value\",\n    \"ppo/loss/total\",\n    \"ppo/policy/entropy\",\n    \"ppo/policy/approxkl\",\n    \"ppo/policy/policykl\",\n    \"ppo/policy/clipfrac\",\n    \"ppo/policy/advantages\",\n    \"ppo/policy/advantages_mean\",\n    \"ppo/policy/ratio\",\n    \"ppo/returns/mean\",\n    \"ppo/returns/var\",\n    \"ppo/val/vpred\",\n    \"ppo/val/error\",\n    \"ppo/val/clipfrac\",\n    \"ppo/val/mean\",\n    \"ppo/val/var\",\n    \"ppo/val/var_explained\",\n    \"time/ppo/forward_pass\",\n    \"time/ppo/compute_rewards\",\n    \"time/ppo/optimize_step\",\n    \"time/ppo/calc_stats\",\n    \"time/ppo/total\",\n    \"ppo/learning_rate\",\n]\n\n\nclass DummyDataset(torch.utils.data.Dataset):\n    def __init__(self, query_data, response_data):\n        self.query_data = query_data\n        self.response_data = response_data\n\n    def __len__(self):\n        return len(self.query_data)\n\n    def __getitem__(self, idx):\n        return self.query_data[idx], self.response_data[idx]\n\n\ndef apply_mask(values, mask):\n    unmasked_values = []\n    for v, m in zip(values, mask):\n        if m == 1:\n            unmasked_values.append(v)\n    return torch.Tensor(unmasked_values)\n\n\ndef abs_diff_masked_tensors(tensor_1, tensor_2, mask_1, mask_2):\n    diffs = []\n    for l1, l2, m1, m2 in zip(tensor_1, tensor_2, mask_1, mask_2):\n        diff = apply_mask(l1, m1) - apply_mask(l2, m2)\n        diffs.append(diff.sum())\n    return abs(sum(diffs))\n\n\nclass PPOTrainerTester(unittest.TestCase):\n    \"\"\"\n    A wrapper class for testing PPOTrainer\n    \"\"\"\n\n    @classmethod\n    def setUpClass(cls):\n        cls._api = HfApi(endpoint=CI_HUB_ENDPOINT)\n\n    def setUp(self):\n        set_seed(42)\n\n        # model_id\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n\n        # get models and tokenizer\n        self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n        self.gpt2_ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n        self.gpt2_tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n\n        self.gpt2_tokenizer.pad_token = self.gpt2_tokenizer.eos_token\n\n        # get bloom as right padding examples:\n        model_id = \"trl-internal-testing/tiny-BloomForCausalLM-correct-vocab\"\n        self.bloom_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id)\n        self.bloom_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab\"\n        self.t5_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        # initialize trainer\n        self.ppo_config = PPOConfig(batch_size=2, mini_batch_size=1, log_with=None)\n\n    @classmethod\n    def tearDownClass(cls):\n        for model in [f\"{CI_HUB_USER}/test-ppo-trainer\"]:\n            try:\n                cls._api.delete_repo(repo_id=model)\n            except HTTPError:\n                pass\n\n    def tearDown(self):\n        # free memory\n        gc.collect()\n\n    def _init_dummy_dataset(self):\n        # encode a query\n        query_txt = \"This morning I went to the \"\n        query_tensor = self.gpt2_tokenizer.encode(query_txt, return_tensors=\"pt\")\n        assert query_tensor.shape == (1, 7)\n        # get model response\n        response_tensor = respond_to_batch(self.gpt2_model, query_tensor)\n        assert response_tensor.shape == (1, 20)\n\n        # create a dummy dataset\n        min_length = min(len(query_tensor[0]), len(response_tensor[0]))\n        dummy_dataset = DummyDataset(\n            [query_tensor[:, :min_length].squeeze(0) for _ in range(2)],\n            [response_tensor[:, :min_length].squeeze(0) for _ in range(2)],\n        )\n\n        return dummy_dataset\n\n    def test_drop_last_dataloader(self):\n        self.ppo_config = PPOConfig(batch_size=3, mini_batch_size=1, log_with=None)\n\n        dummy_dataset = self._init_dummy_dataset()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=self.gpt2_ref_model,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        dummy_dataloader = ppo_trainer.dataloader\n\n        assert len(dummy_dataloader) == 0\n\n    def test_ppo_step(self):\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=self.gpt2_ref_model,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        for param in ppo_trainer.model.parameters():\n            assert param.grad is not None\n\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n    def test_ppo_step_with_masks(self):\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=self.gpt2_ref_model,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n\n            response_mask = [torch.ones_like(r) for r in response_tensor]\n\n            # train model\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward, response_mask)\n            break\n\n        for param in ppo_trainer.model.parameters():\n            assert param.grad is not None\n\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n    def test_ppo_step_with_no_ref_sgd(self):\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n        optimizer = torch.optim.SGD(self.gpt2_model.parameters(), lr=0.01)\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            optimizer=optimizer,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n\n        assert isinstance(ppo_trainer.optimizer.optimizer, torch.optim.SGD)\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n        # Finally check stats\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n    def test_ppo_step_with_no_ref_sgd_lr_scheduler(self):\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n        optimizer = torch.optim.SGD(self.gpt2_model.parameters(), lr=0.01)\n        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            optimizer=optimizer,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n            lr_scheduler=lr_scheduler,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n\n        assert isinstance(ppo_trainer.optimizer.optimizer, torch.optim.SGD)\n        assert isinstance(ppo_trainer.lr_scheduler.scheduler, torch.optim.lr_scheduler.ExponentialLR)\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n        # Finally check stats\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n        # assert that the LR has increased for exponential decay\n        assert train_stats[\"ppo/learning_rate\"] > self.ppo_config.learning_rate\n\n    def test_ppo_step_with_no_ref(self):\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n        self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n        # initialize a new gpt2 model:\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            if \"v_head\" not in name:\n                name = name.replace(\"pretrained_model.\", \"\")\n\n                assert torch.allclose(\n                    param.cpu(), model.state_dict()[name].cpu()\n                ), f\"Parameter {name} has changed from the original model\"\n\n        # Finally check stats\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n    def test_ppo_step_with_no_ref_custom_layers(self):\n        \"\"\"\n        Test PPO step with no reference model and custom layers\n        For shared layers configuration, all the layers after the `num_shared_layers` are considered as custom layers\n        therefore the gradients should be computed for these layers only.\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n        self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n        num_shared_layers = 1\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n            num_shared_layers=num_shared_layers,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            train_stats = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        pattern = r\".*transformer\\.h\\.(\\d+)\\..*\"\n        final_layers = [\"ln_f\", \"v_head\", \"lm_head\"]\n\n        for name, param in ppo_trainer.model.named_parameters():\n            if re.match(pattern, name):\n                layer_number = int(re.match(pattern, name).groups(0)[0])\n                if layer_number < num_shared_layers:\n                    assert param.grad is None, f\"Parameter {name} has a gradient\"\n                else:\n                    assert param.grad is not None, f\"Parameter {name} has no gradient\"\n            elif any(layer in name for layer in final_layers):\n                assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats.keys()\n\n    def test_ppo_step_with_ref_and_custom_layers_warning(self):\n        \"\"\"\n        Test PPO step with a reference model and custom layers\n        The trainer should raise a warning if the argument `num_shared_layers` is set\n        together with a reference model.\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        num_shared_layers = 6\n\n        with self.assertWarns(UserWarning):\n            _ = PPOTrainer(\n                config=self.ppo_config,\n                model=self.gpt2_model,\n                ref_model=self.gpt2_ref_model,\n                tokenizer=self.gpt2_tokenizer,\n                dataset=dummy_dataset,\n                num_shared_layers=num_shared_layers,\n            )\n\n    def test_ppo_step_rewards_shape(self):\n        \"\"\"\n        Test if the rewards shape is correct by asserting that if a wrong reward shape is passed, we get\n        a value error.\n        \"\"\"\n\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor([[1.0]]), torch.tensor([[0.0]])]\n            # train model - this should raise an error\n            with pytest.raises(ValueError):\n                _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n\n            reward = [torch.tensor([1.0]), torch.tensor([0.0])]\n            # train model - this should work\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        # check if the gradients are computed for the model\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n    def test_ppo_step_input_shape(self):\n        \"\"\"\n        Test if the shape of the expected inputs are correct\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor([1.0]), torch.tensor([0.0])]\n            # train model - this should raise an error\n            bs = ppo_trainer.config.batch_size\n\n            queries, responses, _, _ = ppo_trainer._step_safety_checker(\n                bs, list(query_tensor), list(response_tensor), reward\n            )\n\n            assert isinstance(queries, list), f\"queries should be a list, got {type(queries)}\"\n            assert isinstance(responses, list), f\"responses should be a list, got {type(responses)}\"\n\n            # check the shapes\n            for i in range(bs):\n                assert queries[i].shape == torch.Size([7])\n                assert responses[i].size() == torch.Size([7])\n            break\n\n    def test_ppo_step_no_dataset(self):\n        \"\"\"\n        Test if the training loop works fine without passing a dataset\n        \"\"\"\n        query_txt = \"This morning I went to the \"\n        query_tensor = self.gpt2_tokenizer.encode(query_txt, return_tensors=\"pt\")\n        self.ppo_config.batch_size = 1\n\n        response_tensor = respond_to_batch(self.gpt2_model, query_tensor)\n\n        # Check that this warns the user about batch size\n        with self.assertWarns(UserWarning):\n            ppo_trainer = PPOTrainer(\n                config=self.ppo_config,\n                model=self.gpt2_model,\n                ref_model=self.gpt2_ref_model,\n                tokenizer=self.gpt2_tokenizer,\n            )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        # train model with ppo\n        reward = [torch.tensor([1.0])]\n        # train model - this should work fine\n        train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)\n\n        # check gradients\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n\n        # ref model should not be trained\n        for name, param in ppo_trainer.ref_model.named_parameters():\n            assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n        # check train stats\n        for stat in EXPECTED_STATS:\n            assert stat in train_stats, f\"Train stats should contain {stat}\"\n\n    def test_loss_trainer(self):\n        \"\"\"\n        Test if the loss trainer works fine\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        self.gpt2_model.eval()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        dummy_queries = [torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3, 4, 5, 6, 7])]\n        dummy_responses = [torch.tensor([5, 6, 7, 8, 9]), torch.tensor([8, 9, 10, 11, 12, 13])]\n        dummy_scores = torch.Tensor([1, 2])\n\n        ppo_trainer.config.mini_batch_size = 1\n        ppo_trainer.config.batch_size = 1\n        model_inputs = ppo_trainer.prepare_model_inputs(dummy_queries, dummy_responses)\n        all_logprobs, _, values, mask = ppo_trainer.batched_forward_pass(\n            self.gpt2_model, dummy_queries, dummy_responses, model_inputs\n        )\n\n        # dummy values\n        ref_logprobs = all_logprobs + 1\n        logits = torch.exp(all_logprobs)\n        vpreds = values + 0.1\n\n        score, non_score, kls = ppo_trainer.compute_rewards(dummy_scores, all_logprobs, ref_logprobs, mask)\n        values, advantages, returns = ppo_trainer.compute_advantages(values, score, mask)\n\n        # just make sure a dummy loss is computed\n        idx = 0\n        pg_loss, v_loss, _ = ppo_trainer.loss(\n            all_logprobs[idx].unsqueeze(0),\n            values[idx].unsqueeze(0),\n            logits[idx].unsqueeze(0),\n            vpreds[idx].unsqueeze(0),\n            ref_logprobs[idx].unsqueeze(0),\n            mask[idx].unsqueeze(0),\n            advantages[idx].unsqueeze(0),\n            returns[idx].unsqueeze(0),\n        )\n\n        assert abs(pg_loss.item() - 1.8226) < 0.0001\n        assert abs(v_loss.item() - 0.1260) < 0.0001\n\n        # check if we get same results with masked parts removed\n        pg_loss_unmasked, v_loss_unmasked, _ = ppo_trainer.loss(\n            apply_mask(all_logprobs[idx], mask[idx]).unsqueeze(0),\n            apply_mask(values[idx], mask[idx]).unsqueeze(0),\n            apply_mask(logits[idx], mask[idx]).unsqueeze(0),\n            apply_mask(vpreds[idx], mask[idx]).unsqueeze(0),\n            apply_mask(ref_logprobs[idx], mask[idx]).unsqueeze(0),\n            apply_mask(mask[idx], mask[idx]).unsqueeze(0),\n            apply_mask(advantages[idx], mask[idx]).unsqueeze(0),\n            apply_mask(returns[idx], mask[idx]).unsqueeze(0),\n        )\n        assert abs(pg_loss_unmasked.item() - 1.8226) < 0.0001\n        assert abs(v_loss_unmasked.item() - 0.1260) < 0.0001\n\n    @parameterized.expand(\n        [\n            [\"gpt2\"],\n            [\"bloom\"],\n            [\"t5\"],\n        ]\n    )\n    def test_batched_forward_pass(self, name):\n        \"\"\"\n        Test if the loss trainer works fine\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        dummy_queries = [torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3, 4, 5, 6, 7])]\n        dummy_responses = [torch.tensor([5, 6, 7, 8, 9]), torch.tensor([8, 9, 10, 11, 12, 13])]\n\n        if name == \"gpt2\":\n            model = self.gpt2_model\n            tokenizer = self.gpt2_tokenizer\n        elif name == \"bloom\":\n            model = self.bloom_model\n            tokenizer = self.bloom_tokenizer\n        elif name == \"t5\":\n            model = self.t5_model\n            tokenizer = self.t5_tokenizer\n\n        model.eval()\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=model,\n            ref_model=None,\n            tokenizer=tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        # we test all combinations of fwd_bs and bs:\n        # if fwd_bs=bs=1: no padding is applied and only one forward pass\n        # if fwd_bs=1/bs=2: padding is applied and results computed in two fwd passes\n        # if fwd_bs=bs=2: padding is applied and results computed in one fwd pass\n\n        ppo_trainer.config.mini_batch_size = 1\n        ppo_trainer.config.batch_size = 1\n\n        model_inputs = ppo_trainer.prepare_model_inputs([dummy_queries[0]], [dummy_responses[0]])\n        logprobs_0, logits_0, values_0, mask_0 = ppo_trainer.batched_forward_pass(\n            model, [dummy_queries[0]], [dummy_responses[0]], model_inputs\n        )\n\n        ppo_trainer.config.batch_size = 2\n        model_inputs = ppo_trainer.prepare_model_inputs(dummy_queries, dummy_responses)\n        logprobs_1, logits_1, values_1, mask_1 = ppo_trainer.batched_forward_pass(\n            model, dummy_queries, dummy_responses, model_inputs\n        )\n\n        ppo_trainer.config.mini_batch_size = 2\n        model_inputs = ppo_trainer.prepare_model_inputs(dummy_queries, dummy_responses)\n        logprobs_2, logits_2, values_2, mask_2 = ppo_trainer.batched_forward_pass(\n            model, dummy_queries, dummy_responses, model_inputs\n        )\n\n        assert abs_diff_masked_tensors(logprobs_1, logprobs_2, mask_1, mask_2) <= 0.0001\n        assert abs_diff_masked_tensors(values_1, values_2, mask_1, mask_2) <= 0.0001\n\n        assert abs_diff_masked_tensors(logprobs_0, logprobs_2[:1], mask_0, mask_2[:1]) <= 0.0001\n        assert abs_diff_masked_tensors(values_0, values_2[:1], mask_0, mask_2[:1]) <= 0.0001\n\n    def test_ppo_trainer_max_grad_norm(self):\n        \"\"\"\n        Test if the `max_grad_norm` feature works as expected\n        \"\"\"\n        # initialize dataset\n        dummy_dataset = self._init_dummy_dataset()\n\n        self.ppo_config.max_grad_norm = 0.00001\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        # check gradients\n        for name, param in ppo_trainer.model.named_parameters():\n            assert param.grad is not None, f\"Parameter {name} has no gradient\"\n            assert torch.all(\n                param.grad.abs() <= self.ppo_config.max_grad_norm\n            ), f\"Parameter {name} has a gradient larger than max_grad_norm\"\n\n    def test_ppo_trainer_kl_penalty(self):\n        dummy_dataset = self._init_dummy_dataset()\n\n        log_probs = torch.Tensor([[0.5, 0.2, 0.1], [0.6, 0.2, 0.1]])\n        ref_log_probs = torch.Tensor([[0.4, 0.3, 0.0], [0.7, 0.1, 0.3]])\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        expected_output = torch.Tensor([[0.1000, -0.1000, 0.1000], [-0.1000, 0.1000, -0.2000]])\n        assert torch.allclose(ppo_trainer._kl_penalty(log_probs, ref_log_probs), expected_output)\n\n        self.ppo_config.kl_penalty = \"abs\"\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        expected_output = torch.Tensor([[0.1000, 0.1000, 0.1000], [0.1000, 0.1000, 0.2000]])\n        assert torch.allclose(ppo_trainer._kl_penalty(log_probs, ref_log_probs), expected_output)\n\n        self.ppo_config.kl_penalty = \"mse\"\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        expected_output = torch.Tensor([[0.0050, 0.0050, 0.0050], [0.0050, 0.0050, 0.0200]])\n        assert torch.allclose(ppo_trainer._kl_penalty(log_probs, ref_log_probs), expected_output)\n\n    def test_ppo_trainer_full_kl_penalty(self):\n        # a few more extensive tests for the full kl option as it is more involved\n        dummy_dataset = self._init_dummy_dataset()\n\n        self.ppo_config.kl_penalty = \"full\"\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        # Test on tensors for size B,S,T = (1,2,3)\n        # test for when the two dists are the same\n        log_probs = torch.Tensor(\n            [\n                [\n                    [0.1, 0.2, 0.7],\n                    [0.3, 0.4, 0.3],\n                ]\n            ]\n        ).exp()\n\n        ref_log_probs = torch.Tensor(\n            [\n                [\n                    [0.1, 0.2, 0.7],\n                    [0.3, 0.4, 0.3],\n                ]\n            ]\n        ).exp()\n\n        expected_output = torch.Tensor(\n            [[0.0, 0.0]],\n        )\n        output = ppo_trainer._kl_penalty(log_probs, ref_log_probs)\n        assert output.shape == (1, 2)\n        assert torch.allclose(output, expected_output)\n\n        # test for when the two dists are almost not overlapping\n        log_probs = torch.Tensor(\n            [\n                [\n                    [0.98, 0.01, 0.01],\n                    [0.01, 0.98, 0.01],\n                ]\n            ]\n        ).log()\n\n        ref_log_probs = torch.Tensor(\n            [\n                [\n                    [0.01, 0.01, 0.98],\n                    [0.01, 0.01, 0.98],\n                ]\n            ]\n        ).log()\n\n        expected_output = torch.Tensor(\n            [[4.4474, 4.4474]],\n        )\n        output = ppo_trainer._kl_penalty(log_probs, ref_log_probs)\n        assert output.shape == (1, 2)\n        assert torch.allclose(output, expected_output)\n\n        # test for when the two dists are almost not overlapping\n        log_probs = torch.Tensor(\n            [\n                [\n                    [0.49, 0.02, 0.49],\n                    [0.49, 0.02, 0.49],\n                ]\n            ]\n        ).log()\n\n        ref_log_probs = torch.Tensor(\n            [\n                [\n                    [0.01, 0.98, 0.01],\n                    [0.49, 0.02, 0.49],\n                ]\n            ]\n        ).log()\n\n        expected_output = torch.Tensor(\n            [[3.7361, 0.0]],\n        )\n        output = ppo_trainer._kl_penalty(log_probs, ref_log_probs)\n        assert output.shape == (1, 2)\n        assert torch.allclose(output, expected_output, atol=0.0001)\n\n    @require_peft\n    def test_peft_model_ppo_trainer(self):\n        from peft import LoraConfig, get_peft_model\n        from transformers import AutoModelForCausalLM\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n        gpt2_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n\n        # this line is very important\n        def make_inputs_require_grad(module, input, output):\n            output.requires_grad_(True)\n\n        gpt2_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        peft_model = get_peft_model(gpt2_model, lora_config)\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(peft_model)\n\n        dummy_dataset = self._init_dummy_dataset()\n        self.ppo_config.batch_size = 2\n        self.ppo_config.mini_batch_size = 1\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        assert ppo_trainer.ref_model is None\n\n        dummy_dataloader = ppo_trainer.dataloader\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model by running a step twice\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n\n            ppo_trainer.model.train()\n            ppo_trainer.model.gradient_checkpointing_enable()\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        # check gradients\n        for name, param in model.named_parameters():\n            if \"lora\" in name or \"v_head\" in name:\n                assert param.grad is not None, f\"Parameter {name} has a no gradient\"\n            else:\n                assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n    @require_peft\n    def test_peft_model_ppo_adapter_rm_trainer(self):\n        from peft import LoraConfig, get_peft_model\n        from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification\n\n        dummy_inputs = torch.LongTensor([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]])\n        rm_lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"SEQ_CLS\",\n        )\n\n        reward_model = AutoModelForSequenceClassification.from_pretrained(self.model_id)\n        reward_model = get_peft_model(reward_model, rm_lora_config)\n        dummy_optim = torch.optim.Adam(filter(lambda p: p.requires_grad, reward_model.parameters()), lr=1e-3)\n\n        previous_rm_logits = reward_model(dummy_inputs).logits\n        loss = previous_rm_logits.mean()\n        loss.backward()\n\n        dummy_optim.step()\n        reward_model.eval()\n\n        original_rm_logits = reward_model(dummy_inputs).logits\n\n        with tempfile.TemporaryDirectory() as tmpdirname:\n            reward_model.save_pretrained(tmpdirname)\n\n            lora_config = LoraConfig(\n                r=16,\n                lora_alpha=32,\n                lora_dropout=0.05,\n                bias=\"none\",\n                task_type=\"CAUSAL_LM\",\n            )\n            gpt2_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n\n            # this line is very important\n            def make_inputs_require_grad(module, input, output):\n                output.requires_grad_(True)\n\n            gpt2_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            peft_model = get_peft_model(gpt2_model, lora_config)\n            model = AutoModelForCausalLMWithValueHead.from_pretrained(\n                peft_model,\n                reward_adapter=tmpdirname,\n            )\n\n            dummy_dataset = self._init_dummy_dataset()\n            self.ppo_config.batch_size = 2\n            self.ppo_config.mini_batch_size = 1\n\n            ppo_trainer = PPOTrainer(\n                config=self.ppo_config,\n                model=model,\n                ref_model=None,\n                tokenizer=self.gpt2_tokenizer,\n                dataset=dummy_dataset,\n            )\n            ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n            assert ppo_trainer.ref_model is None\n\n            dummy_dataloader = ppo_trainer.dataloader\n\n            # train model with ppo\n            for query_tensor, response_tensor in dummy_dataloader:\n                # define a reward for response\n                # (this could be any reward such as human feedback or output from another model)\n                reward = [torch.tensor(1.0), torch.tensor(0.0)]\n                # train model by running a step twice\n                _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n\n                ppo_trainer.model.train()\n                ppo_trainer.model.gradient_checkpointing_enable()\n                _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n                break\n\n            dummy_inputs = dummy_inputs.to(ppo_trainer.accelerator.device)\n            new_logits = ppo_trainer.model.compute_reward_score(dummy_inputs)\n            assert not torch.allclose(previous_rm_logits.to(ppo_trainer.accelerator.device), new_logits[:, -1, :])\n            assert torch.allclose(original_rm_logits.to(ppo_trainer.accelerator.device), new_logits[:, -1, :])\n\n            # check gradients\n            for name, param in model.named_parameters():\n                if (\"lora\" in name or \"v_head\" in name) and (\"reward\" not in name):\n                    assert param.grad is not None, f\"Parameter {name} has a no gradient\"\n                else:\n                    assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n    @unittest.skip(\"Fix by either patching `whomai()` to work in the staging endpoint or use a dummy prod user.\")\n    def test_push_to_hub(self):\n        REPO_NAME = \"test-ppo-trainer\"\n        repo_id = f\"{CI_HUB_USER}/{REPO_NAME}\"\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=self.gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=self._init_dummy_dataset(),\n        )\n        with tempfile.TemporaryDirectory():\n            url = ppo_trainer.push_to_hub(repo_id=repo_id, token=self._token, api_endpoint=CI_HUB_ENDPOINT)\n            # Extract repo_name from the url\n            re_search = re.search(CI_HUB_ENDPOINT + r\"/([^/]+/[^/]+)/\", url)\n            assert re_search is not None\n            hub_repo_id = re_search.groups()[0]\n            # Check we created a Hub repo\n            assert hub_repo_id == repo_id\n            # Ensure all files are present\n            files = sorted(self._api.list_repo_files(hub_repo_id))\n            assert all(\n                fnmatch.fnmatch(file, expected_file)\n                for file, expected_file in zip(\n                    files,\n                    [\n                        \".gitattributes\",\n                        \"README.md\",\n                        \"config.json\",\n                        \"merges.txt\",\n                        \"pytorch_model.bin\",\n                        \"special_tokens_map.json\",\n                        \"tokenizer_config.json\",\n                        \"vocab.json\",\n                    ],\n                )\n            )\n\n    @require_peft\n    @require_torch_multi_accelerator\n    def test_peft_model_ppo_trainer_multi_gpu(self):\n        from peft import LoraConfig, get_peft_model\n        from transformers import AutoModelForCausalLM\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n        gpt2_model = AutoModelForCausalLM.from_pretrained(\n            \"gpt2\", device_map=\"balanced\", max_memory={0: \"500MB\", 1: \"500MB\"}\n        )\n\n        assert set(gpt2_model.hf_device_map.values()) == {0, 1}\n\n        # this line is very important\n        def make_inputs_require_grad(module, input, output):\n            output.requires_grad_(True)\n\n        gpt2_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        peft_model = get_peft_model(gpt2_model, lora_config)\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(peft_model)\n\n        assert model.is_sequential_parallel\n\n        dummy_dataset = self._init_dummy_dataset()\n        self.ppo_config.batch_size = 2\n        self.ppo_config.mini_batch_size = 1\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        assert ppo_trainer.ref_model is None\n\n        dummy_dataloader = ppo_trainer.dataloader\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model by running a step twice\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n\n            ppo_trainer.model.train()\n            ppo_trainer.model.gradient_checkpointing_enable()\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        # check gradients\n        for name, param in model.named_parameters():\n            if \"lora\" in name or \"v_head\" in name:\n                assert param.grad is not None, f\"Parameter {name} has a no gradient\"\n            else:\n                assert param.grad is None, f\"Parameter {name} has a gradient\"\n\n    def test_generation(self):\n        dummy_dataset = self._init_dummy_dataset()\n\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=model,\n            ref_model=None,\n            tokenizer=tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        input_texts = [\"this is a test\", \"this is another, longer test\"]\n\n        generation_kwargs = {\"do_sample\": False, \"max_new_tokens\": 4, \"pad_token_id\": tokenizer.eos_token_id}\n\n        tokenizer.pad_token = tokenizer.eos_token\n\n        model_inputs = [tokenizer(txt, return_tensors=\"pt\").input_ids.squeeze() for txt in input_texts]\n        model_inputs = [input_ids.to(ppo_trainer.accelerator.device) for input_ids in model_inputs]\n\n        generations_batched = ppo_trainer.generate(model_inputs, batch_size=2, **generation_kwargs)\n        generations_batched = tokenizer.batch_decode(generations_batched)\n\n        generations_single = [ppo_trainer.generate(inputs, **generation_kwargs).squeeze() for inputs in model_inputs]\n        generations_single = tokenizer.batch_decode(generations_single)\n\n        assert generations_single == generations_batched\n\n    def test_generation_with_ref_model(self):\n        dummy_dataset = self._init_dummy_dataset()\n        model = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\n        tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n        # Negate the weights in the last layer of the ref model so it never\n        # outputs the same things as the primary model\n        ref_model = copy.deepcopy(model)\n        lm_head_weight = ref_model.pretrained_model.lm_head.weight\n        lm_head_weight.data = -lm_head_weight.data\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=model,\n            ref_model=ref_model,\n            tokenizer=tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        input_texts = [\"this is a test\", \"this is another, longer test\"]\n\n        generation_kwargs = {\"do_sample\": False, \"max_new_tokens\": 4, \"pad_token_id\": tokenizer.eos_token_id}\n\n        tokenizer.pad_token = tokenizer.eos_token\n\n        model_inputs = [tokenizer(txt, return_tensors=\"pt\").input_ids.squeeze() for txt in input_texts]\n        model_inputs = [input_ids.to(ppo_trainer.accelerator.device) for input_ids in model_inputs]\n\n        generations_batched, ref_generations_batched = ppo_trainer.generate(\n            model_inputs, batch_size=2, generate_ref_response=True, **generation_kwargs\n        )\n        generations_batched = tokenizer.batch_decode(generations_batched)\n        ref_generations_batched = tokenizer.batch_decode(ref_generations_batched)\n\n        generations_single = []\n        ref_generations_single = []\n        for inputs in model_inputs:\n            generation, ref_generation = ppo_trainer.generate(inputs, generate_ref_response=True, **generation_kwargs)\n            generations_single.append(generation.squeeze())\n            ref_generations_single.append(ref_generation.squeeze())\n\n        generations_single = tokenizer.batch_decode(generations_single)\n        ref_generations_single = tokenizer.batch_decode(ref_generations_single)\n\n        assert generations_single == generations_batched\n        assert ref_generations_single == ref_generations_batched\n\n        assert generations_batched != ref_generations_batched\n        assert generations_single != ref_generations_single\n\n    def test_grad_accumulation(self):\n        dummy_dataset = self._init_dummy_dataset()\n\n        torch.manual_seed(0)\n        gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id, summary_dropout_prob=0.0)\n        gpt2_model_clone = copy.deepcopy(gpt2_model)\n\n        self.ppo_config.mini_batch_size = 2\n        self.ppo_config.ppo_epochs = 1\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=gpt2_model,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        dummy_dataloader = ppo_trainer.dataloader\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(1.0)]\n            # train model by running a step twice\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        model_grad = gpt2_model.v_head.summary.weight\n\n        self.ppo_config.mini_batch_size = 1\n        self.ppo_config.gradient_accumulation_steps = 2\n\n        ppo_trainer = PPOTrainer(\n            config=self.ppo_config,\n            model=gpt2_model_clone,\n            ref_model=None,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        dummy_dataloader = ppo_trainer.dataloader\n\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(1.0)]\n            # train model by running a step twice\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n        model_grad_acc = gpt2_model_clone.v_head.summary.weight\n        assert torch.allclose(model_grad_acc, model_grad, rtol=0.001, atol=0.001)\n\n    @unittest.skip(\"Fix by either patching `whomai()` to work in the staging endpoint or use a dummy prod user.\")\n    def test_push_to_hub_if_best_reward(self):\n        REPO_NAME = \"test-ppo-trainer\"\n        repo_id = f\"{CI_HUB_USER}/{REPO_NAME}\"\n\n        dummy_dataset = self._init_dummy_dataset()\n\n        push_to_hub_if_best_kwargs = {\"repo_id\": repo_id}\n\n        ppo_config = PPOConfig(\n            batch_size=2,\n            mini_batch_size=1,\n            log_with=None,\n            push_to_hub_if_best_kwargs=push_to_hub_if_best_kwargs,\n            compare_steps=1,\n        )\n\n        ppo_trainer = PPOTrainer(\n            config=ppo_config,\n            model=self.gpt2_model,\n            ref_model=self.gpt2_ref_model,\n            tokenizer=self.gpt2_tokenizer,\n            dataset=dummy_dataset,\n        )\n\n        ppo_trainer.optimizer.zero_grad = partial(ppo_trainer.optimizer.zero_grad, set_to_none=False)\n        dummy_dataloader = ppo_trainer.dataloader\n        # train model with ppo\n        for query_tensor, response_tensor in dummy_dataloader:\n            # define a reward for response\n            # (this could be any reward such as human feedback or output from another model)\n            reward = [torch.tensor(1.0), torch.tensor(0.0)]\n            # train model\n            _ = ppo_trainer.step(list(query_tensor), list(response_tensor), reward)\n            break\n\n    def test_batch_size_check(self):\n        with pytest.raises(ValueError):\n            PPOConfig(batch_size=2, mini_batch_size=2, gradient_accumulation_steps=2)\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom unittest.mock import patch\n\nimport torch\nfrom transformers import AutoTokenizer\n\nfrom trl import AutoModelForCausalLMWithValueHead, TextEnvironment, TextHistory\n\n\nclass DummyTool:\n    def __call__(self, text):\n        return text\n\n\ndef dummy_generate(histories):\n    for i in range(len(histories)):\n        histories[i].append_segment(\"<request><DummyTool>test<call>\", torch.tensor([1, 2, 3]), system=False)\n    return histories\n\n\nclass TextHistoryTest(unittest.TestCase):\n    def test_text_history_init(self):\n        text = \"Hello there!\"\n        tokens = torch.tensor([1, 2, 3])\n\n        history = TextHistory(text, tokens)\n        assert history.text == text\n        assert torch.equal(history.tokens, tokens)\n        assert torch.equal(history.token_masks, torch.zeros_like(tokens))\n\n        history = TextHistory(text, tokens, system=False)\n        assert torch.equal(history.token_masks, torch.ones_like(tokens))\n\n    def test_text_history_append_segment(self):\n        text = \"Hello there!\"\n        tokens = torch.tensor([1, 2, 3])\n\n        history = TextHistory(text, tokens)\n        history.append_segment(\"General Kenobi!\", torch.tensor([4, 5, 6]), system=False)\n        assert history.text == (text + \"General Kenobi!\")\n        assert torch.equal(history.tokens, torch.tensor([1, 2, 3, 4, 5, 6]))\n        assert torch.equal(history.token_masks, torch.tensor([0, 0, 0, 1, 1, 1]))\n\n        history.append_segment(\"You are a bold one!\", torch.tensor([7, 8, 9]))\n        assert history.text == ((text + \"General Kenobi!\") + \"You are a bold one!\")\n        assert torch.equal(history.tokens, torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9]))\n        assert torch.equal(history.token_masks, torch.tensor([0, 0, 0, 1, 1, 1, 0, 0, 0]))\n\n    def test_text_history_complete(self):\n        text = \"Hello there!\"\n        tokens = torch.tensor([1, 2, 3])\n        history = TextHistory(text, tokens)\n        history.complete()\n        assert history.completed\n        assert not history.truncated\n\n        history.complete(truncated=True)\n        assert history.completed\n        assert history.truncated\n\n    def test_text_history_last_segment(self):\n        text = \"Hello there!\"\n        tokens = torch.tensor([1, 2, 3])\n        history = TextHistory(text, tokens)\n        history.append_segment(\"General Kenobi!\", torch.tensor([4, 5, 6]))\n        history.append_segment(\"You are a bold one!\", torch.tensor([7, 8, 9]))\n        assert history.last_text_segment == \"You are a bold one!\"\n\n    def test_text_history_split_query_response(self):\n        text = \"Hello there!\"\n        tokens = torch.tensor([1, 2, 3])\n        history = TextHistory(text, tokens)\n        history.append_segment(\"General Kenobi!\", torch.tensor([4, 5, 6]), system=False)\n        history.append_segment(\"You are a bold one!\", torch.tensor([7, 8, 9]), system=True)\n        query, response, mask = history.split_query_response_tokens()\n\n        assert torch.equal(query, torch.tensor([1, 2, 3]))\n        assert torch.equal(response, torch.tensor([4, 5, 6, 7, 8, 9]))\n        assert torch.equal(mask, torch.tensor([1, 1, 1, 0, 0, 0]))\n\n\nclass TextEnvironmentTester(unittest.TestCase):\n    def setUp(self):\n        # model_id\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n\n        # get models and tokenizer\n        self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id)\n        self.gpt2_tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.gpt2_tokenizer.pad_token = self.gpt2_tokenizer.eos_token\n\n    def test_text_environment_setup(self):\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools=[DummyTool()],\n            reward_fn=lambda x: torch.tensor(1),\n            prompt=\"I am a prompt!\\n\",\n        )\n        assert env.prompt == \"I am a prompt!\\n\"\n        assert list(env.tools.keys()) == [\"DummyTool\"]\n        assert isinstance(env.tools[\"DummyTool\"], DummyTool)\n        assert env.reward_fn(\"Hello there!\") == 1\n\n    def test_text_environment_generate(self):\n        generation_kwargs = {\"do_sample\": False, \"max_new_tokens\": 4, \"pad_token_id\": self.gpt2_tokenizer.eos_token_id}\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools=[DummyTool()],\n            reward_fn=lambda x: torch.tensor(1),\n            prompt=\"I am a prompt!\\n\",\n            generation_kwargs=generation_kwargs,\n        )\n\n        input_texts = [\"this is a test\", \"this is another, longer test\"]\n\n        model_inputs = [self.gpt2_tokenizer(txt, return_tensors=\"pt\").input_ids.squeeze() for txt in input_texts]\n\n        generations_batched = env._generate_batched(model_inputs, batch_size=2)\n        generations_batched = self.gpt2_tokenizer.batch_decode(generations_batched)\n\n        generations_single = [env._generate_batched([inputs], batch_size=1)[0] for inputs in model_inputs]\n        generations_single = self.gpt2_tokenizer.batch_decode(generations_single)\n\n        assert generations_single == generations_batched\n\n    def test_text_environment_tool_call_parsing(self):\n        string_valid = \"Something something <request><Tool1>Hello there!<call>\"\n        string_invalid_request = \"Something something <Tool1>Hello there!<call>\"\n        string_invalid_call = \"Something something <request><Tool1>Hello there!\"\n        string_invalid_tool = \"Something something <request>|Tool2|Hello there!<call>\"\n        string_invalid_random = \"<>abcdefghijklm<>nopqrstuvwxyz<>\"\n\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools=[DummyTool()],\n            reward_fn=lambda x: torch.tensor(1),\n            prompt=\"I am a prompt!\\n\",\n        )\n        tool, response = env.parse_tool_call(string_valid)\n        assert tool == \"Tool1\"\n        assert response == \"Hello there!\"\n\n        tool, response = env.parse_tool_call(string_invalid_request)\n        assert tool is None\n        assert response is None\n\n        tool, response = env.parse_tool_call(string_invalid_call)\n        assert tool is None\n        assert response is None\n\n        tool, response = env.parse_tool_call(string_invalid_tool)\n        assert tool is None\n        assert response is None\n\n        tool, response = env.parse_tool_call(string_invalid_random)\n        assert tool is None\n        assert response is None\n\n    def test_text_environment_tool_truncation(self):\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools={\"dummy\": lambda x: \"a\" * 1000},\n            reward_fn=lambda x: torch.tensor(1),\n            prompt=\"I am a prompt!\\n\",\n        )\n\n        env.max_tool_response = 100\n        history = env.step(TextHistory(\"<request><dummy>Hello there!<call>\", torch.tensor([1, 2, 3])))\n        assert (len(history.last_text_segment) - len(env.response_token)) == 100\n\n        env.max_tool_response = 500\n        history = env.step(TextHistory(\"<request><dummy>Hello there!<call>\", torch.tensor([1, 2, 3])))\n        assert (len(history.last_text_segment) - len(env.response_token)) == 500\n\n        env.max_tool_response = 1001\n        history = env.step(TextHistory(\"<request><dummy>Hello there!<call>\", torch.tensor([1, 2, 3])))\n        assert (len(history.last_text_segment) - len(env.response_token)) == 1000\n\n        env.max_tool_response = 2000\n        history = env.step(TextHistory(\"<request><dummy>Hello there!<call>\", torch.tensor([1, 2, 3])))\n        assert (len(history.last_text_segment) - len(env.response_token)) == 1000\n\n    @patch.object(TextEnvironment, \"generate\", side_effect=dummy_generate)\n    def test_text_environment_max_calls(self, mock_generate):\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools={\"DummyTool\": DummyTool()},\n            reward_fn=lambda x: [torch.tensor(1) for _ in x],\n            prompt=\"I am a prompt!\\n\",\n        )\n\n        env.max_turns = 1\n        _, _, _, _, histories = env.run([\"test\"])\n        assert histories[0].text == (\n            (\"I am a prompt!\\n\" + \"test\") + (1 * \"<request><DummyTool>test<call>test<response>\")\n        )\n\n        env.max_turns = 2\n        _, _, _, _, histories = env.run([\"test\"])\n        assert histories[0].text == (\n            (\"I am a prompt!\\n\" + \"test\") + (2 * \"<request><DummyTool>test<call>test<response>\")\n        )\n\n        env.max_turns = 4\n        _, _, _, _, histories = env.run([\"test\"])\n        assert histories[0].text == (\n            (\"I am a prompt!\\n\" + \"test\") + (4 * \"<request><DummyTool>test<call>test<response>\")\n        )\n\n    def test_text_environment_compute_rewards(self):\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools={\"DummyTool\": DummyTool()},\n            reward_fn=lambda x: [torch.tensor(i) for i, _ in enumerate(x)],\n            prompt=\"I am a prompt!\\n\",\n        )\n\n        histories = [TextHistory(\"<request><DummyTool>test<call>\", torch.tensor([1, 2, 3])) for _ in range(8)]\n        histories = env.compute_reward(histories)\n\n        for i in range(8):\n            assert histories[i].reward == i\n\n    @patch.object(TextEnvironment, \"generate\", side_effect=dummy_generate)\n    def test_text_environment_run(self, mock_generate):\n        env = TextEnvironment(\n            self.gpt2_model,\n            self.gpt2_tokenizer,\n            tools={\"DummyTool\": DummyTool()},\n            reward_fn=lambda x: [torch.tensor(i) for i, _ in enumerate(x)],\n            prompt=\"I am a prompt!\\n\",\n            max_turns=2,\n        )\n        task_1 = \"Hello there!\"\n        task_2 = \"Hello there! General Kenobi!\"\n\n        query, response, response_mask, reward, histories = env.run([task_1, task_2])\n        assert len(query[0]) == 9\n        assert len(query[1]) == 12\n        assert len(response[0]) == 14\n        assert len(response[1]) == 14\n        assert response_mask[0].sum() == (2 * 3)\n        # mocked generate always adds 3 toknes\n        assert response_mask[1].sum() == (2 * 3)\n        # mocked generate always adds 3 toknes\n        assert reward[0] == 0\n        assert reward[1] == 1\n        assert histories[0].text == (\n            (\"I am a prompt!\\n\" + \"Hello there!\") + (2 * \"<request><DummyTool>test<call>test<response>\")\n        )\n        assert histories[1].text == (\n            (\"I am a prompt!\\n\" + \"Hello there! General Kenobi!\")\n            + (2 * \"<request><DummyTool>test<call>test<response>\")\n        )\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\nfrom transformers.testing_utils import require_peft\nfrom transformers.utils import is_peft_available\n\nfrom trl import NashMDConfig, NashMDTrainer\n\n\nif is_peft_available():\n    from peft import LoraConfig, get_peft_model\n\n\nclass TestNashMDTrainer(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.reward_model = AutoModelForSequenceClassification.from_pretrained(\"EleutherAI/pythia-14m\", num_labels=1)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n    @parameterized.expand([(\"standard_prompt_only\",), (\"conversational_prompt_only\",)])\n    def test_nash_md_trainer_training(self, config_name):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = NashMDConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", config_name)\n\n            trainer = NashMDTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = NashMDConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = NashMDTrainer(\n                model=self.model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft_and_ref_model(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = NashMDConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = NashMDTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    def test_training_with_peft_model_and_peft_config(self):\n        model_lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.1, bias=\"none\", task_type=\"CAUSAL_LM\")\n        model = get_peft_model(self.model, model_lora_config)\n        # we want only the \"train adapter\" to be trained\n        lora_train_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = NashMDConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = NashMDTrainer(\n                model=model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_train_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport tempfile\nimport unittest\n\nimport torch\nimport torch.nn.functional as F\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n\nfrom trl import GKDConfig, GKDTrainer\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\nclass TestGKDTrainer(unittest.TestCase):\n    @classmethod\n    def setUpClass(cls):\n        cls.tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n        cls.tokenizer.pad_token = cls.tokenizer.eos_token\n        cls.model = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n        cls.generation_config = GenerationConfig(\n            max_new_tokens=20,\n            num_return_sequences=1,\n            pad_token_id=cls.tokenizer.pad_token_id,\n            eos_token_id=cls.tokenizer.eos_token_id,\n        )\n\n    def test_generate_on_policy_outputs_deterministic(self):\n        prompts = [\"Hello, how are you?\", \"What's the weather like today?\"]\n        tokenized_prompts = self.tokenizer(prompts, return_tensors=\"pt\", padding=True)\n\n        inputs = {\n            \"prompts\": tokenized_prompts[\"input_ids\"],\n            \"prompt_attention_mask\": tokenized_prompts[\"attention_mask\"],\n        }\n\n        # Set temperature to 0 for deterministic output\n        deterministic_generation_config = GenerationConfig(\n            max_new_tokens=30,\n            num_return_sequences=1,\n            pad_token_id=self.tokenizer.pad_token_id,\n            eos_token_id=self.tokenizer.eos_token_id,\n            temperature=0.0,\n        )\n\n        outputs = GKDTrainer.generate_on_policy_outputs(\n            self.model, inputs, deterministic_generation_config, self.tokenizer.pad_token_id\n        )\n\n        new_input_ids, new_attention_mask, new_labels = outputs\n\n        # Decode the generated outputs\n        generated_texts = self.tokenizer.batch_decode(new_input_ids, skip_special_tokens=True)\n\n        # Check if the generated texts start with the original prompts\n        for prompt, generated_text in zip(prompts, generated_texts):\n            self.assertTrue(\n                generated_text.startswith(prompt),\n                f\"Generated text '{generated_text}' does not start with prompt '{prompt}'\",\n            )\n\n        # Run the generation twice and check if the outputs are identical\n        outputs2 = GKDTrainer.generate_on_policy_outputs(\n            self.model, inputs, deterministic_generation_config, self.tokenizer.pad_token_id\n        )\n\n        new_input_ids2, new_attention_mask2, new_labels2 = outputs2\n\n        # Check if the two generations are identical\n        self.assertTrue(torch.all(new_input_ids.eq(new_input_ids2)), \"Deterministic generations are not identical\")\n        self.assertTrue(\n            torch.all(new_attention_mask.eq(new_attention_mask2)),\n            \"Attention masks for deterministic generations are not identical\",\n        )\n        self.assertTrue(\n            torch.all(new_labels.eq(new_labels2)),\n            \"Labels for deterministic generations are not identical\",\n        )\n\n    def test_generate_on_policy_outputs(self):\n        prompts = [\"Hello, how are you?\", \"What's the weather like today?\"]\n        tokenized_prompts = self.tokenizer(prompts, return_tensors=\"pt\", padding=True)\n\n        inputs = {\n            \"prompts\": tokenized_prompts[\"input_ids\"],\n            \"attention_mask\": tokenized_prompts[\"attention_mask\"],\n        }\n\n        outputs = GKDTrainer.generate_on_policy_outputs(\n            self.model, inputs, self.generation_config, self.tokenizer.pad_token_id\n        )\n\n        # Check that outputs is a tuple of three tensors\n        self.assertIsInstance(outputs, tuple)\n        self.assertEqual(len(outputs), 3)\n\n        new_input_ids, new_attention_mask, new_labels = outputs\n\n        # Check shapes\n        batch_size = len(prompts)\n        self.assertEqual(new_input_ids.shape[0], batch_size)\n        self.assertEqual(new_attention_mask.shape[0], batch_size)\n        self.assertEqual(new_labels.shape[0], batch_size)\n\n        # Check types\n        self.assertIsInstance(new_input_ids, torch.Tensor)\n        self.assertIsInstance(new_attention_mask, torch.Tensor)\n        self.assertIsInstance(new_labels, torch.Tensor)\n\n        # Check that new_input_ids and new_attention_mask have the same shape\n        self.assertEqual(new_input_ids.shape, new_attention_mask.shape)\n        self.assertEqual(new_labels.shape, new_attention_mask.shape)\n\n\nclass TestGeneralizedJSDLoss(unittest.TestCase):\n    def setUp(self):\n        self.batch_size = 2\n        self.seq_length = 3\n        self.vocab_size = 5\n        self.student_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)\n        self.teacher_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)\n\n    def test_uniform_distribution(self):\n        logits = torch.ones(1, 1, self.vocab_size)\n        loss = GKDTrainer.generalized_jsd_loss(logits, logits)\n        self.assertAlmostEqual(loss.item(), 0, places=5)\n\n    def test_generalized_jsd_loss_edge_cases(self):\n        # Setup\n        student_logits = torch.log(torch.tensor([[0.1, 0.9]])).unsqueeze(0)\n        teacher_logits = torch.log(torch.tensor([[0.9, 0.1]])).unsqueeze(0)\n\n        # Case 1: beta = 1 (should be equivalent to KL(student || teacher))\n        loss_beta_1 = GKDTrainer.generalized_jsd_loss(student_logits, teacher_logits, beta=1)\n        expected_loss_beta_1 = F.kl_div(\n            F.log_softmax(student_logits, dim=-1), F.softmax(teacher_logits, dim=-1), reduction=\"batchmean\"\n        )\n        self.assertAlmostEqual(loss_beta_1.item(), expected_loss_beta_1.item(), places=5)\n\n        # Case 2: beta = 0 (should be equivalent to KL(teacher || student))\n        loss_beta_0 = GKDTrainer.generalized_jsd_loss(student_logits, teacher_logits, beta=0)\n        expected_loss_beta_0 = F.kl_div(\n            F.log_softmax(teacher_logits, dim=-1), F.softmax(student_logits, dim=-1), reduction=\"batchmean\"\n        )\n        self.assertAlmostEqual(loss_beta_0.item(), expected_loss_beta_0.item(), places=5)\n\n    def test_output_shape(self):\n        loss = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits)\n        self.assertTrue(torch.is_tensor(loss))\n        self.assertEqual(loss.shape, torch.Size([]))\n\n    def test_beta_values(self):\n        loss_beta_0 = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0)\n        loss_beta_1 = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=1)\n        self.assertNotEqual(loss_beta_0, loss_beta_1)\n\n    def test_temperature_scaling(self):\n        loss_temp_1 = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, temperature=1)\n        loss_temp_2 = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, temperature=2)\n        self.assertNotEqual(loss_temp_1, loss_temp_2)\n\n    def test_reduction_methods(self):\n        loss_batchmean = GKDTrainer.generalized_jsd_loss(\n            self.student_logits, self.teacher_logits, reduction=\"batchmean\"\n        )\n        loss_sum = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction=\"sum\")\n        loss_mean = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction=\"mean\")\n        loss_none = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction=\"none\")\n\n        self.assertEqual(loss_batchmean.shape, torch.Size([]))\n        self.assertEqual(loss_sum.shape, torch.Size([]))\n        self.assertEqual(loss_mean.shape, torch.Size([]))\n        self.assertEqual(loss_none.shape, self.student_logits.shape)\n\n    def test_symmetry(self):\n        student_teacher = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0.1)\n        teacher_student = GKDTrainer.generalized_jsd_loss(self.teacher_logits, self.student_logits, beta=0.1)\n        self.assertNotEqual(student_teacher, teacher_student)\n\n        student_teacher = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0.5)\n        teacher_student = GKDTrainer.generalized_jsd_loss(self.teacher_logits, self.student_logits, beta=0.5)\n        self.assertEqual(student_teacher, teacher_student)\n\n    def test_zero_loss_for_identical_inputs(self):\n        identical_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)\n        loss = GKDTrainer.generalized_jsd_loss(identical_logits, identical_logits)\n        self.assertAlmostEqual(loss.item(), 0, places=6)\n\n\nclass GKDTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.teacher_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # Ensure the tokenizer has a chat template\n        if not hasattr(self.tokenizer, \"chat_template\") or self.tokenizer.chat_template is None:\n            self.tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n\n    def test_gkd_trainer(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = GKDConfig(\n                output_dir=tmp_dir,\n                dataloader_drop_last=True,\n                eval_strategy=\"steps\",\n                max_steps=4,\n                eval_steps=2,\n                save_steps=2,\n                per_device_train_batch_size=2,\n                per_device_eval_batch_size=2,\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"conversational_language_modeling\")\n\n            trainer = GKDTrainer(\n                model=self.model_id,\n                teacher_model=self.model_id,\n                args=training_args,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                tokenizer=self.tokenizer,\n            )\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[(-1)][\"train_loss\"])\n            self.assertIsNotNone(trainer.state.log_history[0][\"eval_loss\"])\n            self.assertIn(\"model.safetensors\", os.listdir(tmp_dir + \"/checkpoint-2\"))\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\nfrom transformers.testing_utils import require_peft\nfrom transformers.utils import is_peft_available\n\nfrom trl import XPOConfig, XPOTrainer\n\n\nif is_peft_available():\n    from peft import LoraConfig, get_peft_model\n\n\nclass TestXPOTrainer(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.reward_model = AutoModelForSequenceClassification.from_pretrained(\"EleutherAI/pythia-14m\", num_labels=1)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n    @parameterized.expand([(\"standard_prompt_only\",), (\"conversational_prompt_only\",)])\n    def test_xpo_trainer_training(self, config_name):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = XPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", config_name)\n\n            trainer = XPOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = XPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = XPOTrainer(\n                model=self.model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft_and_ref_model(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = XPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = XPOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    def test_training_with_peft_model_and_peft_config(self):\n        model_lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.1, bias=\"none\", task_type=\"CAUSAL_LM\")\n        model = get_peft_model(self.model, model_lora_config)\n        # we want only the \"train adapter\" to be trained\n        lora_train_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = XPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = XPOTrainer(\n                model=model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_train_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nimport torch\nfrom transformers import AutoTokenizer\nfrom transformers.testing_utils import require_peft\nfrom transformers.utils import is_peft_available\n\nfrom trl.trainer.model_config import ModelConfig\nfrom trl.trainer.utils import decode_and_strip_padding, get_peft_config, pad\n\n\nif is_peft_available():\n    from peft import LoraConfig\n\n\nclass TestPad(unittest.TestCase):\n    def test_pad_1_dim_left(self):\n        x = torch.tensor([1, 2, 3])\n        y = torch.tensor([4, 5])\n        output = pad((x, y), padding_value=0, padding_side=\"left\")\n        expected = torch.tensor([[1, 2, 3], [0, 4, 5]])\n        self.assertTrue(torch.equal(output, expected))\n\n    def test_pad_1_dim_right(self):\n        x = torch.tensor([1, 2, 3])\n        y = torch.tensor([4, 5])\n        output = pad((x, y), padding_value=0, padding_side=\"right\")\n        expected = torch.tensor([[1, 2, 3], [4, 5, 0]])\n        self.assertTrue(torch.equal(output, expected))\n\n    def test_pad_2_dim_left(self):\n        x = torch.tensor([[1, 2], [3, 4]])\n        y = torch.tensor([[5, 6]])\n        output = pad((x, y), padding_value=0, padding_side=\"left\")\n        expected = torch.tensor(\n            [\n                [[1, 2], [3, 4]],\n                [[0, 0], [5, 6]],\n            ]\n        )\n        self.assertTrue(torch.equal(output, expected))\n\n    def test_pad_2_dim_right(self):\n        x = torch.tensor([[1, 2], [3, 4]])\n        y = torch.tensor([[5, 6]])\n        output = pad((x, y), padding_value=0, padding_side=\"right\")\n        expected = torch.tensor(\n            [\n                [[1, 2], [3, 4]],\n                [[5, 6], [0, 0]],\n            ]\n        )\n        self.assertTrue(torch.equal(output, expected))\n\n    def test_pad_2_dim_right_multidim(self):\n        x = torch.tensor([[1, 2], [3, 4]])\n        y = torch.tensor([[5]])\n        output = pad((x, y), padding_value=0, padding_side=\"right\")\n        expected = torch.tensor(\n            [\n                [[1, 2], [3, 4]],\n                [[5, 0], [0, 0]],\n            ]\n        )\n        self.assertTrue(torch.equal(output, expected))\n\n\n@require_peft\nclass TestGetPEFTConfig(unittest.TestCase):\n    def test_create_peft_config_use_peft_false(self):\n        \"\"\"Test that when use_peft is False, the function returns None.\"\"\"\n        model_config = ModelConfig(use_peft=False)\n        peft_config = get_peft_config(model_config)\n        self.assertIsNone(peft_config)\n\n    def test_create_peft_config_use_peft_true(self):\n        \"\"\"Test that when use_peft is True, the function returns a LoraConfig object.\"\"\"\n        # Provide non-default values to the model config for testing\n        peft_kwargs = {\n            \"lora_r\": 8,\n            \"lora_alpha\": 16,\n            \"lora_dropout\": 0.1,\n            \"lora_task_type\": \"SEQ_CLS\",\n            \"use_rslora\": True,\n            \"lora_target_modules\": [\"up_proj\", \"down_proj\"],\n            \"lora_modules_to_save\": [\"up_proj\"],\n        }\n        model_config = ModelConfig(use_peft=True, **peft_kwargs)\n        peft_config = get_peft_config(model_config)\n        self.assertTrue(isinstance(peft_config, LoraConfig))\n        for arg, value in peft_kwargs.items():\n            # Test that lists of modules are converted to sets\n            if arg == \"lora_target_modules\":\n                value = set(value)\n            # Rename the argument to match the LoraConfig attribute name\n            if arg in [\"lora_r\", \"lora_task_type\", \"lora_target_modules\", \"lora_modules_to_save\"]:\n                arg = arg[len(\"lora_\") :] if arg.startswith(\"lora_\") else arg\n\n            self.assertEqual(getattr(peft_config, arg), value)\n\n\nclass TestDecodeAndStripPadding(unittest.TestCase):\n    def setUp(self):\n        self.tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\n\n    def test_example_with_padding(self):\n        inputs = self.tokenizer([\"Hello world\", \"Hello\"], padding=True, return_tensors=\"pt\")\n        decoded = decode_and_strip_padding(inputs[\"input_ids\"], self.tokenizer)\n        self.assertEqual(decoded, [\"Hello world\", \"Hello\"])\n\n    def test_example_without_padding(self):\n        inputs = self.tokenizer([\"Hello\", \"Hello\"], padding=False, return_tensors=\"pt\")\n        decoded = decode_and_strip_padding(inputs[\"input_ids\"], self.tokenizer)\n        self.assertEqual(decoded, [\"Hello\", \"Hello\"])\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\nfrom functools import partial\n\nimport torch\nfrom datasets import Dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, TrainingArguments\n\nfrom trl import IterativeSFTTrainer\n\n\nclass IterativeTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab-calibrated\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    def _init_tensor_dummy_dataset(self):\n        dummy_dataset_dict = {\n            \"input_ids\": [\n                torch.tensor([5303, 3621, 3666, 1438, 318]),\n                torch.tensor([3666, 1438, 318, 3666, 1438, 318]),\n                torch.tensor([5303, 3621, 3666, 1438, 318]),\n            ],\n            \"attention_mask\": [\n                torch.tensor([1, 1, 1, 1, 1]),\n                torch.tensor([1, 1, 1, 1, 1, 1]),\n                torch.tensor([1, 1, 1, 1, 1]),\n            ],\n            \"labels\": [\n                torch.tensor([5303, 3621, 3666, 1438, 318]),\n                torch.tensor([3666, 1438, 318, 3666, 1438, 318]),\n                torch.tensor([5303, 3621, 3666, 1438, 318]),\n            ],\n        }\n\n        dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n        dummy_dataset.set_format(\"torch\")\n        return dummy_dataset\n\n    def _init_textual_dummy_dataset(self):\n        dummy_dataset_dict = {\n            \"texts\": [\"Testing the IterativeSFTTrainer.\", \"This is a test of the IterativeSFTTrainer\"],\n            \"texts_labels\": [\"Testing the IterativeSFTTrainer.\", \"This is a test of the IterativeSFTTrainer\"],\n        }\n\n        dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n        dummy_dataset.set_format(\"torch\")\n        return dummy_dataset\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", \"tensor\"],\n            [\"gpt2\", \"text\"],\n            [\"t5\", \"tensor\"],\n            [\"t5\", \"text\"],\n        ]\n    )\n    def test_iterative_step_from_tensor(self, model_name, input_name):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            # initialize dataset\n            if input_name == \"tensor\":\n                dummy_dataset = self._init_tensor_dummy_dataset()\n                inputs = {\n                    \"input_ids\": dummy_dataset[\"input_ids\"],\n                    \"attention_mask\": dummy_dataset[\"attention_mask\"],\n                    \"labels\": dummy_dataset[\"labels\"],\n                }\n            else:\n                dummy_dataset = self._init_textual_dummy_dataset()\n                inputs = {\n                    \"texts\": dummy_dataset[\"texts\"],\n                    \"texts_labels\": dummy_dataset[\"texts_labels\"],\n                }\n\n            if model_name == \"gpt2\":\n                model = self.model\n                tokenizer = self.tokenizer\n            else:\n                model = self.t5_model\n                tokenizer = self.t5_tokenizer\n\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=2,\n                learning_rate=1e-3,\n                report_to=\"none\",\n            )\n            iterative_trainer = IterativeSFTTrainer(model=model, args=training_args, tokenizer=tokenizer)\n            iterative_trainer.optimizer.zero_grad = partial(iterative_trainer.optimizer.zero_grad, set_to_none=False)\n\n            iterative_trainer.step(**inputs)\n\n            for param in iterative_trainer.model.parameters():\n                assert param.grad is not None\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nimport torch\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer\nfrom transformers.testing_utils import require_peft\n\nfrom trl import ORPOConfig, ORPOTrainer\n\n\nclass ORPOTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    @parameterized.expand([[\"gpt2\"], [\"t5\"]])\n    def test_orpo_trainer(self, name):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = ORPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            if name == \"gpt2\":\n                model = self.model\n                tokenizer = self.tokenizer\n            elif name == \"t5\":\n                model = self.t5_model\n                tokenizer = self.t5_tokenizer\n                training_args.is_encoder_decoder = True\n\n            trainer = ORPOTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.equal(param, new_param)\n\n    @require_peft\n    def test_orpo_trainer_with_lora(self):\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = ORPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            trainer = ORPOTrainer(\n                model=self.model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                if \"lora\" in n:\n                    new_param = trainer.model.get_parameter(n)\n                    # check the params have changed - ignore 0 biases\n                    if param.sum() != 0:\n                        assert not torch.equal(param, new_param)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport itertools\nimport unittest\n\nfrom datasets import Dataset, DatasetDict\nfrom parameterized import parameterized\nfrom transformers import AutoTokenizer\n\nfrom trl.data_utils import (\n    apply_chat_template,\n    extract_prompt,\n    is_conversational,\n    maybe_apply_chat_template,\n    maybe_extract_prompt,\n    maybe_unpair_preference_dataset,\n    unpair_preference_dataset,\n)\n\n\nclass IsConversationalTester(unittest.TestCase):\n    conversational_examples = [\n        {  # Language modeling\n            \"messages\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n            ],\n        },\n        {  # Prompt only\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n        },\n        {  # Pompt-completion\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"completion\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n        },\n        {  # Preference\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"chosen\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n            \"rejected\": [{\"role\": \"assistant\", \"content\": \"It is green.\"}],\n        },\n        {  # Preference with implicit prompt\n            \"chosen\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n            ],\n            \"rejected\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is green.\"},\n            ],\n        },\n        {  # Unpaired preference\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"completion\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n            \"label\": True,\n        },\n    ]\n\n    non_conversational_examples = [\n        {\"prompt\": \"The sky is\", \"completion\": \" blue.\"},\n        {\"text\": \"The sky is blue.\"},\n        {\"prompt\": \"The sky is\"},\n        {\"prompt\": \"The sky is\", \"chosen\": \" blue.\", \"rejected\": \" green.\"},\n        {\"prompt\": \"The sky is\", \"completion\": \" blue.\", \"label\": True},\n    ]\n\n    @parameterized.expand(itertools.product(conversational_examples))\n    def test_conversational(self, example):\n        self.assertTrue(is_conversational(example))\n\n    @parameterized.expand(itertools.product(non_conversational_examples))\n    def test_non_conversational(self, example):\n        self.assertFalse(is_conversational(example))\n\n\nclass ApplyChatTemplateTester(unittest.TestCase):\n    tokenizers = [\n        \"trl-internal-testing/tiny-random-Qwen2-7B-Instruct\",\n        \"trl-internal-testing/tiny-random-Meta-Llama-3.1-8B-Instruct\",\n        \"trl-internal-testing/tiny-random-Meta-Llama-3-8B-Instruct\",\n        \"trl-internal-testing/tiny-random-DeepSeek-Coder-V2-Instruct\",\n        \"trl-internal-testing/tiny-random-Phi-3-mini-128k-instruct\",\n        \"trl-internal-testing/tiny-random-gemma-2-9b-it\",\n        \"trl-internal-testing/tiny-random-Mistral-7B-Instruct-v0.1\",\n        \"trl-internal-testing/tiny-random-Mistral-7B-Instruct-v0.2\",\n        \"trl-internal-testing/tiny-random-Mistral-7B-Instruct-v0.3\",\n    ]\n\n    conversational_examples = [\n        {  # Language modeling\n            \"messages\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n            ],\n        },\n        {  # Prompt only\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n        },\n        {  # Pompt-completion\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"completion\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n        },\n        {  # Preference\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"chosen\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n            \"rejected\": [{\"role\": \"assistant\", \"content\": \"It is green.\"}],\n        },\n        {  # Preference with implicit prompt\n            \"chosen\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n            ],\n            \"rejected\": [\n                {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"It is green.\"},\n            ],\n        },\n        {  # Unpaired preference\n            \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n            \"completion\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n            \"label\": True,\n        },\n    ]\n\n    non_conversational_examples = [\n        {\"prompt\": \"The sky is\", \"completion\": \" blue.\"},\n        {\"text\": \"The sky is blue.\"},\n        {\"prompt\": \"The sky is\"},\n        {\"prompt\": \"The sky is\", \"chosen\": \" blue.\", \"rejected\": \" green.\"},\n        {\"chosen\": \"The sky is blue.\", \"rejected\": \"The sky is green.\"},\n        {\"prompt\": \"The sky is\", \"completion\": \" blue.\", \"label\": True},\n    ]\n\n    @parameterized.expand(itertools.product(tokenizers, conversational_examples))\n    def test_apply_chat_template(self, tokenizer_id, example):\n        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)\n        result = apply_chat_template(example, tokenizer)\n\n        # Checking if the result is a dictionary\n        self.assertIsInstance(result, dict)\n\n        # The chat template should be applied to the the following keys\n        for key in [\"prompt\", \"chosen\", \"rejected\", \"completion\"]:\n            if key in example:\n                self.assertIn(key, result)\n                self.assertIsInstance(result[key], str)\n\n        # Exception for messages, the key is \"text\" once the chat template is applied\n        if \"messages\" in example:\n            self.assertIn(\"text\", result)\n            self.assertIsInstance(result[\"text\"], str)\n\n        # The label should be kept\n        if \"label\" in example:\n            self.assertIn(\"label\", result)\n            self.assertIsInstance(result[\"label\"], bool)\n            self.assertEqual(result[\"label\"], example[\"label\"])\n\n    # both conversational and non-conversational examples\n    @parameterized.expand(itertools.product(tokenizers, conversational_examples + non_conversational_examples))\n    def test_maybe_apply_chat_template(self, tokenizer_id, example):\n        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)\n        result = maybe_apply_chat_template(example, tokenizer)\n\n        # Checking if the result is a dictionary\n        self.assertIsInstance(result, dict)\n\n        # The chat template should be applied to the the following keys\n        for key in [\"prompt\", \"chosen\", \"rejected\", \"completion\"]:\n            if key in example:\n                self.assertIn(key, result)\n                self.assertIsInstance(result[key], str)\n\n        # Exception for messages, the key is \"text\" once the chat template is applied\n        if \"messages\" in example:\n            self.assertIn(\"text\", result)\n            self.assertIsInstance(result[\"text\"], str)\n\n        # The label should be kept\n        if \"label\" in example:\n            self.assertIn(\"label\", result)\n            self.assertIsInstance(result[\"label\"], bool)\n            self.assertEqual(result[\"label\"], example[\"label\"])\n\n\nclass UnpairPreferenceDatasetTester(unittest.TestCase):\n    paired_dataset = Dataset.from_dict(\n        {\n            \"prompt\": [\"The sky is\", \"The sun is\"],\n            \"chosen\": [\" blue.\", \" in the sky.\"],\n            \"rejected\": [\" green.\", \" in the sea.\"],\n        }\n    )\n\n    unpaired_dataset = Dataset.from_dict(\n        {\n            \"prompt\": [\"The sky is\", \"The sun is\", \"The sky is\", \"The sun is\"],\n            \"completion\": [\" blue.\", \" in the sky.\", \" green.\", \" in the sea.\"],\n            \"label\": [True, True, False, False],\n        }\n    )\n\n    def test_unpair_preference_dataset(self):\n        # Test that a paired-formatted dataset is correctly converted to unpaired format\n        unpaired_dataset = unpair_preference_dataset(self.paired_dataset)\n        self.assertEqual(\n            unpaired_dataset.to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The paired-formatted dataset should be reformatted to unpaired format.\",\n        )\n\n    def test_unpair_preference_dataset_dict(self):\n        # Test that a paired-formatted dataset dict is correctly converted to unpaired format\n        paired_dataset_dict = DatasetDict({\"abc\": self.paired_dataset})\n        unpaired_dataset_dict = unpair_preference_dataset(paired_dataset_dict)\n        self.assertEqual(\n            unpaired_dataset_dict[\"abc\"].to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The paired-formatted dataset should be reformatted to unpaired format.\",\n        )\n\n    def test_maybe_unpair_preference_dataset(self):\n        # Test that a paired-formatted dataset is correctly reformatted to unpaired format with maybe_unpair_preference_dataset\n        unpaired_dataset = maybe_unpair_preference_dataset(self.paired_dataset)\n        self.assertEqual(\n            unpaired_dataset.to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The paired-formatted dataset should be reformatted to unpaired format.\",\n        )\n\n    def test_maybe_unpair_preference_dataset_dict(self):\n        # Test that a paired-formatted dataset dict is correctly converted to unpaired format with maybe_unpair_preference_dataset\n        paired_dataset_dict = DatasetDict({\"abc\": self.paired_dataset})\n        unpaired_dataset_dict = maybe_unpair_preference_dataset(paired_dataset_dict)\n        self.assertEqual(\n            unpaired_dataset_dict[\"abc\"].to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The paired-formatted dataset should be reformatted to unpaired format.\",\n        )\n\n    def test_maybe_unpair_preference_dataset_already_paired(self):\n        # Test that a paired-formatted dataset remains unchanged with maybe_unpair_preference_dataset\n        unpaired_dataset = maybe_unpair_preference_dataset(self.unpaired_dataset)\n        self.assertEqual(\n            unpaired_dataset.to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The unpaired-formatted dataset should remain unchanged.\",\n        )\n\n    def test_maybe_unpair_preference_dataset_dict_already_paired(self):\n        # Test that a paired-formatted dataset dict remains unchanged with maybe_unpair_preference_dataset\n        unpaired_dataset_dict = maybe_unpair_preference_dataset(DatasetDict({\"abc\": self.unpaired_dataset}))\n        self.assertEqual(\n            unpaired_dataset_dict[\"abc\"].to_dict(),\n            self.unpaired_dataset.to_dict(),\n            \"The unpaired-formatted dataset should remain unchanged.\",\n        )\n\n\nclass ExtractPromptTester(unittest.TestCase):\n    example_implicit_prompt = {\n        \"chosen\": [\n            {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n            {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n        ],\n        \"rejected\": [\n            {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n            {\"role\": \"assistant\", \"content\": \"It is green.\"},\n        ],\n    }\n\n    example_explicit_prompt = {\n        \"prompt\": [\n            {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n        ],\n        \"chosen\": [\n            {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n        ],\n        \"rejected\": [\n            {\"role\": \"assistant\", \"content\": \"It is green.\"},\n        ],\n    }\n\n    def test_extract_prompt(self):\n        # Test that the prompt is correctly extracted from the dataset\n        example_extracted_prompt = extract_prompt(self.example_implicit_prompt)\n        self.assertEqual(\n            example_extracted_prompt,\n            self.example_explicit_prompt,\n            \"The prompt is not correctly extracted from the dataset.\",\n        )\n\n    def test_maybe_extract_prompt(self):\n        # Test that the prompt is correctly extracted from the dataset with maybe_extract_prompt\n        example_extracted_prompt = maybe_extract_prompt(self.example_implicit_prompt)\n        self.assertEqual(\n            example_extracted_prompt,\n            self.example_explicit_prompt,\n            \"The prompt is not correctly extracted from the dataset.\",\n        )\n\n    def test_maybe_extract_prompt_already_explicit(self):\n        # Test that the prompt remains unchanged with maybe_extract_prompt\n        example_extracted_prompt = maybe_extract_prompt(self.example_explicit_prompt)\n        self.assertEqual(\n            example_extracted_prompt,\n            self.example_explicit_prompt,\n            \"The prompt should remain unchanged.\",\n        )\n\n\n# Run the tests\nif __name__ == \"__main__\":\n    unittest.main()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport subprocess\n\n\ndef test_hello_world():\n    subprocess.run(\n        \"python examples/hello_world.py\",\n        shell=True,\n        check=True,\n    )\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nimport torch\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer\nfrom transformers.testing_utils import require_peft\n\nfrom trl import CPOConfig, CPOTrainer\n\n\nclass CPOTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", \"sigmoid\"],\n            [\"t5\", \"hinge\"],\n            [\"gpt2\", \"ipo\"],\n            [\"t5\", \"ipo\"],\n            [\"gpt2\", \"simpo\"],\n            [\"t5\", \"simpo\"],\n        ]\n    )\n    def test_cpo_trainer(self, name, loss_type):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = CPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                loss_type=loss_type,\n                cpo_alpha=1.0,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            if name == \"gpt2\":\n                model = self.model\n                tokenizer = self.tokenizer\n            elif name == \"t5\":\n                model = self.t5_model\n                tokenizer = self.t5_tokenizer\n                training_args.is_encoder_decoder = True\n\n            trainer = CPOTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.equal(param, new_param)\n\n    @require_peft\n    def test_cpo_trainer_with_lora(self):\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = CPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                cpo_alpha=1.0,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_preference\")\n\n            trainer = CPOTrainer(\n                model=self.model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[-1][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                if \"lora\" in n:\n                    new_param = trainer.model.get_parameter(n)\n                    # check the params have changed - ignore 0 biases\n                    if param.sum() != 0:\n                        assert not torch.equal(param, new_param)\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport platform\nimport subprocess\n\nimport torch\n\n\ndef test():\n    command = \"\"\"\\\npython examples/scripts/rloo/rloo.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/rloo \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10 \\\n    --model_name_or_path EleutherAI/pythia-14m \\\n    --missing_eos_penalty 1.0 \\\n    --save_strategy no \\\n    --stop_token eos\n\"\"\"\n    if platform.system() == \"Windows\":\n        # windows CI does not work with subprocesses for some reason\n        # e.g., https://github.com/huggingface/trl/actions/runs/9600036224/job/26475286210?pr=1743\n        return\n    subprocess.run(\n        command,\n        shell=True,\n        check=True,\n    )\n\n\ndef test_rloo_reward():\n    local_batch_size = 3\n    rloo_k = 4\n    # fmt: off\n    rlhf_reward = torch.tensor([\n        1, 2, 3, # first rlhf reward for three prompts\n        2, 3, 4, # second rlhf reward for three prompts\n        5, 6, 7, # third rlhf reward for three prompts\n        8, 9, 10, # fourth rlhf reward for three prompts\n    ]).float()\n    # fmt: on\n\n    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)\n    advantages = torch.zeros_like(rlhf_reward)\n    for i in range(0, len(advantages), local_batch_size):\n        other_response_rlhf_rewards = []\n        for j in range(0, len(advantages), local_batch_size):\n            if i != j:\n                other_response_rlhf_rewards.append(rlhf_reward[j : j + local_batch_size])\n        advantages[i : i + local_batch_size] = rlhf_reward[i : i + local_batch_size] - torch.stack(\n            other_response_rlhf_rewards\n        ).mean(0)\n    assert (1 - (2 + 5 + 8) / 3 - advantages[0].item()) < 1e-6\n    assert (6 - (3 + 2 + 9) / 3 - advantages[7].item()) < 1e-6\n\n    # vectorized impl\n    rlhf_reward = rlhf_reward.reshape(rloo_k, local_batch_size)\n    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)\n    vec_advantages = rlhf_reward - baseline\n    torch.testing.assert_close(vec_advantages.flatten(), advantages)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\nfrom transformers.testing_utils import require_peft\nfrom transformers.utils import is_peft_available\n\nfrom trl import OnlineDPOConfig, OnlineDPOTrainer\n\n\nif is_peft_available():\n    from peft import LoraConfig, get_peft_model\n\n\nclass TestOnlineDPOTrainer(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.reward_model = AutoModelForSequenceClassification.from_pretrained(\"EleutherAI/pythia-14m\", num_labels=1)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n    @parameterized.expand([(\"standard_prompt_only\",), (\"conversational_prompt_only\",)])\n    def test_training(self, config_name):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", config_name)\n\n            trainer = OnlineDPOTrainer(\n                model=self.model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    def test_training_with_ref_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = OnlineDPOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    def test_ref_model_is_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            with self.assertRaises(ValueError):\n                OnlineDPOTrainer(\n                    model=self.model,\n                    ref_model=self.model,  # ref_model can't be the same as model\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n    @require_peft\n    def test_training_with_peft(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = OnlineDPOTrainer(\n                model=self.model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft_and_ref_model(self):\n        lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = OnlineDPOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n    @require_peft\n    def test_training_with_peft_model_and_peft_config(self):\n        model_lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.1, bias=\"none\", task_type=\"CAUSAL_LM\")\n        model = get_peft_model(self.model, model_lora_config)\n        # we want only the \"train adapter\" to be trained\n        lora_train_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\")\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = OnlineDPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                learning_rate=5.0e-7,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_prompt_only\")\n\n            trainer = OnlineDPOTrainer(\n                model=model,\n                reward_model=self.reward_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_train_config,\n            )\n\n            trainer.train()\n\n            # Check if training loss is available\n            self.assertIn(\"train_loss\", trainer.state.log_history[-1])\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nimport torch\nfrom transformers import AutoTokenizer, GenerationConfig\n\nfrom trl import AutoModelForCausalLMWithValueHead\nfrom trl.core import LengthSampler\nfrom trl.extras import BestOfNSampler\n\n\ndef queries_to_scores(list_of_strings):\n    return [torch.rand(1).item() for _ in list_of_strings]\n\n\nclass BestOfNSamplerTester(unittest.TestCase):\n    \"\"\"\n    Tests the BestOfNSampler class\n    \"\"\"\n\n    ref_model_name = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n    output_length_sampler = LengthSampler(2, 6)\n    model = AutoModelForCausalLMWithValueHead.from_pretrained(ref_model_name)\n    tokenizer = AutoTokenizer.from_pretrained(ref_model_name)\n    tokenizer.pad_token = tokenizer.eos_token\n    output_length_sampler = LengthSampler(2, 6)\n\n    def test_different_input_types(self):\n        r\"\"\"\n        Tests if the different input types normalizer works\n        \"\"\"\n\n        generation_config = GenerationConfig(\n            min_length=-1,\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n            pad_token_id=self.tokenizer.eos_token_id,\n        )\n\n        output_length_sampler = LengthSampler(2, 6)\n\n        best_of_n = BestOfNSampler(\n            self.model,\n            self.tokenizer,\n            queries_to_scores,\n            length_sampler=output_length_sampler,\n            generation_config=generation_config,\n        )\n\n        queries = [\"hello world\", \"goodbye world\"]\n        tokenized_queries = [self.tokenizer.encode(query) for query in queries]\n\n        various_queries_formats = [\n            (tokenized_queries[0], 1),\n            (tokenized_queries, 2),\n            (torch.tensor(tokenized_queries[1]), 1),\n            ([torch.tensor(query) for query in tokenized_queries], 2),\n        ]\n\n        for q, expected_length in various_queries_formats:\n            results = best_of_n.generate(q)\n            assert isinstance(results, list)\n            assert len(results) == expected_length\n\n    def test_different_sample_sizes_and_n_candidates_values(self):\n        r\"\"\"\n        Tests different sample sizes and n_candidates values\n        \"\"\"\n        generation_config = GenerationConfig(\n            min_length=-1,\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n            pad_token_id=self.tokenizer.eos_token_id,\n        )\n\n        output_length_sampler = LengthSampler(6, 10)\n\n        for sample_value, n_candidates_values, expected in [\n            (4, 2, 2),\n            (10, 3, 3),\n            (6, 4, 4),\n        ]:\n            best_of_n = BestOfNSampler(\n                self.model,\n                self.tokenizer,\n                queries_to_scores,\n                length_sampler=output_length_sampler,\n                generation_config=generation_config,\n                sample_size=sample_value,\n                n_candidates=n_candidates_values,\n            )\n\n            queries = [\"hello world\", \"troll the world\"]\n            tokenized_queries = [self.tokenizer.encode(query) for query in queries]\n            results = best_of_n.generate(tokenized_queries)\n            for result in results:\n                assert len(result) == expected\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\nfrom functools import partial\n\nimport torch\nfrom accelerate import Accelerator\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer\nfrom transformers.testing_utils import require_peft\n\nfrom trl import BCOConfig, BCOTrainer\nfrom trl.trainer.bco_trainer import _process_tokens, _tokenize\n\nfrom .testing_utils import require_no_wandb\n\n\nclass BCOTrainerTester(unittest.TestCase):\n    def setUp(self):\n        self.model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n\n        # get t5 as seq2seq example:\n        model_id = \"trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab\"\n        self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_ref_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)\n        self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        # get embedding model\n        model_id = \"facebook/bart-base\"\n        self.embedding_model = AutoModel.from_pretrained(model_id)\n        self.embedding_tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n    @parameterized.expand(\n        [\n            [\"gpt2\", True, True],\n            [\"gpt2\", True, False],\n            [\"gpt2\", False, True],\n            [\"gpt2\", False, False],\n        ]\n    )\n    def test_bco_trainer(self, name, pre_compute, eval_dataset):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                precompute_ref_log_probs=pre_compute,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            if name == \"gpt2\":\n                model = self.model\n                ref_model = self.ref_model\n                tokenizer = self.tokenizer\n            elif name == \"t5\":\n                model = self.t5_model\n                ref_model = self.t5_ref_model\n                tokenizer = self.t5_tokenizer\n\n            trainer = BCOTrainer(\n                model=model,\n                ref_model=ref_model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"] if eval_dataset else None,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    self.assertFalse(torch.equal(param.cpu(), new_param.cpu()))\n\n    def test_bco_trainer_with_ref_model_is_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            with self.assertRaises(ValueError):\n                BCOTrainer(\n                    model=self.model,\n                    ref_model=self.model,  # ref_model can't be the same as model\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                )\n\n    def test_tokenize_and_process_tokens(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = BCOTrainer(\n                model=self.model,\n                ref_model=self.ref_model,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            train_dataset = dummy_dataset[\"train\"]\n            tokenized_dataset = train_dataset.map(\n                _tokenize,\n                fn_kwargs={\"tokenizer\": trainer.tokenizer},\n                batched=True,\n                batch_size=2,\n            )\n            self.assertListEqual(tokenized_dataset[\"prompt\"], train_dataset[\"prompt\"])\n            self.assertListEqual(tokenized_dataset[\"completion\"], train_dataset[\"completion\"])\n            self.assertListEqual(tokenized_dataset[\"label\"], train_dataset[\"label\"])\n            self.assertListEqual(tokenized_dataset[\"prompt_input_ids\"][0], [5377, 11141])\n            self.assertListEqual(tokenized_dataset[\"prompt_attention_mask\"][0], [1, 1])\n            self.assertListEqual(tokenized_dataset[\"answer_input_ids\"][0], [318, 1365, 621, 8253, 13])\n            self.assertListEqual(tokenized_dataset[\"answer_attention_mask\"][0], [1, 1, 1, 1, 1])\n\n            fn_kwargs = {\n                \"prefix\": \"\",\n                \"is_encoder_decoder\": trainer.is_encoder_decoder,\n                \"tokenizer\": trainer.tokenizer,\n                \"max_length\": trainer.max_length,\n                \"truncation_mode\": trainer.truncation_mode,\n                \"label_pad_token_id\": trainer.label_pad_token_id,\n                \"max_prompt_length\": trainer.max_prompt_length,\n            }\n            processed_dataset = tokenized_dataset.map(_process_tokens, fn_kwargs=fn_kwargs, num_proc=2)\n            self.assertListEqual(processed_dataset[\"prompt\"], train_dataset[\"prompt\"])\n            self.assertListEqual(processed_dataset[\"completion\"], train_dataset[\"completion\"])\n            self.assertListEqual(processed_dataset[\"label\"], train_dataset[\"label\"])\n            self.assertListEqual(processed_dataset[\"prompt_input_ids\"][0], [50256, 5377, 11141])\n            self.assertListEqual(processed_dataset[\"prompt_attention_mask\"][0], [1, 1, 1])\n            self.assertListEqual(\n                processed_dataset[\"completion_input_ids\"][0], [50256, 5377, 11141, 318, 1365, 621, 8253, 13, 50256]\n            )\n            self.assertListEqual(processed_dataset[\"completion_attention_mask\"][0], [1, 1, 1, 1, 1, 1, 1, 1, 1])\n            self.assertListEqual(\n                processed_dataset[\"completion_labels\"][0], [-100, -100, -100, 318, 1365, 621, 8253, 13, 50256]\n            )\n\n    def test_bco_trainer_without_providing_ref_model(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = BCOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    self.assertFalse(torch.equal(param.cpu(), new_param.cpu()))\n\n    def test_bco_trainer_udm(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            def embed_prompt(input_ids, attention_mask, model):\n                outputs = model(input_ids=input_ids, attention_mask=attention_mask)\n\n                return outputs.last_hidden_state.mean(dim=1)\n\n            embedding_model = Accelerator().prepare_model(self.embedding_model)\n            embedding_func = partial(embed_prompt, model=embedding_model)\n\n            trainer = BCOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                embedding_func=embedding_func,\n                embedding_tokenizer=self.embedding_tokenizer,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    self.assertFalse(torch.equal(param.cpu(), new_param.cpu()))\n\n    @require_peft\n    def test_bco_trainer_without_providing_ref_model_with_lora(self):\n        from peft import LoraConfig\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            trainer = BCOTrainer(\n                model=self.model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            self.assertIsNotNone(trainer.state.log_history[-1][\"train_loss\"])\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                if \"lora\" in n:\n                    new_param = trainer.model.get_parameter(n)\n                    # check the params have changed - ignore 0 biases\n                    if param.sum() != 0:\n                        self.assertFalse(torch.equal(param.cpu(), new_param.cpu()))\n\n    @require_no_wandb\n    def test_bco_trainer_generate_during_eval_no_wandb(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=1,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                generate_during_eval=True,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            with self.assertRaisesRegex(\n                ValueError,\n                expected_regex=\"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install with `pip install wandb` to resolve.\",\n            ):\n                BCOTrainer(\n                    model=self.model,\n                    ref_model=None,\n                    args=training_args,\n                    tokenizer=self.tokenizer,\n                    train_dataset=dummy_dataset[\"train\"],\n                    eval_dataset=dummy_dataset[\"test\"],\n                )\n\n    @require_peft\n    def test_bco_lora_save(self):\n        from peft import LoraConfig, get_peft_model\n\n        lora_config = LoraConfig(\n            r=16,\n            lora_alpha=32,\n            lora_dropout=0.05,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n        # lora model\n        model = AutoModelForCausalLM.from_pretrained(self.model_id)\n        model_peft = get_peft_model(model, lora_config)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = BCOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                beta=0.1,\n                report_to=\"none\",\n            )\n\n            dummy_dataset = load_dataset(\"trl-internal-testing/zen\", \"standard_unpaired_preference\")\n\n            # bco train lora model with a lora config\n            trainer = BCOTrainer(\n                model=model_peft,\n                ref_model=None,\n                args=training_args,\n                tokenizer=self.tokenizer,\n                train_dataset=dummy_dataset[\"train\"],\n                eval_dataset=dummy_dataset[\"test\"],\n                peft_config=lora_config,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save peft adapter\n            trainer.save_model()\n\n            # assert that the model is loaded without giving OSError\n            try:\n                AutoModelForCausalLM.from_pretrained(tmp_dir)\n            except OSError:\n                self.fail(\"Loading the saved peft adapter failed\")\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nCI_HUB_USER = \"__DUMMY_TRANSFORMERS_USER__\"\nCI_HUB_USER_FULL_NAME = \"Dummy User\"\n\nCI_HUB_ENDPOINT = \"https://hub-ci.huggingface.co\"\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport platform\nimport subprocess\n\n\ndef test():\n    command = \"\"\"\\\npython examples/scripts/ppo/ppo.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10 \\\n    --model_name_or_path EleutherAI/pythia-14m \\\n    --missing_eos_penalty 1.0 \\\n    --save_strategy no \\\n    --stop_token eos\n\"\"\"\n    if platform.system() == \"Windows\":\n        # windows CI does not work with subprocesses for some reason\n        # e.g., https://github.com/huggingface/trl/actions/runs/9600036224/job/26475286210?pr=1743\n        return\n    subprocess.run(\n        command,\n        shell=True,\n        check=True,\n    )\n\n\ndef test_num_train_epochs():\n    command = \"\"\"\\\npython examples/scripts/ppo/ppo.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 1 \\\n    --num_train_epochs 0.003 \\\n    --model_name_or_path EleutherAI/pythia-14m \\\n    --missing_eos_penalty 1.0 \\\n    --save_strategy no \\\n    --stop_token eos\n\"\"\"\n    if platform.system() == \"Windows\":\n        # windows CI does not work with subprocesses for some reason\n        # e.g., https://github.com/huggingface/trl/actions/runs/9600036224/job/26475286210?pr=1743\n        return\n    subprocess.run(\n        command,\n        shell=True,\n        check=True,\n    )\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport tempfile\nimport unittest\n\nimport pytest\nimport torch\nfrom datasets import Dataset\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, EvalPrediction\nfrom transformers.testing_utils import require_peft\n\nfrom trl import RewardConfig, RewardTrainer\nfrom trl.trainer import compute_accuracy\n\n\nclass RewardTrainerTester(unittest.TestCase):\n    def test_accuracy_metrics(self):\n        dummy_eval_predictions = EvalPrediction(torch.FloatTensor([[0.1, 0.9], [0.9, 0.1]]), torch.LongTensor([0, 0]))\n        accuracy = compute_accuracy(dummy_eval_predictions)\n        assert accuracy[\"accuracy\"] == 0.5\n\n    def test_reward_trainer(self):\n        model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        model = AutoModelForSequenceClassification.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n        tokenizer.pad_token = tokenizer.eos_token\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n\n            # fmt: off\n            dummy_dataset_dict = {\n                \"input_ids_chosen\": [\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                ],\n                \"attention_mask_chosen\": [\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                ],\n                \"input_ids_rejected\": [\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                ],\n                \"attention_mask_rejected\": [\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 0]),\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 1]),\n                ],\n            }\n            # fmt: on\n            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n\n            trainer = RewardTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset,\n                eval_dataset=dummy_dataset,\n            )\n\n            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                # check the params have changed - ignore 0 biases\n                if param.sum() != 0:\n                    assert not torch.equal(param, new_param)\n\n            preds = trainer.predict(dummy_dataset)\n            assert preds.predictions.shape == (4, 2)\n\n    @require_peft\n    def test_reward_trainer_peft(self):\n        from peft import LoraConfig, TaskType\n\n        model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        model = AutoModelForSequenceClassification.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n        tokenizer.pad_token = tokenizer.eos_token\n\n        peft_config = LoraConfig(\n            task_type=TaskType.SEQ_CLS,\n            inference_mode=False,\n            r=8,\n            lora_alpha=32,\n            lora_dropout=0.1,\n        )\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=6,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=2,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n\n            # fmt: off\n            dummy_dataset_dict = {\n                \"input_ids_chosen\": [\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                ],\n                \"attention_mask_chosen\": [\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                ],\n                \"input_ids_rejected\": [\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                ],\n                \"attention_mask_rejected\": [\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 0]),\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 1]),\n                ],\n            }\n            # fmt: on\n            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n\n            trainer = RewardTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset,\n                eval_dataset=dummy_dataset,\n                peft_config=peft_config,\n            )\n            previous_trainable_params = {}\n            previous_non_trainable_params = {}\n\n            # due to a change in the way the modules to save are dealt in PEFT.\n            trainable_params_name = [\"lora\", \"modules_to_save\"]\n\n            # check gradients are not None\n            for n, param in trainer.model.named_parameters():\n                if any(t in n for t in trainable_params_name):\n                    previous_trainable_params[n] = param.clone()\n                else:\n                    previous_non_trainable_params[n] = param.clone()\n\n            trainer.train()\n\n            assert trainer.state.log_history[(-1)][\"train_loss\"] is not None\n\n            # check the params have changed\n            for n, param in previous_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                assert not torch.allclose(param, new_param, atol=1e-12, rtol=1e-12)\n\n            # check the non trainable params have not changed\n            for n, param in previous_non_trainable_params.items():\n                new_param = trainer.model.get_parameter(n)\n                assert torch.allclose(param, new_param, atol=1e-12, rtol=1e-12)\n\n            preds = trainer.predict(dummy_dataset)\n            assert preds.predictions.shape == (4, 2)\n\n    def test_reward_trainer_assert_value_error(self):\n        model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        model = AutoModelForSequenceClassification.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n        tokenizer.pad_token = tokenizer.eos_token\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=1,\n                remove_unused_columns=False,\n                report_to=\"none\",\n            )\n\n            # fmt: off\n            dummy_dataset_dict = {\n                \"input_ids_b\": [\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                ],\n                \"attention_mask_c\": [\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                ],\n                \"input_ids_f\": [\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                ],\n                \"attention_mask_g\": [\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 0]),\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 1]),\n                ],\n            }\n            # fmt: on\n            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n\n            trainer = RewardTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset,\n            )\n\n            with pytest.raises(ValueError):\n                trainer.train()\n\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=1,\n                remove_unused_columns=True,\n                report_to=\"none\",\n            )\n\n            with self.assertWarns(UserWarning):\n                trainer = RewardTrainer(\n                    model=model,\n                    args=training_args,\n                    tokenizer=tokenizer,\n                    train_dataset=dummy_dataset,\n                )\n\n    def test_reward_trainer_margin(self):\n        model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        model = AutoModelForSequenceClassification.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n        tokenizer.pad_token = tokenizer.eos_token\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n\n            # fmt: off\n            dummy_dataset_dict = {\n                \"input_ids_chosen\": [\n                    torch.LongTensor([0, 1, 2]),\n                ],\n                \"attention_mask_chosen\": [\n                    torch.LongTensor([1, 1, 1]),\n                ],\n                \"input_ids_rejected\": [\n                    torch.LongTensor([0, 2]),\n                ],\n                \"attention_mask_rejected\": [\n                    torch.LongTensor([1, 1]),\n                ],\n                \"margin\": [\n                    torch.FloatTensor([1.0]),\n                ]\n            }\n            # fmt: on\n            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n\n            trainer = RewardTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset,\n                eval_dataset=dummy_dataset,\n            )\n\n            batch = [dummy_dataset[0]]\n            batch = trainer.data_collator(batch)\n            batch = {k: v.to(trainer.model.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n            loss, outputs = trainer.compute_loss(trainer.model, batch, return_outputs=True)\n\n            l_val = -torch.nn.functional.logsigmoid(\n                outputs[\"rewards_chosen\"] - outputs[\"rewards_rejected\"] - batch[\"margin\"]\n            ).mean()\n\n            assert abs(loss - l_val) < 1e-6\n\n    def test_reward_trainer_tags(self):\n        model_id = \"trl-internal-testing/dummy-GPT2-correct-vocab\"\n        model = AutoModelForSequenceClassification.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n        tokenizer.pad_token = tokenizer.eos_token\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = RewardConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=3,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=4,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                report_to=\"none\",\n            )\n\n            # fmt: off\n            dummy_dataset_dict = {\n                \"input_ids_chosen\": [\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                    torch.LongTensor([0, 1, 2]),\n                    torch.LongTensor([1, 2]),\n                ],\n                \"attention_mask_chosen\": [\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                    torch.LongTensor([1, 1, 1]),\n                    torch.LongTensor([1, 0]),\n                ],\n                \"input_ids_rejected\": [\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                    torch.LongTensor([0, 2]),\n                    torch.LongTensor([1, 2, 0]),\n                ],\n                \"attention_mask_rejected\": [\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 0]),\n                    torch.LongTensor([1, 1]),\n                    torch.LongTensor([1, 1, 1]),\n                ],\n            }\n            # fmt: on\n            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)\n\n            trainer = RewardTrainer(\n                model=model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=dummy_dataset,\n                eval_dataset=dummy_dataset,\n            )\n\n            assert trainer.model.model_tags == trainer._tag_names\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport unittest\n\nimport torch\nfrom transformers import AutoTokenizer\n\nfrom trl import DataCollatorForCompletionOnlyLM\n\n\nclass DataCollatorForCompletionOnlyLMTester(unittest.TestCase):\n    def test_data_collator_finds_response_template_llama2_tokenizer(self):\n        # this should ideally be tested with meta-llama/Llama-2-7b-hf\n        self.tokenizer = AutoTokenizer.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        self.instruction = \"\"\"### System: You are a helpful assistant.\n\n### User: How much is 2+2?\n\n### Assistant: 2+2 equals 4\"\"\"\n        self.instruction_template = \"\\n### User:\"\n        self.response_template = \"\\n### Assistant:\"\n\n        # GPT2Tokenizer: [198, 21017, 11787, 25] -> [21017, 11787, 25]\n        # Llama2Tokenizer: [29871, 13, 2277, 29937, 4911, 29901] -> [2277, 29937, 4911, 29901]\n        # Note: If this test is ever switched to Llama2Tokenizer, this should be double checked,\n        # and possibly switched back to [2:] instead of [1:].\n        # With GPT2Tokenizer, [1:] is correct - we want the 21017 token included, which is ###.\n        self.tokenized_instruction_w_context = self.tokenizer.encode(\n            self.instruction_template, add_special_tokens=False\n        )[1:]\n\n        # GPT2Tokenizer: [198, 21017, 15286, 25] -> [15286, 25]\n        # Llama2Tokenizer: [29871, 13, 2277, 29937, 4007, 22137, 29901] -> [2277, 29937, 4007, 22137, 29901]\n        self.tokenized_response_w_context = self.tokenizer.encode(self.response_template, add_special_tokens=False)[2:]\n\n        # Plain check on string\n        assert self.response_template in self.instruction\n        self.tokenized_instruction = self.tokenizer.encode(self.instruction, add_special_tokens=False)\n\n        # Test the fix for #598\n        # Pass already tokenized (w context) and truncated response_template so token_ids are like in the instruction + response\n        self.collator = DataCollatorForCompletionOnlyLM(self.tokenized_response_w_context, tokenizer=self.tokenizer)\n        self.collator.torch_call([self.tokenized_instruction])\n\n        # Test for PR #749\n        # Pass already tokenized (w context) instruction and response both so token_ids are like in the instruction + response\n        self.collator = DataCollatorForCompletionOnlyLM(\n            self.tokenized_response_w_context, self.tokenized_instruction_w_context, tokenizer=self.tokenizer\n        )\n        self.collator.torch_call([self.tokenized_instruction])\n\n        # Test for PR #1185\n        # We pass in a string where the first user template is different than the rest.\n        # Usually this would happen due to context-sensitive tokenization, but here we\n        # explicitly change the template to test the fix.\n        self.instruction = \"\"\"## User: First instruction\n\n### Assistant: First response\n\n### User: Second instruction\n\n### Assistant: Second response\"\"\"\n        self.tokenized_instruction = self.tokenizer.encode(self.instruction, add_special_tokens=False)\n        self.collator = DataCollatorForCompletionOnlyLM(\n            self.tokenized_response_w_context, self.tokenized_instruction_w_context, tokenizer=self.tokenizer\n        )\n        collator_output = self.collator.torch_call([self.tokenized_instruction])\n        collator_text = self.tokenizer.decode(\n            collator_output[\"labels\"][torch.where(collator_output[\"labels\"] != -100)]\n        )\n        expected_text = \" First response\\n\\n Second response\" \"\"\n        assert collator_text == expected_text\n\n    def test_data_collator_handling_of_long_sequences(self):\n        self.tokenizer = AutoTokenizer.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        self.instruction = \"\"\"### System: You are a helpful assistant.\n\n### User: How much is 2+2? I'm asking because I'm not sure. And I'm not sure because I'm not good at math.\n\"\"\"\n        self.response_template = \"\\n### Assistant:\"\n        # check DataCollatorForCompletionOnlyLM using response template only\n        self.tokenized_instruction = self.tokenizer.encode(self.instruction, add_special_tokens=False)\n        self.collator = DataCollatorForCompletionOnlyLM(self.response_template, tokenizer=self.tokenizer)\n        encoded_instance = self.collator.torch_call([self.tokenized_instruction])\n        result = torch.all(encoded_instance[\"labels\"] == -100)\n        assert result, \"Not all values in the tensor are -100.\"\n\n        # check DataCollatorForCompletionOnlyLM using response template and instruction template\n        self.instruction_template = \"\\n### User:\"\n        self.collator = DataCollatorForCompletionOnlyLM(\n            self.response_template, self.instruction_template, tokenizer=self.tokenizer\n        )\n        encoded_instance = self.collator.torch_call([self.tokenized_instruction])\n        result = torch.all(encoded_instance[\"labels\"] == -100)\n        assert result, \"Not all values in the tensor are -100.\"\n\n    def test_padding_free(self):\n        tokenizer = AutoTokenizer.from_pretrained(\"trl-internal-testing/dummy-GPT2-correct-vocab\")\n        if tokenizer.pad_token_id is None:\n            tokenizer.pad_token = tokenizer.eos_token\n            tokenizer.pad_token_id = tokenizer.eos_token_id\n        inst1 = \"### System: You are a helpful assistant.\\n\\n### User: How much is 2+2?\\n\\n### Assistant: 2+2 equals 4\"\n        inst2 = \"### System: You are a honest and helpful assistant.\\n\\n### User: What is the answer of 22x22?\\n\\n### Assistant: 22x22 equals 484\"\n\n        response_template = \"\\n### Assistant:\"\n        collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)\n        collator_paddingfree = DataCollatorForCompletionOnlyLM(\n            response_template, tokenizer=tokenizer, padding_free=True\n        )\n\n        tokenized_instruction = [tokenizer(x, add_special_tokens=False) for x in [inst1, inst2]]\n        batch = collator(tokenized_instruction)\n        batch_paddingfree = collator_paddingfree(tokenized_instruction)\n\n        self.assertNotIn(\"attention_mask\", batch_paddingfree)\n        self.assertIn(\"input_ids\", batch_paddingfree)\n        self.assertIn(\"labels\", batch_paddingfree)\n        self.assertIn(\"position_ids\", batch_paddingfree)\n        self.assertEqual(batch_paddingfree[\"input_ids\"].size(), batch_paddingfree[\"labels\"].size())\n        self.assertEqual(batch_paddingfree[\"labels\"].size(), batch_paddingfree[\"position_ids\"].size())\n\n        attn_mask = batch[\"attention_mask\"]\n        input_ids_remove_pad = batch[\"input_ids\"][attn_mask.bool()].unsqueeze(0)\n        expected_position_ids = attn_mask.cumsum(1)[attn_mask.bool()].unsqueeze(0) - 1\n        expected_labels = []\n        for idx in range(batch[\"input_ids\"].size(0)):\n            expected_labels.append(batch[\"labels\"][idx][attn_mask[idx].bool()])\n            expected_labels[-1][0] = collator.ignore_index\n        expected_labels = torch.cat(expected_labels).unsqueeze(0)\n\n        self.assertTrue((input_ids_remove_pad == batch_paddingfree[\"input_ids\"]).all())\n        self.assertTrue((expected_position_ids == batch_paddingfree[\"position_ids\"]).all())\n        self.assertTrue((expected_labels == batch_paddingfree[\"labels\"]).all())\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport subprocess\nimport sys\nimport unittest\n\n\n@unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\ndef test_sft_cli():\n    try:\n        subprocess.run(\n            \"trl sft --max_steps 1 --output_dir tmp-sft --model_name_or_path trl-internal-testing/tiny-random-LlamaForCausalLM --dataset_name stanfordnlp/imdb --learning_rate 1e-4 --lr_scheduler_type cosine --dataset_text_field text\",\n            shell=True,\n            check=True,\n        )\n    except BaseException as exc:\n        raise AssertionError(\"An error occured while running the CLI, please double check\") from exc\n\n\n@unittest.skipIf(sys.platform.startswith(\"win\"), \"Skipping on Windows\")\ndef test_dpo_cli():\n    try:\n        subprocess.run(\n            \"trl dpo --max_steps 1 --output_dir tmp-dpo --model_name_or_path trl-internal-testing/tiny-random-LlamaForCausalLM --dataset_name trl-lib/ultrafeedback_binarized --learning_rate 1e-4 --lr_scheduler_type cosine\",\n            shell=True,\n            check=True,\n        )\n    except BaseException as exc:\n        raise AssertionError(\"An error occured while running the CLI, please double check\") from exc\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport unittest\n\nfrom transformers import is_wandb_available\n\nfrom trl import is_diffusers_available, is_liger_kernel_available\n\n\ndef require_diffusers(test_case):\n    \"\"\"\n    Decorator marking a test that requires diffusers. Skips the test if diffusers is not available.\n    \"\"\"\n    return unittest.skipUnless(is_diffusers_available(), \"test requires diffusers\")(test_case)\n\n\ndef require_no_wandb(test_case):\n    \"\"\"\n    Decorator marking a test that requires no wandb. Skips the test if wandb is available.\n    \"\"\"\n    return unittest.skipUnless(not is_wandb_available(), \"test requires no wandb\")(test_case)\n\n\ndef require_liger_kernel(test_case):\n    \"\"\"\n    Decorator marking a test that requires liger_kernel. Skips the test if liger_kernel is not available.\n    \"\"\"\n    return unittest.skipUnless(is_liger_kernel_available(), \"test requires liger_kernel\")(test_case)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tempfile\nimport unittest\n\nimport torch\nimport torch.nn as nn\nfrom datasets import Dataset\nfrom transformers import Trainer, TrainingArguments\n\nfrom trl.trainer.callbacks import RichProgressCallback\n\n\nclass DummyModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.a = nn.Parameter(torch.tensor(1.0))\n\n    def forward(self, x):\n        return self.a * x\n\n\nclass TestRichProgressCallback(unittest.TestCase):\n    def setUp(self):\n        self.dummy_model = DummyModel()\n        self.dummy_train_dataset = Dataset.from_list([{\"x\": 1.0, \"y\": 2.0}] * 5)\n        self.dummy_val_dataset = Dataset.from_list([{\"x\": 1.0, \"y\": 2.0}] * 101)\n\n    def test_rich_progress_callback_logging(self):\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = TrainingArguments(\n                output_dir=tmp_dir,\n                per_device_eval_batch_size=2,\n                per_device_train_batch_size=2,\n                num_train_epochs=4,\n                eval_strategy=\"steps\",\n                eval_steps=1,\n                logging_strategy=\"steps\",\n                logging_steps=1,\n                save_strategy=\"no\",\n                report_to=\"none\",\n                disable_tqdm=True,\n            )\n            callbacks = [RichProgressCallback()]\n            trainer = Trainer(\n                model=self.dummy_model,\n                train_dataset=self.dummy_train_dataset,\n                eval_dataset=self.dummy_val_dataset,\n                args=training_args,\n                callbacks=callbacks,\n            )\n\n            trainer.train()\n            trainer.train()\n\n\n# Copyright 2023 metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport unittest\n\nimport torch\nfrom parameterized import parameterized\nfrom transformers.utils import is_peft_available\n\nfrom trl import is_diffusers_available\n\nfrom .testing_utils import require_diffusers\n\n\nif is_diffusers_available() and is_peft_available():\n    from trl import AlignPropConfig, AlignPropTrainer, DefaultDDPOStableDiffusionPipeline\n\n\ndef scorer_function(images, prompts, metadata):\n    return torch.randn(1) * 3.0, {}\n\n\ndef prompt_function():\n    return (\"cabbages\", {})\n\n\n@require_diffusers\nclass AlignPropTrainerTester(unittest.TestCase):\n    \"\"\"\n    Test the AlignPropTrainer class.\n    \"\"\"\n\n    def setUp(self):\n        training_args = AlignPropConfig(\n            num_epochs=2,\n            train_gradient_accumulation_steps=1,\n            train_batch_size=2,\n            truncated_backprop_rand=False,\n            mixed_precision=None,\n            save_freq=1000000,\n        )\n        pretrained_model = \"hf-internal-testing/tiny-stable-diffusion-torch\"\n        pretrained_revision = \"main\"\n        pipeline_with_lora = DefaultDDPOStableDiffusionPipeline(\n            pretrained_model, pretrained_model_revision=pretrained_revision, use_lora=True\n        )\n        pipeline_without_lora = DefaultDDPOStableDiffusionPipeline(\n            pretrained_model, pretrained_model_revision=pretrained_revision, use_lora=False\n        )\n        self.trainer_with_lora = AlignPropTrainer(training_args, scorer_function, prompt_function, pipeline_with_lora)\n        self.trainer_without_lora = AlignPropTrainer(\n            training_args, scorer_function, prompt_function, pipeline_without_lora\n        )\n\n    def tearDown(self) -> None:\n        gc.collect()\n\n    @parameterized.expand([True, False])\n    def test_generate_samples(self, use_lora):\n        trainer = self.trainer_with_lora if use_lora else self.trainer_without_lora\n        output_pairs = trainer._generate_samples(2, with_grad=True)\n        assert len(output_pairs.keys()) == 3\n        assert len(output_pairs[\"images\"]) == 2\n\n    @parameterized.expand([True, False])\n    def test_calculate_loss(self, use_lora):\n        trainer = self.trainer_with_lora if use_lora else self.trainer_without_lora\n        sample = trainer._generate_samples(2)\n\n        images = sample[\"images\"]\n        prompts = sample[\"prompts\"]\n\n        assert images.shape == (2, 3, 128, 128)\n        assert len(prompts) == 2\n\n        rewards = trainer.compute_rewards(sample)\n        loss = trainer.calculate_loss(rewards)\n\n        assert torch.isfinite(loss.cpu())\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport itertools\nimport tempfile\nimport unittest\n\nimport torch\nfrom accelerate.utils.memory import release_memory\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\nfrom transformers.testing_utils import (\n    require_bitsandbytes,\n    require_peft,\n    require_torch_accelerator,\n    require_torch_multi_accelerator,\n)\nfrom transformers.utils import is_peft_available\n\nfrom trl import SFTConfig, SFTTrainer\nfrom trl.models.utils import setup_chat_format\n\nfrom ..testing_utils import require_liger_kernel\nfrom .testing_constants import DEVICE_MAP_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS, MODELS_TO_TEST, PACKING_OPTIONS\n\n\nif is_peft_available():\n    from peft import LoraConfig, PeftModel\n\n\n@require_torch_accelerator\nclass SFTTrainerSlowTester(unittest.TestCase):\n    def setUp(self):\n        self.train_dataset = load_dataset(\"stanfordnlp/imdb\", split=\"train[:10%]\")\n        self.eval_dataset = load_dataset(\"stanfordnlp/imdb\", split=\"test[:10%]\")\n        self.dataset_text_field = \"text\"\n        self.max_seq_length = 128\n        self.peft_config = LoraConfig(\n            lora_alpha=16,\n            lora_dropout=0.1,\n            r=8,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n\n    def tearDown(self):\n        gc.collect()\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    def test_sft_trainer_str(self, model_name, packing):\n        \"\"\"\n        Simply tests if passing a simple str to `SFTTrainer` loads and runs the trainer\n        as expected.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n            )\n\n            trainer = SFTTrainer(\n                model_name,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    def test_sft_trainer_transformers(self, model_name, packing):\n        \"\"\"\n        Simply tests if passing a transformers model to `SFTTrainer` loads and runs the trainer\n        as expected.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    @require_peft\n    def test_sft_trainer_peft(self, model_name, packing):\n        \"\"\"\n        Simply tests if passing a transformers model + peft config to `SFTTrainer` loads and runs the trainer\n        as expected.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                fp16=True,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    def test_sft_trainer_transformers_mp(self, model_name, packing):\n        \"\"\"\n        Simply tests if passing a transformers model to `SFTTrainer` loads and runs the trainer\n        as expected in mixed precision.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                fp16=True,  # this is sufficient to enable amp\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS)))\n    def test_sft_trainer_transformers_mp_gc(self, model_name, packing, gradient_checkpointing_kwargs):\n        \"\"\"\n        Simply tests if passing a transformers model to `SFTTrainer` loads and runs the trainer\n        as expected in mixed precision + different scenarios of gradient_checkpointing.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n                fp16=True,  # this is sufficient to enable amp\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS)))\n    @require_peft\n    def test_sft_trainer_transformers_mp_gc_peft(self, model_name, packing, gradient_checkpointing_kwargs):\n        \"\"\"\n        Simply tests if passing a transformers model + PEFT to `SFTTrainer` loads and runs the trainer\n        as expected in mixed precision + different scenarios of gradient_checkpointing.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n                fp16=True,  # this is sufficient to enable amp\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(\n        list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS, DEVICE_MAP_OPTIONS))\n    )\n    @require_torch_multi_accelerator\n    def test_sft_trainer_transformers_mp_gc_device_map(\n        self, model_name, packing, gradient_checkpointing_kwargs, device_map\n    ):\n        \"\"\"\n        Simply tests if passing a transformers model to `SFTTrainer` loads and runs the trainer\n        as expected in mixed precision + different scenarios of gradient_checkpointing (single, multi-gpu, etc).\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n                fp16=True,  # this is sufficient to enable amp\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n            )\n\n            model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device_map)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS, GRADIENT_CHECKPOINTING_KWARGS)))\n    @require_peft\n    @require_bitsandbytes\n    def test_sft_trainer_transformers_mp_gc_peft_qlora(self, model_name, packing, gradient_checkpointing_kwargs):\n        \"\"\"\n        Simply tests if passing a transformers model + PEFT + bnb to `SFTTrainer` loads and runs the trainer\n        as expected in mixed precision + different scenarios of gradient_checkpointing.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n                fp16=True,  # this is sufficient to enable amp\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n            )\n\n            quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)\n\n            model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    @require_peft\n    @require_bitsandbytes\n    def test_sft_trainer_with_chat_format_qlora(self, model_name, packing):\n        \"\"\"\n        Simply tests if using setup_chat_format with a transformers model + peft + bnb config to `SFTTrainer` loads and runs the trainer\n        as expected.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            train_dataset = load_dataset(\"trl-internal-testing/dolly-chatml-sft\", split=\"train\")\n\n            training_args = SFTConfig(\n                packing=packing,\n                max_seq_length=self.max_seq_length,\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=10,\n                fp16=True,\n            )\n\n            quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)\n\n            model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n            model, tokenizer = setup_chat_format(model, tokenizer)\n\n            trainer = SFTTrainer(\n                model,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=train_dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n\n            trainer.train()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, PACKING_OPTIONS)))\n    @require_liger_kernel\n    def test_sft_trainer_with_liger(self, model_name, packing):\n        \"\"\"\n        Tests if passing use_liger=True to SFTConfig loads and runs the trainer\n        with AutoLigerKernelForCausalLM as expected.\n        \"\"\"\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = SFTConfig(\n                output_dir=tmp_dir,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                per_device_train_batch_size=2,\n                max_steps=2,\n                packing=packing,\n                dataset_text_field=self.dataset_text_field,\n                max_seq_length=self.max_seq_length,\n                use_liger=True,\n            )\n\n            trainer = SFTTrainer(\n                model_name,\n                args=training_args,\n                train_dataset=self.train_dataset,\n                eval_dataset=self.eval_dataset,\n            )\n\n            # check that the components of the trainer.model are monkey patched:\n            self.assertTrue(any(\"Liger\" in type(module).__name__ for module in trainer.model.model.modules()))\n            trainer.train()\n\n        release_memory(trainer.model, trainer)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport itertools\nimport tempfile\nimport unittest\n\nimport torch\nfrom accelerate.utils.memory import release_memory\nfrom datasets import load_dataset\nfrom parameterized import parameterized\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\nfrom transformers.testing_utils import require_bitsandbytes, require_peft, require_torch_accelerator, torch_device\nfrom transformers.utils import is_peft_available\n\nfrom trl import DPOConfig, DPOTrainer\n\nfrom .testing_constants import DPO_LOSS_TYPES, DPO_PRECOMPUTE_LOGITS, GRADIENT_CHECKPOINTING_KWARGS, MODELS_TO_TEST\n\n\nif is_peft_available():\n    from peft import LoraConfig, PeftModel\n\n\n@require_torch_accelerator\nclass DPOTrainerSlowTester(unittest.TestCase):\n    def setUp(self):\n        self.dataset = load_dataset(\"trl-internal-testing/mlabonne-chatml-dpo-pairs-copy\", split=\"train[:10%]\")\n        self.peft_config = LoraConfig(\n            lora_alpha=16,\n            lora_dropout=0.1,\n            r=8,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\",\n        )\n        self.max_length = 128\n\n    def tearDown(self):\n        gc.collect()\n        if torch_device == \"cpu\":\n            torch.cuda.empty_cache()\n        elif torch_device == \"xpu\":\n            torch.xpu.empty_cache()\n        gc.collect()\n\n    @parameterized.expand(list(itertools.product(MODELS_TO_TEST, DPO_LOSS_TYPES, DPO_PRECOMPUTE_LOGITS)))\n    def test_dpo_bare_model(self, model_id, loss_type, pre_compute_logits):\n        \"\"\"\n        A test that tests the simple usage of `DPOTrainer` using a bare model in full precision.\n        \"\"\"\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=2,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=2,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                fp16=True,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                beta=0.1,\n                loss_type=loss_type,\n                precompute_ref_log_probs=pre_compute_logits,\n                max_length=self.max_length,\n            )\n\n            # dpo train lora model\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.dataset,\n                eval_dataset=self.dataset,\n            )\n\n            # train the model\n            trainer.train()\n\n            # save trained model or adapter\n            trainer.save_model()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(\n        list(\n            itertools.product(\n                MODELS_TO_TEST,\n                DPO_LOSS_TYPES,\n                DPO_PRECOMPUTE_LOGITS,\n                GRADIENT_CHECKPOINTING_KWARGS,\n            )\n        )\n    )\n    @require_peft\n    def test_dpo_peft_model(self, model_id, loss_type, pre_compute_logits, gradient_checkpointing_kwargs):\n        \"\"\"\n        A test that tests the simple usage of `DPOTrainer` using a peft model in full precision + different scenarios of gradient checkpointing.\n        \"\"\"\n        model = AutoModelForCausalLM.from_pretrained(model_id)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=2,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=2,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                fp16=True,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n                generate_during_eval=False,\n                loss_type=loss_type,\n                precompute_ref_log_probs=pre_compute_logits,\n                beta=0.1,\n                max_length=self.max_length,\n            )\n\n            # dpo train lora model\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.dataset,\n                eval_dataset=self.dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n            assert trainer.ref_model is None\n\n            # train the model\n            trainer.train()\n\n            # save trained model or adapter\n            trainer.save_model()\n\n        release_memory(model, trainer)\n\n    @parameterized.expand(\n        list(\n            itertools.product(\n                MODELS_TO_TEST,\n                DPO_LOSS_TYPES,\n                DPO_PRECOMPUTE_LOGITS,\n                GRADIENT_CHECKPOINTING_KWARGS,\n            )\n        )\n    )\n    @require_bitsandbytes\n    @require_peft\n    def test_dpo_peft_model_qlora(self, model_id, loss_type, pre_compute_logits, gradient_checkpointing_kwargs):\n        \"\"\"\n        A test that tests the simple usage of `DPOTrainer` using QLoRA + different scenarios of gradient checkpointing.\n        \"\"\"\n        quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)\n\n        model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)\n        tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n        with tempfile.TemporaryDirectory() as tmp_dir:\n            training_args = DPOConfig(\n                output_dir=tmp_dir,\n                per_device_train_batch_size=2,\n                max_steps=2,\n                remove_unused_columns=False,\n                gradient_accumulation_steps=2,\n                learning_rate=9e-1,\n                eval_strategy=\"steps\",\n                fp16=True,\n                logging_strategy=\"no\",\n                report_to=\"none\",\n                gradient_checkpointing=True,\n                gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,\n                beta=0.1,\n                generate_during_eval=False,\n                loss_type=loss_type,\n                precompute_ref_log_probs=pre_compute_logits,\n                max_length=self.max_length,\n            )\n\n            # dpo train lora model\n            trainer = DPOTrainer(\n                model=model,\n                ref_model=None,\n                args=training_args,\n                tokenizer=tokenizer,\n                train_dataset=self.dataset,\n                eval_dataset=self.dataset,\n                peft_config=self.peft_config,\n            )\n\n            assert isinstance(trainer.model, PeftModel)\n            assert trainer.ref_model is None\n\n            # train the model\n            trainer.train()\n\n            # save trained model or adapter\n            trainer.save_model()\n\n        release_memory(model, trainer)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# TODO: push them under trl-org\nMODELS_TO_TEST = [\n    \"trl-internal-testing/tiny-random-LlamaForCausalLM\",\n    \"HuggingFaceM4/tiny-random-MistralForCausalLM\",\n]\n\n# We could have also not declared these variables but let's be verbose\nPACKING_OPTIONS = [True, False]\nGRADIENT_CHECKPOINTING_KWARGS = [None, {\"use_reentrant\": False}, {\"use_reentrant\": True}]\nDEVICE_MAP_OPTIONS = [{\"\": 0}, \"auto\"]\n\nDPO_LOSS_TYPES = [\"sigmoid\", \"ipo\"]\nDPO_PRECOMPUTE_LOGITS = [True, False]\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n\n# Examples\n\nPlease check out https://huggingface.co/docs/trl/example_overview for documentation on our examples.\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# 0. imports\nimport torch\nfrom transformers import GPT2Tokenizer\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer\n\n\n# 1. load a pretrained model\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\nref_model = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\ntokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\ntokenizer.pad_token = tokenizer.eos_token\n\n# 2. initialize trainer\nppo_config = {\"mini_batch_size\": 1, \"batch_size\": 1}\nconfig = PPOConfig(**ppo_config)\nppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)\n\n# 3. encode a query\nquery_txt = \"This morning I went to the \"\nquery_tensor = tokenizer.encode(query_txt, return_tensors=\"pt\").to(model.pretrained_model.device)\n\n# 4. generate model response\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n    \"max_new_tokens\": 20,\n}\nresponse_tensor = ppo_trainer.generate(list(query_tensor), return_prompt=False, **generation_kwargs)\nresponse_txt = tokenizer.decode(response_tensor[0])\n\n# 5. define a reward for response\n# (this could be any reward such as human feedback or output from another model)\nreward = [torch.tensor(1.0, device=model.pretrained_model.device)]\n\n# 6. train model with ppo\ntrain_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)\n\n\n# Research projects that use TRL\n\nWelcome to the research projects folder! Here you can find the scripts used for some research projects that used TRL and maintained by the developers and the community (LM de-toxification, Stack-Llama, etc.). Check out the READMEs in the subfolders for more information!\n\n- [De-detoxifying language models](https://github.com/huggingface/trl/tree/main/examples/research_projects/toxicity)\n- [Stack-Llama](https://github.com/huggingface/trl/tree/main/examples/research_projects/stack_llama)\n- [Stack-Llama-2](https://github.com/huggingface/trl/tree/main/examples/research_projects/stack_llama_2)\n\n# DPO pipeline for the creation of StackLlaMa 2: a Stack exchange llama-v2-7b model\n\n## Prerequisites\n\nInstall all the dependencies in the `requirements.txt`:\n\n```\n$ pip install -U -r requirements.txt\n```\n\nSince we will use `accelerate` for training, make sure to run:\n```\n$ accelerate config\n```\n\n## Training\n\nThere were two main steps to the DPO training process:\n1. Supervised fine-tuning of the base llama-v2-7b model to create llama-v2-7b-se:\n\n    ```\n    accelerate launch examples/research_projects/stack_llama_2/scripts/sft_llama2.py \\\n        --output_dir=\"./sft\" \\\n        --max_steps=500 \\\n        --logging_steps=10 \\\n        --save_steps=10 \\\n        --per_device_train_batch_size=4 \\\n        --per_device_eval_batch_size=1 \\\n        --gradient_accumulation_steps=2 \\\n        --gradient_checkpointing=False \\\n        --group_by_length=False \\\n        --learning_rate=1e-4 \\\n        --lr_scheduler_type=\"cosine\" \\\n        --warmup_steps=100 \\\n        --weight_decay=0.05 \\\n        --optim=\"paged_adamw_32bit\" \\\n        --bf16=True \\\n        --remove_unused_columns=False \\\n        --run_name=\"sft_llama2\" \\\n        --report_to=\"wandb\"\n    ```\n1. Run the DPO trainer using the model saved by the previous step:\n    ```\n    accelerate launch examples/research_projects/stack_llama_2/scripts/dpo_llama2.py \\\n        --model_name_or_path=\"sft/final_checkpoint\" \\\n        --output_dir=\"dpo\"\n    ```\n\n\n## Merging the adaptors\n\nTo merge the adaptors into the base model we can use the `merge_peft_adapter.py` helper script that comes with TRL:\n\n```\npython examples/research_projects/stack_llama/scripts/merge_peft_adapter.py --base_model_name=\"meta-llama/Llama-2-7b-hf\" --adapter_model_name=\"dpo/final_checkpoint/\" --output_name=\"stack-llama-2\"\n```\n\nwhich will also push the model to your HuggingFace hub account.\n\n## Running the model\n\nWe can load the DPO-trained LoRA adaptors which were saved by the DPO training step and load them via:\n\n```py\nfrom peft import AutoPeftModelForCausalLM\n\n\nmodel = AutoPeftModelForCausalLM.from_pretrained(\n    \"dpo/final_checkpoint\",\n    low_cpu_mem_usage=True,\n    torch_dtype=torch.float16,\n    load_in_4bit=True,\n)\n\nmodel.generate(...)\n```\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# 0. imports\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Dict, Optional\n\nimport torch\nfrom accelerate import Accelerator\nfrom datasets import Dataset, load_dataset\nfrom peft import LoraConfig\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, set_seed\n\nfrom trl import DPOConfig, DPOTrainer\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The arguments for the DPO training script.\n    \"\"\"\n\n    # data parameters\n    beta: Optional[float] = field(default=0.1, metadata={\"help\": \"the beta parameter for DPO loss\"})\n\n    # training parameters\n    model_name_or_path: Optional[str] = field(\n        default=\"../sft/results/final_checkpoint\",\n        metadata={\"help\": \"the location of the SFT model name or path\"},\n    )\n    learning_rate: Optional[float] = field(default=5e-4, metadata={\"help\": \"optimizer learning rate\"})\n    lr_scheduler_type: Optional[str] = field(default=\"cosine\", metadata={\"help\": \"the lr scheduler type\"})\n    warmup_steps: Optional[int] = field(default=100, metadata={\"help\": \"the number of warmup steps\"})\n    weight_decay: Optional[float] = field(default=0.05, metadata={\"help\": \"the weight decay\"})\n    optimizer_type: Optional[str] = field(default=\"paged_adamw_32bit\", metadata={\"help\": \"the optimizer type\"})\n\n    per_device_train_batch_size: Optional[int] = field(default=4, metadata={\"help\": \"train batch size per device\"})\n    per_device_eval_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"eval batch size per device\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=4, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    gradient_checkpointing: Optional[bool] = field(\n        default=True, metadata={\"help\": \"whether to use gradient checkpointing\"}\n    )\n\n    gradient_checkpointing_use_reentrant: Optional[bool] = field(\n        default=False, metadata={\"help\": \"whether to use reentrant for gradient checkpointing\"}\n    )\n\n    lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n    lora_dropout: Optional[float] = field(default=0.05, metadata={\"help\": \"the lora dropout parameter\"})\n    lora_r: Optional[int] = field(default=8, metadata={\"help\": \"the lora r parameter\"})\n\n    max_prompt_length: Optional[int] = field(default=512, metadata={\"help\": \"the maximum prompt length\"})\n    max_length: Optional[int] = field(default=1024, metadata={\"help\": \"the maximum sequence length\"})\n    max_steps: Optional[int] = field(default=1000, metadata={\"help\": \"max number of training steps\"})\n    logging_steps: Optional[int] = field(default=10, metadata={\"help\": \"the logging frequency\"})\n    save_steps: Optional[int] = field(default=100, metadata={\"help\": \"the saving frequency\"})\n    eval_steps: Optional[int] = field(default=100, metadata={\"help\": \"the evaluation frequency\"})\n\n    output_dir: Optional[str] = field(default=\"./results\", metadata={\"help\": \"the output directory\"})\n    log_freq: Optional[int] = field(default=1, metadata={\"help\": \"the logging frequency\"})\n    load_in_4bit: Optional[bool] = field(default=True, metadata={\"help\": \"whether to load the model in 4bit\"})\n    model_dtype: Optional[str] = field(\n        default=\"float16\", metadata={\"help\": \"model_dtype[float16, bfloat16, float] for loading.\"}\n    )\n\n    # instrumentation\n    report_to: Optional[str] = field(\n        default=\"wandb\",\n        metadata={\n            \"help\": 'The list of integrations to report the results and logs to. Supported platforms are `\"azure_ml\"`,'\n            '`\"comet_ml\"`, `\"mlflow\"`, `\"neptune\"`, `\"tensorboard\"`,`\"clearml\"` and `\"wandb\"`. '\n            'Use `\"all\"` to report to all integrations installed, `\"none\"` for no integrations.'\n        },\n    )\n    # debug argument for distributed training\n    ignore_bias_buffers: Optional[bool] = field(\n        default=False,\n        metadata={\n            \"help\": \"fix for DDP issues with LM bias/mask buffers - invalid scalar type,`inplace operation. See\"\n            \"https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992\"\n        },\n    )\n    seed: Optional[int] = field(\n        default=0, metadata={\"help\": \"Random seed that will be set at the beginning of training.\"}\n    )\n\n\ndef get_stack_exchange_paired(\n    data_dir: str = \"data/rl\",\n    cache_dir: Optional[str] = None,\n    num_proc=24,\n) -> Dataset:\n    \"\"\"Load the stack-exchange-paired dataset from Hugging Face and convert it to the necessary format.\n\n    The dataset is converted to a dictionary with the following structure:\n    {\n        'prompt': List[str],\n        'chosen': List[str],\n        'rejected': List[str],\n    }\n\n    Prompts are structured as follows:\n      \"Question: \" + <prompt> + \"\\n\\nAnswer: \"\n    \"\"\"\n    dataset = load_dataset(\n        \"lvwerra/stack-exchange-paired\",\n        split=\"train\",\n        cache_dir=cache_dir,\n        data_dir=data_dir,\n        verification_mode=\"no_checks\",\n    )\n    original_columns = dataset.column_names\n\n    def return_prompt_and_responses(samples) -> Dict[str, str]:\n        return {\n            \"prompt\": [\"Question: \" + question + \"\\n\\nAnswer: \" for question in samples[\"question\"]],\n            \"chosen\": samples[\"response_j\"],\n            \"rejected\": samples[\"response_k\"],\n        }\n\n    return dataset.map(\n        return_prompt_and_responses,\n        batched=True,\n        num_proc=num_proc,\n        remove_columns=original_columns,\n    )\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    script_args = parser.parse_args_into_dataclasses()[0]\n\n    set_seed(script_args.seed)\n\n    # 1. load a pretrained model\n    torch_dtype = torch.float\n    if script_args.model_dtype == \"float16\":\n        torch_dtype = torch.float16\n    elif script_args.model_dtype == \"bfloat16\":\n        torch_dtype = torch.bfloat16\n\n    model = AutoModelForCausalLM.from_pretrained(\n        script_args.model_name_or_path,\n        low_cpu_mem_usage=True,\n        torch_dtype=torch_dtype,\n        load_in_4bit=script_args.load_in_4bit,\n        device_map={\"\": Accelerator().local_process_index},\n    )\n    model.config.use_cache = False\n\n    if script_args.ignore_bias_buffers:\n        # torch distributed hack\n        model._ddp_params_and_buffers_to_ignore = [\n            name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool\n        ]\n\n    tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n    tokenizer.pad_token = tokenizer.eos_token\n\n    # 2. Load the Stack-exchange paired dataset\n    train_dataset = get_stack_exchange_paired(data_dir=\"data/rl\")\n    train_dataset = train_dataset.filter(\n        lambda x: len(x[\"prompt\"]) + len(x[\"chosen\"]) <= script_args.max_length\n        and len(x[\"prompt\"]) + len(x[\"rejected\"]) <= script_args.max_length,\n        num_proc=script_args.num_proc,\n    )\n\n    # 3. Load evaluation dataset\n    eval_dataset = get_stack_exchange_paired(data_dir=\"data/evaluation\")\n    eval_dataset = eval_dataset.filter(\n        lambda x: len(x[\"prompt\"]) + len(x[\"chosen\"]) <= script_args.max_length\n        and len(x[\"prompt\"]) + len(x[\"rejected\"]) <= script_args.max_length,\n        num_proc=script_args.num_proc,\n    )\n\n    # 4. initialize training arguments:\n    training_args = DPOConfig(\n        per_device_train_batch_size=script_args.per_device_train_batch_size,\n        per_device_eval_batch_size=script_args.per_device_eval_batch_size,\n        max_steps=script_args.max_steps,\n        logging_steps=script_args.logging_steps,\n        save_steps=script_args.save_steps,\n        gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n        gradient_checkpointing=script_args.gradient_checkpointing,\n        learning_rate=script_args.learning_rate,\n        eval_strategy=\"steps\",\n        eval_steps=script_args.eval_steps,\n        output_dir=script_args.output_dir,\n        report_to=script_args.report_to,\n        lr_scheduler_type=script_args.lr_scheduler_type,\n        warmup_steps=script_args.warmup_steps,\n        optim=script_args.optimizer_type,\n        bf16=True,\n        remove_unused_columns=False,\n        run_name=\"dpo_llama2\",\n        gradient_checkpointing_kwargs=dict(use_reentrant=script_args.gradient_checkpointing_use_reentrant),\n        seed=script_args.seed,\n    )\n\n    peft_config = LoraConfig(\n        r=script_args.lora_r,\n        lora_alpha=script_args.lora_alpha,\n        lora_dropout=script_args.lora_dropout,\n        target_modules=[\n            \"q_proj\",\n            \"v_proj\",\n            \"k_proj\",\n            \"out_proj\",\n            \"fc_in\",\n            \"fc_out\",\n            \"wte\",\n        ],\n        bias=\"none\",\n        task_type=\"CAUSAL_LM\",\n    )\n\n    # 5. initialize the DPO trainer\n    dpo_trainer = DPOTrainer(\n        model,\n        ref_model=None,\n        args=training_args,\n        beta=script_args.beta,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n        tokenizer=tokenizer,\n        peft_config=peft_config,\n        max_prompt_length=script_args.max_prompt_length,\n        max_length=script_args.max_length,\n    )\n\n    # 6. train\n    dpo_trainer.train()\n    dpo_trainer.save_model(script_args.output_dir)\n\n    # 7. save\n    output_dir = os.path.join(script_args.output_dir, \"final_checkpoint\")\n    dpo_trainer.model.save_pretrained(output_dir)\n\n\ntransformers\ntrl\npeft\naccelerate\ndatasets\nbitsandbytes\nwandb\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Fine-Tune Llama2-7b on SE paired dataset\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom accelerate import Accelerator\nfrom datasets import load_dataset\nfrom peft import AutoPeftModelForCausalLM, LoraConfig\nfrom tqdm import tqdm\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoTokenizer,\n    BitsAndBytesConfig,\n    HfArgumentParser,\n    is_torch_npu_available,\n    is_torch_xpu_available,\n    set_seed,\n)\n\nfrom trl import SFTConfig, SFTTrainer\nfrom trl.trainer import ConstantLengthDataset\n\n\n@dataclass\nclass ScriptArguments:\n    model_name: Optional[str] = field(default=\"meta-llama/Llama-2-7b-hf\", metadata={\"help\": \"the model name\"})\n    dataset_name: Optional[str] = field(default=\"lvwerra/stack-exchange-paired\", metadata={\"help\": \"the dataset name\"})\n    subset: Optional[str] = field(default=\"data/finetune\", metadata={\"help\": \"the subset to use\"})\n    split: Optional[str] = field(default=\"train\", metadata={\"help\": \"the split to use\"})\n    size_valid_set: Optional[int] = field(default=4000, metadata={\"help\": \"the size of the validation set\"})\n    streaming: Optional[bool] = field(default=True, metadata={\"help\": \"whether to stream the dataset\"})\n    shuffle_buffer: Optional[int] = field(default=5000, metadata={\"help\": \"the shuffle buffer size\"})\n    seq_length: Optional[int] = field(default=1024, metadata={\"help\": \"the sequence length\"})\n    num_workers: Optional[int] = field(default=4, metadata={\"help\": \"the number of workers\"})\n    use_bnb: Optional[bool] = field(default=True, metadata={\"help\": \"whether to use BitsAndBytes\"})\n\n    # LoraConfig\n    lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n    lora_dropout: Optional[float] = field(default=0.05, metadata={\"help\": \"the lora dropout parameter\"})\n    lora_r: Optional[int] = field(default=8, metadata={\"help\": \"the lora r parameter\"})\n\n\nparser = HfArgumentParser((ScriptArguments, SFTConfig))\nscript_args, training_args = parser.parse_args_into_dataclasses()\npeft_config = LoraConfig(\n    r=script_args.lora_r,\n    lora_alpha=script_args.lora_alpha,\n    lora_dropout=script_args.lora_dropout,\n    target_modules=[\"q_proj\", \"v_proj\"],\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\n\nif training_args.group_by_length and training_args.packing:\n    raise ValueError(\"Cannot use both packing and group by length\")\n\n# `gradient_checkpointing` was True by default until `1f3314`, but it's actually not used.\n# `gradient_checkpointing=True` will cause `Variable._execution_engine.run_backward`.\nif training_args.gradient_checkpointing:\n    raise ValueError(\"gradient_checkpointing not supported\")\n\nset_seed(training_args.seed)\n\n\ndef chars_token_ratio(dataset, tokenizer, nb_examples=400):\n    \"\"\"\n    Estimate the average number of characters per token in the dataset.\n    \"\"\"\n    total_characters, total_tokens = 0, 0\n    for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):\n        text = prepare_sample_text(example)\n        total_characters += len(text)\n        if tokenizer.is_fast:\n            total_tokens += len(tokenizer(text).tokens())\n        else:\n            total_tokens += len(tokenizer.tokenize(text))\n\n    return total_characters / total_tokens\n\n\ndef print_trainable_parameters(model):\n    \"\"\"\n    Prints the number of trainable parameters in the model.\n    \"\"\"\n    trainable_params = 0\n    all_param = 0\n    for _, param in model.named_parameters():\n        all_param += param.numel()\n        if param.requires_grad:\n            trainable_params += param.numel()\n    print(\n        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n    )\n\n\ndef prepare_sample_text(example):\n    \"\"\"Prepare the text from a sample of the dataset.\"\"\"\n    text = f\"Question: {example['question']}\\n\\nAnswer: {example['response_j']}\"\n    return text\n\n\ndef create_datasets(tokenizer, args, seed=None):\n    dataset = load_dataset(\n        args.dataset_name,\n        data_dir=args.subset,\n        split=args.split,\n        use_auth_token=True,\n        num_proc=args.num_workers if not args.streaming else None,\n        streaming=args.streaming,\n    )\n    if args.streaming:\n        print(\"Loading the dataset in streaming mode\")\n        valid_data = dataset.take(args.size_valid_set)\n        train_data = dataset.skip(args.size_valid_set)\n        train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=seed)\n    else:\n        dataset = dataset.train_test_split(test_size=0.005, seed=seed)\n        train_data = dataset[\"train\"]\n        valid_data = dataset[\"test\"]\n        print(f\"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}\")\n\n    chars_per_token = chars_token_ratio(train_data, tokenizer)\n    print(f\"The character to token ratio of the dataset is: {chars_per_token:.2f}\")\n\n    train_dataset = ConstantLengthDataset(\n        tokenizer,\n        train_data,\n        formatting_func=prepare_sample_text,\n        infinite=True,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    valid_dataset = ConstantLengthDataset(\n        tokenizer,\n        valid_data,\n        formatting_func=prepare_sample_text,\n        infinite=False,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    return train_dataset, valid_dataset\n\n\nbnb_config = None\nif script_args.use_bnb:\n    bnb_config = BitsAndBytesConfig(\n        load_in_4bit=True,\n        bnb_4bit_quant_type=\"nf4\",\n        bnb_4bit_compute_dtype=torch.bfloat16,\n    )\n\nbase_model = AutoModelForCausalLM.from_pretrained(\n    script_args.model_name,\n    quantization_config=bnb_config,\n    device_map={\"\": Accelerator().local_process_index},\n    trust_remote_code=True,\n    use_auth_token=True,\n)\nbase_model.config.use_cache = False\n\n\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name, trust_remote_code=True)\ntokenizer.pad_token = tokenizer.eos_token\ntokenizer.padding_side = \"right\"  # Fix weird overflow issue with fp16 training\n\ntrain_dataset, eval_dataset = create_datasets(tokenizer, script_args, seed=training_args.seed)\n\ntrainer = SFTTrainer(\n    model=base_model,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n    peft_config=peft_config,\n    max_seq_length=None,\n    formatting_func=prepare_sample_text,\n    tokenizer=tokenizer,\n    args=training_args,\n)\ntrainer.train()\ntrainer.save_model(training_args.output_dir)\n\noutput_dir = os.path.join(training_args.output_dir, \"final_checkpoint\")\ntrainer.model.save_pretrained(output_dir)\n\n# Free memory for merging weights\ndel base_model\nif is_torch_xpu_available():\n    torch.xpu.empty_cache()\nelif is_torch_npu_available():\n    torch.npu.empty_cache()\nelse:\n    torch.cuda.empty_cache()\n\nmodel = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map=\"auto\", torch_dtype=torch.bfloat16)\nmodel = model.merge_and_unload()\n\noutput_merged_dir = os.path.join(training_args.output_dir, \"final_merged_checkpoint\")\nmodel.save_pretrained(output_merged_dir, safe_serialization=True)\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom transformers import AutoTokenizer, HfArgumentParser, load_tool\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, TextEnvironment\n\n\nos.environ[\"HF_ALLOW_CODE_EVAL\"] = \"1\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\n\n@dataclass\nclass ScriptArguments:\n    model_name: Optional[str] = field(default=\"bigcode/starcoderbase\", metadata={\"help\": \"the model name\"})\n    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n    learning_rate: Optional[float] = field(default=1e-5, metadata={\"help\": \"the learning rate\"})\n    mini_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the PPO minibatch size\"})\n    batch_size: Optional[int] = field(default=32, metadata={\"help\": \"the batch size\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=16, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    max_new_tokens: Optional[int] = field(default=256, metadata={\"help\": \"max number of generated tokens per turn\"})\n    ppo_epochs: Optional[int] = field(default=1, metadata={\"help\": \"max number of ppo epochs\"})\n    iterations: Optional[int] = field(default=1000, metadata={\"help\": \"the number of iterations\"})\n    seed: Optional[int] = field(default=0, metadata={\"help\": \"the random seed\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nargs = parser.parse_args_into_dataclasses()[0]\n\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=32,\n    lora_dropout=0.05,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n    target_modules=[\"c_proj\", \"c_attn\", \"q_attn\"],\n)\n\n# set up models\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\n    args.model_name,\n    use_auth_token=True,\n    trust_remote_code=True,\n    load_in_4bit=True,\n    peft_config=lora_config,\n)\ntokenizer = AutoTokenizer.from_pretrained(args.model_name, use_auth_token=True)\ntokenizer.pad_token = tokenizer.eos_token\n\n# system prompt\nprompt = \"\"\"\\\nAnswer the following question:\n\nQ: In which branch of the arts is Patricia Neary famous?\nA: Ballets\nA2: <request><Wiki>Patricia Neary<call>Patricia Neary (born October 27, 1942) is an American ballerina, choreographer and ballet director, who has been particularly active in Switzerland. She has also been a highly successful ambassador for the Balanchine Trust, bringing George Balanchine's ballets to 60 cities around the globe.<response>\nResult=Ballets<submit>\n\nQ: Who won Super Bowl XX?\nA: Chicago Bears\nA2: <request><Wiki>Super Bowl XX<call>Super Bowl XX was an American football game between the National Football Conference (NFC) champion Chicago Bears and the American Football Conference (AFC) champion New England Patriots to decide the National Football League (NFL) champion for the 1985 season. The Bears defeated the Patriots by the score of 46–10, capturing their first NFL championship (and Chicago's first overall sports victory) since 1963, three years prior to the birth of the Super Bowl. Super Bowl XX was played on January 26, 1986 at the Louisiana Superdome in New Orleans.<response>\nResult=Chicago Bears<submit>\n\nQ: \"\"\"\n\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n    \"eos_token_id\": -1,\n    \"max_new_tokens\": args.max_new_tokens,\n}\n\n# trainer\nconfig = PPOConfig(\n    batch_size=args.batch_size,\n    model_name=args.model_name,\n    learning_rate=args.learning_rate,\n    log_with=args.log_with,\n    mini_batch_size=args.mini_batch_size,\n    ppo_epochs=args.ppo_epochs,\n    gradient_accumulation_steps=args.gradient_accumulation_steps,\n    seed=args.seed,\n    optimize_cuda_cache=True,\n)\nppo_trainer = PPOTrainer(config=config, model=model, tokenizer=tokenizer)\ndataset = load_dataset(\"mandarjoshi/trivia_qa\", \"rc\", split=\"train\")\nlocal_seed = args.seed + ppo_trainer.accelerator.process_index * 100003  # Prime\ndataset = dataset.shuffle(local_seed)\n\n\ndef data_generator():\n    for i in range(len(dataset)):\n        yield dataset[i][\"question\"], list(dataset[i][\"answer\"][\"normalized_aliases\"])\n\n\ngen = data_generator()\ngen = iter(gen)\n\n\ndef generate_data(n):\n    tasks, answers = [], []\n    for _i in range(n):\n        q, a = next(gen)\n        tasks.append(q)\n        answers.append(a)\n    return tasks, answers\n\n\ndef exact_match_reward(responses, answers=None):\n    \"\"\"Reward if generated response contains correct answer.\"\"\"\n    rewards = []\n    for response, answer in zip(responses, answers):\n        reward = 0.0\n        for a in answer:\n            if a.lower() in response.lower():\n                reward += 1.0\n                break\n        rewards.append(torch.tensor(reward))\n    return rewards\n\n\ndef tool_fn(x):\n    # limit the amount of tokens\n    return tool(x).split(\"\\n\")[1][:600]\n\n\n# text env\ntool = load_tool(\"vwxyzjn/pyserini-wikipedia-kilt-doc\")\n\ntext_env = TextEnvironment(\n    model,\n    tokenizer,\n    {\"Wiki\": tool_fn},\n    exact_match_reward,\n    prompt,\n    generation_kwargs=generation_kwargs,\n    max_tool_reponse=400,\n)\n\n\ndef print_trainable_parameters(model):\n    trainable_params = 0\n    all_param = 0\n    for _, param in model.named_parameters():\n        all_param += param.numel()\n        if param.requires_grad:\n            trainable_params += param.numel()\n    print(\n        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n    )\n\n\nprint_trainable_parameters(model)\n# main training loop\nfor i in range(args.iterations):\n    tasks, answers = generate_data(config.batch_size)\n    queries, responses, masks, rewards, histories = text_env.run(tasks, answers=answers)\n    train_stats = ppo_trainer.step(queries, responses, rewards, masks)\n    response_texts = [tokenizer.decode(response) for response in responses]\n    query_texts = [tokenizer.decode(query) for query in queries]\n    texts = {\n        \"query\": [qt.split(\"<submit>\")[-1].strip() for qt in query_texts],\n        \"response\": response_texts,\n        \"answer\": [\", \".join(item) for item in answers],\n    }\n    all_rewards = ppo_trainer.accelerator.gather(torch.tensor(rewards, device=ppo_trainer.accelerator.device))\n    ppo_trainer.log_stats(train_stats, texts, list(all_rewards), columns_to_log=[\"query\", \"response\", \"answer\"])\n    if i % 100 == 0:\n        ppo_trainer.save_pretrained(f\"models/{args.model_name}_{args.seed}_{i}_triviaqa\")\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\n\nimport numpy as np\nimport torch\nfrom transformers import AutoTokenizer, load_tool\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, TextEnvironment\n\n\ndef generate_data(n):\n    \"\"\"Generate random arithmetic tasks and answers.\"\"\"\n    tasks, answers = [], []\n    for _ in range(n):\n        a = np.random.randint(0, 50)\n        b = np.random.randint(0, 50)\n        op = np.random.choice([\"-\", \"+\", \"*\"])\n        tasks.append(f\"\\n\\nWhat is {a} {op} {b}?\")\n        if op == \"-\":\n            answers.append(a - b)\n        elif op == \"+\":\n            answers.append(a + b)\n        else:\n            answers.append(a * b)\n    return tasks, answers\n\n\ndef exact_match_reward(responses, answers=None):\n    \"\"\"Reward if generated response contains correct answer.\"\"\"\n    rewards = []\n    pattern = r\"Result\\s*=\\s*(-?\\d+(?:\\.\\d+)?)\\s*<submit>\"  # generated by chatGPT\n    for response, answer in zip(responses, answers):\n        reward = 0.0\n        predicted_number = None\n        match_pattern = re.findall(pattern, response)\n        if match_pattern:\n            predicted_number = float(match_pattern[0])\n        if predicted_number is not None:\n            if np.abs(predicted_number - answer) < 0.01:\n                reward += 1.0\n        rewards.append(torch.tensor(reward))\n    return rewards\n\n\n# set up models\nmodel_id = \"gpt2\"\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(model_id)\nref_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id)\ntokenizer = AutoTokenizer.from_pretrained(model_id)\ntokenizer.pad_token = tokenizer.eos_token\n\n# system prompt\nprompt = \"\"\"\\\nWhat is 13-3?\n\n<request><SimpleCalculatorTool>13-3<call>10.0<response>\n\nResult=10<submit>\n\nWhat is 4*3?\n\n<request><SimpleCalculatorTool>4*3<call>12.0<response>\n\nResult=12<submit>\"\"\"\n\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n    \"eos_token_id\": -1,\n    \"max_new_tokens\": 32,\n}\n\n# trainer\nppo_config = PPOConfig(\n    batch_size=256,\n    learning_rate=1.41e-5,\n    mini_batch_size=64,\n    log_with=\"wandb\",\n)\nppo_trainer = PPOTrainer(ppo_config, model, ref_model, tokenizer)\n\n# text env\ntext_env = TextEnvironment(\n    model,\n    tokenizer,\n    {\"SimpleCalculatorTool\": load_tool(\"ybelkada/simple-calculator\")},\n    exact_match_reward,\n    prompt,\n    generation_kwargs=generation_kwargs,\n)\n\n# main training loop\nfor _step in range(100):\n    tasks, answers = generate_data(ppo_config.batch_size)\n    queries, responses, masks, rewards, histories = text_env.run(tasks, answers=answers)\n    train_stats = ppo_trainer.step(queries, responses, rewards, masks)\n\n    response_texts = [tokenizer.decode(response) for response in responses]\n    query_texts = [tokenizer.decode(query) for query in queries]\n    texts = {\"query\": [qt.split(\"<submit>\")[-1].strip() for qt in query_texts], \"response\": response_texts}\n    ppo_trainer.log_stats(train_stats, texts, rewards, columns_to_log=[\"query\", \"response\", \"answer\"])\nppo_trainer.save_pretrained(model_id + \"-calculator\")\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport re\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom transformers import AutoTokenizer, HfArgumentParser, load_tool\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, TextEnvironment\n\n\nos.environ[\"HF_ALLOW_CODE_EVAL\"] = \"1\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\n\n@dataclass\nclass ScriptArguments:\n    model_name: Optional[str] = field(default=\"bigcode/starcoderbase\", metadata={\"help\": \"the model name\"})\n    learning_rate: Optional[float] = field(default=1e-5, metadata={\"help\": \"the learning rate\"})\n    mini_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the PPO minibatch size\"})\n    batch_size: Optional[int] = field(default=32, metadata={\"help\": \"the batch size\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=16, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    max_new_tokens: Optional[int] = field(default=256, metadata={\"help\": \"max number of generated tokens per turn\"})\n    ppo_epochs: Optional[int] = field(default=1, metadata={\"help\": \"max number of ppo epochs\"})\n    n_epochs: Optional[int] = field(default=32, metadata={\"help\": \"max number of ppo epochs\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nargs = parser.parse_args_into_dataclasses()[0]\n\n\ndef exact_match_reward(responses, answers=None):\n    \"\"\"Reward if generated response contains correct answer.\"\"\"\n    rewards = []\n    pattern = r\"Result\\s*=\\s*(-?\\d+(?:\\.\\d+)?)\\s*<submit>\"  # generated by chatGPT\n    for response, answer in zip(responses, answers):\n        reward = 0.0\n        try:\n            predicted_number = None\n            match_pattern = re.findall(pattern, response)\n            if match_pattern:\n                predicted_number = float(match_pattern[0])\n            if predicted_number is not None:\n                if np.abs(predicted_number - float(answer)) < 0.1:\n                    reward += 1.0\n        except Exception:\n            pass\n        rewards.append(torch.tensor(reward))\n    return rewards\n\n\ndef evaluate(test_dataloader, text_env, ppo_trainer):\n    test_rewards = []\n    for test_batch in test_dataloader:\n        _, _, _, rewards, _ = text_env.run(test_batch[\"query\"], answers=test_batch[\"answer\"])\n        test_rewards.extend(rewards)\n    test_rewards = ppo_trainer.accelerator.gather_for_metrics(\n        torch.stack(test_rewards).to(ppo_trainer.accelerator.device)\n    )\n    return test_rewards.mean()\n\n\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=32,\n    lora_dropout=0.05,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n    target_modules=[\"c_proj\", \"c_attn\", \"q_attn\"],\n)\n\n# set up models\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\n    args.model_name,\n    use_auth_token=True,\n    load_in_4bit=True,\n    peft_config=lora_config,\n)\ntokenizer = AutoTokenizer.from_pretrained(args.model_name, use_auth_token=True)\ntokenizer.pad_token = tokenizer.eos_token\n\nds = load_dataset(\"openai/gsm8k\", \"main\", split=\"train\")\nds = ds.rename_columns({\"question\": \"query\"})\nds = ds.map(lambda x: {\"answer\": x[\"answer\"].split(\"#### \")[1]})\nds = ds.select(range(1, len(ds)))  # skip the first sample which is used in prompt\n\nds_test = load_dataset(\"openai/gsm8k\", \"main\", split=\"test\")\nds_test = ds_test.rename_columns({\"question\": \"query\"})\nds_test = ds_test.map(lambda x: {\"answer\": x[\"answer\"].split(\"#### \")[1]})\n\ntest_dataloader = torch.utils.data.DataLoader(ds_test, batch_size=args.batch_size)\n\n# prompt\nprompt = \"\"\"\\\nExample of using a Python API to solve math questions.\n\nQ: Olivia has $23. She bought five bagels for $3 each. How much money does she have left?\n\n<request><PythonInterpreter>\ndef solution():\n    money_initial = 23\n    bagels = 5\n    bagel_cost = 3\n    money_spent = bagels * bagel_cost\n    money_left = money_initial - money_spent\n    result = money_left\n    return result\nprint(solution())\n<call>72<response>\n\nResult = 72 <submit>\n\nQ: \"\"\"\n\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n    \"eos_token_id\": -1,\n    \"max_new_tokens\": args.max_new_tokens,\n}\n\n# trainer\nppo_config = PPOConfig(\n    batch_size=args.batch_size,\n    learning_rate=args.learning_rate,\n    mini_batch_size=args.mini_batch_size,\n    ppo_epochs=args.ppo_epochs,\n    gradient_accumulation_steps=args.gradient_accumulation_steps,\n    log_with=\"wandb\",\n    tracker_project_name=\"trl-gsm8k\",\n    remove_unused_columns=False,\n    optimize_cuda_cache=True,\n)\n\nppo_trainer = PPOTrainer(config=ppo_config, model=model, tokenizer=tokenizer, dataset=ds)\ntest_dataloader = ppo_trainer.accelerator.prepare(test_dataloader)\n\n# text env\ntext_env = TextEnvironment(\n    model,\n    tokenizer,\n    [load_tool(\"lvwerra/python-interpreter\")],\n    exact_match_reward,\n    prompt,\n    max_turns=2,\n    generation_kwargs=generation_kwargs,\n)\n\n# main training loop\nfor epoch in range(args.n_epochs):\n    for step, batch in enumerate(ppo_trainer.dataloader):\n        if (step == 0) and (epoch % 4 == 0):  # evaluate every 4 epochs\n            reward_mean_test = evaluate(test_dataloader, text_env, ppo_trainer)\n        else:\n            reward_mean_test = None\n\n        queries, responses, masks, rewards, histories = text_env.run(batch[\"query\"], answers=batch[\"answer\"])\n        train_stats = ppo_trainer.step(queries, responses, rewards, masks)\n\n        # logging\n        if reward_mean_test is not None:\n            train_stats[\"env/reward_mean_test\"] = reward_mean_test\n        texts = {\n            \"query\": batch[\"query\"],\n            \"response\": [tokenizer.decode(response) for response in responses],\n            \"answer\": batch[\"answer\"],\n        }\n        ppo_trainer.log_stats(train_stats, texts, rewards, columns_to_log=[\"query\", \"response\", \"answer\"])\n\nreward_mean_test = evaluate(test_dataloader, text_env, ppo_trainer)\nppo_trainer.save_pretrained(f\"model/{args.model_name}-gsm8k\")\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, List, Optional, Union\n\nimport evaluate\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom datasets import load_dataset\nfrom peft import LoraConfig, TaskType, get_peft_model\nfrom transformers import (\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    HfArgumentParser,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainerCallback,\n    TrainingArguments,\n    set_seed,\n)\nfrom transformers.utils import PaddingStrategy\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    These arguments vary depending on how many GPUs you have, what their capacity and features are, and what size model you want to train.\n    \"\"\"\n\n    local_rank: Optional[int] = field(default=-1, metadata={\"help\": \"Used for multi-gpu\"})\n    resume_from_checkpoint: Optional[bool] = field(\n        default=False,\n        metadata={\"help\": \"If you want to resume training where it left off.\"},\n    )\n    deepspeed: Optional[str] = field(\n        default=None,\n        metadata={\n            \"help\": \"Path to deepspeed config if using deepspeed. You may need this if the model that you want to train doesn't fit on a single GPU.\"\n        },\n    )\n    per_device_train_batch_size: Optional[int] = field(default=4)\n    per_device_eval_batch_size: Optional[int] = field(default=1)\n    gradient_accumulation_steps: Optional[int] = field(default=1)\n    learning_rate: Optional[float] = field(default=2e-5)\n    weight_decay: Optional[float] = field(default=0.001)\n    model_name: Optional[str] = field(\n        default=\"gpt2\",\n        metadata={\n            \"help\": \"The model that you want to train from the Hugging Face hub. E.g. gpt2, gpt2-xl, bert, etc.\"\n        },\n    )\n    tokenizer_name: Optional[str] = field(\n        default=None,\n        metadata={\n            \"help\": \"The tokenizer for your model, if left empty will use the default for your model\",\n        },\n    )\n    bf16: Optional[bool] = field(\n        default=True,\n        metadata={\n            \"help\": \"This essentially cuts the training time in half if you want to sacrifice a little precision and have a supported GPU.\"\n        },\n    )\n    num_train_epochs: Optional[int] = field(\n        default=1,\n        metadata={\"help\": \"The number of training epochs for the reward model.\"},\n    )\n    train_subset: Optional[int] = field(\n        default=100000,\n        metadata={\"help\": \"The size of the subset of the training data to use\"},\n    )\n    eval_subset: Optional[int] = field(\n        default=50000,\n        metadata={\"help\": \"The size of the subset of the eval data to use\"},\n    )\n    gradient_checkpointing: Optional[bool] = field(\n        default=False,\n        metadata={\"help\": \"Enables gradient checkpointing.\"},\n    )\n    optim: Optional[str] = field(\n        default=\"adamw_hf\",\n        metadata={\"help\": \"The optimizer to use.\"},\n    )\n    lr_scheduler_type: Optional[str] = field(\n        default=\"linear\",\n        metadata={\"help\": \"The lr scheduler\"},\n    )\n    max_length: Optional[int] = field(default=512)\n    eval_first_step: Optional[bool] = field(\n        default=False,\n        metadata={\"help\": \"Whether to run eval after the first step\"},\n    )\n    seed: Optional[int] = field(\n        default=0, metadata={\"help\": \"Random seed that will be set at the beginning of training.\"}\n    )\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\nset_seed(script_args.seed)\n# Load the human stack-exchange-paired dataset for tuning the reward model.\ntrain_dataset = load_dataset(\n    \"lvwerra/stack-exchange-paired\", data_dir=\"data/reward\", split=\"train\", verification_mode=\"no_checks\"\n)\nif script_args.train_subset > 0:\n    train_dataset = train_dataset.select(range(script_args.train_subset))\neval_dataset = load_dataset(\n    \"lvwerra/stack-exchange-paired\", data_dir=\"data/evaluation\", split=\"train\", verification_mode=\"no_checks\"\n)\nif script_args.eval_subset > 0:\n    eval_dataset = eval_dataset.select(range(script_args.eval_subset))\n# Define the training args. Needs to be done before the model is loaded if you are using deepspeed.\nmodel_name_split = script_args.model_name.split(\"/\")[-1]\noutput_name = (\n    f\"{model_name_split}_peft_stack-exchange-paired_rmts__{script_args.train_subset}_{script_args.learning_rate}\"\n)\n\ntraining_args = TrainingArguments(\n    output_dir=output_name,\n    learning_rate=script_args.learning_rate,\n    per_device_train_batch_size=script_args.per_device_train_batch_size,\n    per_device_eval_batch_size=script_args.per_device_eval_batch_size,\n    num_train_epochs=script_args.num_train_epochs,\n    weight_decay=script_args.weight_decay,\n    eval_strategy=\"steps\",\n    eval_steps=500,\n    save_strategy=\"steps\",\n    save_steps=500,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    gradient_checkpointing=script_args.gradient_checkpointing,\n    deepspeed=script_args.deepspeed,\n    local_rank=script_args.local_rank,\n    remove_unused_columns=False,\n    label_names=[],\n    bf16=script_args.bf16,\n    logging_strategy=\"steps\",\n    logging_steps=10,\n    optim=script_args.optim,\n    lr_scheduler_type=script_args.lr_scheduler_type,\n    seed=script_args.seed,\n)\n\n\n# Load the value-head model and tokenizer.\ntokenizer_name = script_args.tokenizer_name if script_args.tokenizer_name is not None else script_args.model_name\ntokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_auth_token=True)\ntokenizer.pad_token = tokenizer.eos_token\n\n\npeft_config = LoraConfig(\n    task_type=TaskType.SEQ_CLS,\n    inference_mode=False,\n    r=8,\n    lora_alpha=32,\n    lora_dropout=0.1,\n)\n\nmodel = AutoModelForSequenceClassification.from_pretrained(\n    script_args.model_name, num_labels=1, torch_dtype=torch.bfloat16\n)\nmodel = get_peft_model(model, peft_config)\nmodel.print_trainable_parameters()\n\n# Need to do this for gpt2, because it doesn't have an official pad token.\ntokenizer.pad_token = tokenizer.eos_token\nmodel.config.pad_token_id = tokenizer.eos_token_id\nmodel.config.use_cache = not script_args.gradient_checkpointing\nnum_proc = 24  # Can adjust to be higher if you have more processors.\noriginal_columns = train_dataset.column_names\n\n\n# Turn the dataset into pairs of post + summaries, where text_j is the preferred question + answer and text_k is the other.\n# Then tokenize the dataset.\ndef preprocess_function(examples):\n    new_examples = {\n        \"input_ids_j\": [],\n        \"attention_mask_j\": [],\n        \"input_ids_k\": [],\n        \"attention_mask_k\": [],\n    }\n    for question, response_j, response_k in zip(examples[\"question\"], examples[\"response_j\"], examples[\"response_k\"]):\n        tokenized_j = tokenizer(\"Question: \" + question + \"\\n\\nAnswer: \" + response_j, truncation=True)\n        tokenized_k = tokenizer(\"Question: \" + question + \"\\n\\nAnswer: \" + response_k, truncation=True)\n\n        new_examples[\"input_ids_j\"].append(tokenized_j[\"input_ids\"])\n        new_examples[\"attention_mask_j\"].append(tokenized_j[\"attention_mask\"])\n        new_examples[\"input_ids_k\"].append(tokenized_k[\"input_ids\"])\n        new_examples[\"attention_mask_k\"].append(tokenized_k[\"attention_mask\"])\n\n    return new_examples\n\n\n# preprocess the dataset and filter out QAs that are longer than script_args.max_length\ntrain_dataset = train_dataset.map(\n    preprocess_function,\n    batched=True,\n    num_proc=num_proc,\n    remove_columns=original_columns,\n)\ntrain_dataset = train_dataset.filter(\n    lambda x: len(x[\"input_ids_j\"]) <= script_args.max_length and len(x[\"input_ids_k\"]) <= script_args.max_length,\n    num_proc=num_proc,\n)\n\neval_dataset = eval_dataset.map(\n    preprocess_function,\n    batched=True,\n    num_proc=num_proc,\n    remove_columns=original_columns,\n)\neval_dataset = eval_dataset.filter(\n    lambda x: len(x[\"input_ids_j\"]) <= script_args.max_length and len(x[\"input_ids_k\"]) <= script_args.max_length,\n    num_proc=num_proc,\n)\n\n\n# We need to define a special data collator that batches the data in our j vs k format.\n@dataclass\nclass RewardDataCollatorWithPadding:\n    tokenizer: PreTrainedTokenizerBase\n    padding: Union[bool, str, PaddingStrategy] = True\n    max_length: Optional[int] = None\n    pad_to_multiple_of: Optional[int] = None\n    return_tensors: str = \"pt\"\n\n    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:\n        features_j = []\n        features_k = []\n        for feature in features:\n            features_j.append(\n                {\n                    \"input_ids\": feature[\"input_ids_j\"],\n                    \"attention_mask\": feature[\"attention_mask_j\"],\n                }\n            )\n            features_k.append(\n                {\n                    \"input_ids\": feature[\"input_ids_k\"],\n                    \"attention_mask\": feature[\"attention_mask_k\"],\n                }\n            )\n        batch_j = self.tokenizer.pad(\n            features_j,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=self.return_tensors,\n        )\n        batch_k = self.tokenizer.pad(\n            features_k,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=self.return_tensors,\n        )\n        batch = {\n            \"input_ids_j\": batch_j[\"input_ids\"],\n            \"attention_mask_j\": batch_j[\"attention_mask\"],\n            \"input_ids_k\": batch_k[\"input_ids\"],\n            \"attention_mask_k\": batch_k[\"attention_mask\"],\n            \"return_loss\": True,\n        }\n        return batch\n\n\n# Define the metric that we'll use for validation.\naccuracy = evaluate.load(\"accuracy\")\n\n\ndef compute_metrics(eval_pred):\n    predictions, _ = eval_pred\n    # Here, predictions is rewards_j and rewards_k.\n    # We want to see how much of the time rewards_j > rewards_k.\n    predictions = np.argmax(predictions, axis=0)\n    labels = np.zeros(predictions.shape)\n    return accuracy.compute(predictions=predictions, references=labels)\n\n\nclass RewardTrainer(Trainer):\n    # Define how to compute the reward loss. We use the InstructGPT pairwise logloss: https://huggingface.co/papers/2203.02155\n    def compute_loss(self, model, inputs, return_outputs=False):\n        rewards_j = model(input_ids=inputs[\"input_ids_j\"], attention_mask=inputs[\"attention_mask_j\"])[0]\n        rewards_k = model(input_ids=inputs[\"input_ids_k\"], attention_mask=inputs[\"attention_mask_k\"])[0]\n        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()\n        if return_outputs:\n            return loss, {\"rewards_j\": rewards_j, \"rewards_k\": rewards_k}\n        return loss\n\n\n# Train the model, woohoo.\ntrainer = RewardTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n    compute_metrics=compute_metrics,\n    data_collator=RewardDataCollatorWithPadding(tokenizer=tokenizer, max_length=script_args.max_length),\n)\n\n\nif script_args.eval_first_step:\n\n    class EvaluateFirstStepCallback(TrainerCallback):\n        def on_step_end(self, args, state, control, **kwargs):\n            if state.global_step == 1:\n                control.should_evaluate = True\n\n    trainer.add_callback(EvaluateFirstStepCallback())\n\ntrainer.train(script_args.resume_from_checkpoint)\n\nprint(\"Saving last checkpoint of the model\")\nmodel.save_pretrained(output_name + \"_peft_last_checkpoint\")\n\n\n# RLHF pipeline for the creation of StackLLaMa: a Stack exchange llama-7b model.\nThere were three main steps to the training process:\n1. Supervised fine-tuning of the base llama-7b model to create llama-7b-se:\n    - `torchrun --nnodes 1  --nproc_per_node 8 examples/research_projects/stack_llama/scripts/supervised_finetuning.py --model_path=<LLAMA_MODEL_PATH> --streaming --learning_rate 1e-5 --max_steps 5000 --output_dir ./llama-se`\n2. Reward modeling using dialog pairs from the SE dataset using the llama-7b-se to create llama-7b-se-rm:\n    - `torchrun --nnodes 1  --nproc_per_node 8 examples/research_projects/stack_llama/scripts/reward_modeling.py --model_name=<LLAMA_SE_MODEL>`\n3. RL fine-tuning of llama-7b-se with the llama-7b-se-rm reward model:\n    - `accelerate launch --multi_gpu --num_machines 1  --num_processes 8 examples/research_projects/stack_llama/scripts/rl_training.py --log_with=wandb --model_name=<LLAMA_SE_MODEL> --reward_model_name=<LLAMA_SE_RM_MODEL> --adafactor=False --tokenizer_name=<LLAMA_TOKENIZER> --save_freq=100 --output_max_length=128 --batch_size=8 --gradient_accumulation_steps=8 --batched_gen=True --ppo_epochs=4 --seed=0 --learning_rate=1.4e-5 --early_stopping=True --output_dir=llama-se-rl-finetune-128-8-8-1.4e-5_adam`\n\n\nLoRA layers were using at all stages to reduce memory requirements. \nAt each stage the peft adapter layers were merged with the base model, using: \n```shell\npython examples/research_projects/stack_llama/scripts/merge_peft_adapter.py --adapter_model_name=XXX --base_model_name=YYY --output_name=ZZZ\n```\nNote that this script requires `peft>=0.3.0`.\n\nFor access to the base llama-7b model, please see Meta's [release](https://ai.facebook.com/blog/large-language-model-llama-meta-ai/) and [request form](https://docs.google.com/forms/d/e/1FAIpQLSfqNECQnMkycAp2jP4Z9TFX0cGR4uf7b_fBxjY_OjhJILlKGA/viewform).\n\n\n# Copyright 2022 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom accelerate import Accelerator\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import Adafactor, AutoTokenizer, HfArgumentParser, pipeline\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed\nfrom trl.core import LengthSampler\n\n\ntqdm.pandas()\n\n\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The name of the Casual LM model we wish to fine-tune with PPO\n    \"\"\"\n\n    # NOTE: gpt2 models use Conv1D instead of Linear layers which are not yet supported in 8 bit mode\n    # models like gpt-neo* models are more suitable.\n    model_name: Optional[str] = field(default=\"\", metadata={\"help\": \"the model name\"})\n    tokenizer_name: Optional[str] = field(default=\"\", metadata={\"help\": \"the tokenizer name\"})\n    reward_model_name: Optional[str] = field(default=\"\", metadata={\"help\": \"the reward model name\"})\n    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n    learning_rate: Optional[float] = field(default=1.41e-5, metadata={\"help\": \"the learning rate\"})\n    output_max_length: Optional[int] = field(default=128, metadata={\"help\": \"maximum length for generation\"})\n    mini_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the PPO minibatch size\"})\n    batch_size: Optional[int] = field(default=32, metadata={\"help\": \"the batch size\"})\n    ppo_epochs: Optional[int] = field(default=4, metadata={\"help\": \"the number of ppo epochs\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=4, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    adafactor: Optional[bool] = field(default=False, metadata={\"help\": \"whether to use the adafactor optimizer\"})\n    early_stopping: Optional[bool] = field(default=False, metadata={\"help\": \"whether to early stop\"})\n    target_kl: Optional[float] = field(default=0.1, metadata={\"help\": \"kl target for early stopping\"})\n    reward_baseline: Optional[float] = field(\n        default=0.0,\n        metadata={\"help\": \"a baseline value that is subtracted from the reward\"},\n    )\n    batched_gen: Optional[bool] = field(default=False, metadata={\"help\": \"whether to use the batched text gen\"})\n    save_freq: Optional[int] = field(default=None, metadata={\"help\": \"n steps to save the model\"})\n    output_dir: Optional[str] = field(default=\"runs/\", metadata={\"help\": \"n steps to save the model\"})\n    seed: Optional[int] = field(default=0, metadata={\"help\": \"the seed\"})\n    steps: Optional[int] = field(default=20000, metadata={\"help\": \"number of epochs\"})\n    init_kl_coef: Optional[float] = field(\n        default=0.2,\n        metadata={\"help\": \"Initial KL penalty coefficient (used for adaptive and linear control)\"},\n    )\n\n    adap_kl_ctrl: Optional[bool] = field(default=True, metadata={\"help\": \"Use adaptive KL control, otherwise linear\"})\n    load_in_8bit: Optional[bool] = field(default=True, metadata={\"help\": \"whether to load the model in 8bit\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args: ScriptArguments = parser.parse_args_into_dataclasses()[0]\nreward_model_name = script_args.reward_model_name\ndataset_name = \"lvwerra/stack-exchange-paired\"\nconfig = PPOConfig(\n    steps=script_args.steps,\n    model_name=script_args.model_name,\n    learning_rate=script_args.learning_rate,\n    log_with=script_args.log_with,\n    batch_size=script_args.batch_size,\n    mini_batch_size=script_args.mini_batch_size,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    optimize_cuda_cache=True,\n    early_stopping=script_args.early_stopping,\n    target_kl=script_args.target_kl,\n    ppo_epochs=script_args.ppo_epochs,\n    seed=script_args.seed,\n    init_kl_coef=script_args.init_kl_coef,\n    adap_kl_ctrl=script_args.adap_kl_ctrl,\n)\n\ntrain_dataset = load_dataset(\n    \"lvwerra/stack-exchange-paired\", data_dir=\"data/rl\", split=\"train\", verification_mode=\"no_checks\"\n)\ntrain_dataset = train_dataset.select(range(100000))\noriginal_columns = train_dataset.column_names\n\n# We then define the arguments to pass to the sentiment analysis pipeline.\n# We set `return_all_scores` to True to get the sentiment score for each token.\nsent_kwargs = {\n    \"return_all_scores\": True,\n    \"function_to_apply\": \"none\",\n    \"batch_size\": 16,\n    \"truncation\": True,\n}\n\ntokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name)\n# GPT-2 tokenizer has a pad token, but it is not eos_token by default. We need to set it to eos_token.\n# only for this model.\n\nif getattr(tokenizer, \"pad_token\", None) is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\n\n# Below is an example function to build the dataset. In our case, we use the IMDB dataset\n# from the `datasets` library. One should customize this function to train the model on\n# its own dataset.\ndef build_dataset(\n    tokenizer,\n    dataset_name=\"lvwerra/stack-exchange-paired\",\n):\n    \"\"\"\n    Build dataset for training. This builds the dataset from `load_dataset`, one should\n    customize this function to train the model on its own dataset.\n\n    Args:\n        dataset_name (`str`):\n            The name of the dataset to be loaded.\n\n    Returns:\n        dataloader (`torch.utils.data.DataLoader`):\n            The dataloader for the dataset.\n    \"\"\"\n\n    num_proc = 24\n\n    def preprocess_function(examples):\n        new_examples = {\n            \"query\": [],\n            \"input_ids\": [],\n        }\n        for question in examples[\"question\"]:\n            query = \"Question: \" + question + \"\\n\\nAnswer: \"\n            tokenized_question = tokenizer(query, truncation=True)\n            new_examples[\"query\"].append(query)\n            new_examples[\"input_ids\"].append(tokenized_question[\"input_ids\"])\n\n        return new_examples\n\n    ds = train_dataset.map(\n        preprocess_function,\n        batched=True,\n        num_proc=num_proc,\n        remove_columns=original_columns,\n    )\n    ds = ds.filter(lambda x: len(x[\"input_ids\"]) < 512, batched=False, num_proc=num_proc)\n\n    ds.set_format(type=\"torch\")\n    return ds\n\n\n# We retrieve the dataloader by calling the `build_dataset` function.\ndataset = build_dataset(tokenizer)\n\n\ndef collator(data):\n    return {key: [d[key] for d in data] for key in data[0]}\n\n\n# set seed before initializing value head for deterministic eval\nset_seed(config.seed)\n\n# Now let's build the model, the reference model, and the tokenizer.\ncurrent_device = Accelerator().local_process_index\n\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=32,\n    lora_dropout=0.05,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\n    config.model_name,\n    load_in_8bit=script_args.load_in_8bit,\n    device_map={\"\": current_device},\n    peft_config=lora_config,\n)\n\noptimizer = None\nif script_args.adafactor:\n    optimizer = Adafactor(\n        filter(lambda p: p.requires_grad, model.parameters()),\n        scale_parameter=False,\n        relative_step=False,\n        warmup_init=False,\n        lr=config.learning_rate,\n    )\n# We then build the PPOTrainer, passing the model, the reference model, the tokenizer\nppo_trainer = PPOTrainer(\n    config,\n    model,\n    ref_model=None,\n    tokenizer=tokenizer,\n    dataset=dataset,\n    data_collator=collator,\n    optimizer=optimizer,\n)\n\n# We then build the sentiment analysis pipeline using our reward model, passing the\n# model name and the sentiment analysis pipeline arguments. Let's also make sure to\n# set the device to the same device as the PPOTrainer.\ndevice = ppo_trainer.accelerator.device\nif ppo_trainer.accelerator.num_processes == 1:\n    device = 0 if torch.cuda.is_available() else \"cpu\"  # to avoid a ` pipeline` bug\nsentiment_pipe = pipeline(\n    \"sentiment-analysis\",\n    model=reward_model_name,\n    device_map={\"\": current_device},\n    model_kwargs={\"load_in_8bit\": script_args.load_in_8bit},\n    tokenizer=tokenizer,\n    return_token_type_ids=False,\n)\n\nif sentiment_pipe.model.config.pad_token_id is None:\n    sentiment_pipe.model.config.pad_token_id = sentiment_pipe.model.config.eos_token_id\n# We then define the arguments to pass to the `generate` function. These arguments\n# are passed to the `generate` function of the PPOTrainer, which is a wrapper around\n# the `generate` function of the trained model.\ngeneration_kwargs = {\n    # \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.pad_token_id,\n    \"eos_token_id\": 100_000,\n}\noutput_min_length = 32\noutput_max_length = script_args.output_max_length\noutput_length_sampler = LengthSampler(output_min_length, output_max_length)\n\nfor epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n    if epoch >= config.total_ppo_epochs:\n        break\n\n    question_tensors = batch[\"input_ids\"]\n\n    response_tensors = ppo_trainer.generate(\n        question_tensors,\n        return_prompt=False,\n        length_sampler=output_length_sampler,\n        **generation_kwargs,\n    )\n    batch[\"response\"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)\n\n    # Compute reward score (using the sentiment analysis pipeline)\n    texts = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]\n    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)\n    rewards = [torch.tensor(output[0][\"score\"] - script_args.reward_baseline) for output in pipe_outputs]\n\n    # Run PPO step\n    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)\n    ppo_trainer.log_stats(stats, batch, rewards)\n\n    if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:\n        ppo_trainer.save_pretrained(script_args.output_dir + f\"step_{epoch}\")\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom peft import PeftConfig, PeftModel\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The input names representing the Adapter and Base model fine-tuned with PEFT, and the output name representing the\n    merged model.\n    \"\"\"\n\n    adapter_model_name: Optional[str] = field(default=None, metadata={\"help\": \"the adapter name\"})\n    base_model_name: Optional[str] = field(default=None, metadata={\"help\": \"the base model name\"})\n    output_name: Optional[str] = field(default=None, metadata={\"help\": \"the merged model name\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\nassert script_args.adapter_model_name is not None, \"please provide the name of the Adapter you would like to merge\"\nassert script_args.base_model_name is not None, \"please provide the name of the Base model\"\nassert script_args.output_name is not None, \"please provide the output name of the merged model\"\n\npeft_config = PeftConfig.from_pretrained(script_args.adapter_model_name)\nif peft_config.task_type == \"SEQ_CLS\":\n    # The sequence classification task is used for the reward model in PPO\n    model = AutoModelForSequenceClassification.from_pretrained(\n        script_args.base_model_name, num_labels=1, torch_dtype=torch.bfloat16\n    )\nelse:\n    model = AutoModelForCausalLM.from_pretrained(\n        script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16\n    )\n\ntokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name)\n\n# Load the PEFT model\nmodel = PeftModel.from_pretrained(model, script_args.adapter_model_name)\nmodel.eval()\n\nmodel = model.merge_and_unload()\n\nmodel.save_pretrained(f\"{script_args.output_name}\")\ntokenizer.save_pretrained(f\"{script_args.output_name}\")\nmodel.push_to_hub(f\"{script_args.output_name}\", use_temp_dir=False)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport os\n\nfrom accelerate import Accelerator\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, logging, set_seed\n\nfrom trl import SFTTrainer\nfrom trl.trainer import ConstantLengthDataset\n\n\n\"\"\"\nFine-Tune Llama-7b on SE paired dataset\n\"\"\"\n\n\ndef get_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--model_path\", type=str, default=\"\")\n    parser.add_argument(\"--dataset_name\", type=str, default=\"lvwerra/stack-exchange-paired\")\n    parser.add_argument(\"--subset\", type=str, default=\"data/finetune\")\n    parser.add_argument(\"--split\", type=str, default=\"train\")\n    parser.add_argument(\"--size_valid_set\", type=int, default=4000)\n    parser.add_argument(\"--streaming\", action=\"store_true\")\n    parser.add_argument(\"--shuffle_buffer\", type=int, default=5000)\n\n    parser.add_argument(\"--seq_length\", type=int, default=1024)\n    parser.add_argument(\"--max_steps\", type=int, default=10000)\n    parser.add_argument(\"--batch_size\", type=int, default=4)\n    parser.add_argument(\"--gradient_accumulation_steps\", type=int, default=1)\n    parser.add_argument(\"--eos_token_id\", type=int, default=49152)\n\n    parser.add_argument(\"--learning_rate\", type=float, default=1e-4)\n    parser.add_argument(\"--lr_scheduler_type\", type=str, default=\"cosine\")\n    parser.add_argument(\"--num_warmup_steps\", type=int, default=100)\n    parser.add_argument(\"--weight_decay\", type=float, default=0.05)\n\n    parser.add_argument(\"--local_rank\", type=int, default=0)\n    parser.add_argument(\"--fp16\", action=\"store_true\", default=False)\n    parser.add_argument(\"--bf16\", action=\"store_true\", default=False)\n    parser.add_argument(\"--gradient_checkpointing\", action=\"store_true\", default=False)\n    parser.add_argument(\"--seed\", type=int, default=0)\n    parser.add_argument(\"--num_workers\", type=int, default=None)\n    parser.add_argument(\"--output_dir\", type=str, default=\"./checkpoints\")\n    parser.add_argument(\"--log_freq\", default=1, type=int)\n    parser.add_argument(\"--eval_freq\", default=1000, type=int)\n    parser.add_argument(\"--save_freq\", default=1000, type=int)\n\n    return parser.parse_args()\n\n\ndef chars_token_ratio(dataset, tokenizer, nb_examples=400):\n    \"\"\"\n    Estimate the average number of characters per token in the dataset.\n    \"\"\"\n    total_characters, total_tokens = 0, 0\n    for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):\n        text = prepare_sample_text(example)\n        total_characters += len(text)\n        if tokenizer.is_fast:\n            total_tokens += len(tokenizer(text).tokens())\n        else:\n            total_tokens += len(tokenizer.tokenize(text))\n\n    return total_characters / total_tokens\n\n\ndef print_trainable_parameters(model):\n    \"\"\"\n    Prints the number of trainable parameters in the model.\n    \"\"\"\n    trainable_params = 0\n    all_param = 0\n    for _, param in model.named_parameters():\n        all_param += param.numel()\n        if param.requires_grad:\n            trainable_params += param.numel()\n    print(\n        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n    )\n\n\ndef prepare_sample_text(example):\n    \"\"\"Prepare the text from a sample of the dataset.\"\"\"\n    text = f\"Question: {example['question']}\\n\\nAnswer: {example['response_j']}\"\n    return text\n\n\ndef create_datasets(tokenizer, args):\n    dataset = load_dataset(\n        args.dataset_name,\n        data_dir=args.subset,\n        split=args.split,\n        use_auth_token=True,\n        num_proc=args.num_workers if not args.streaming else None,\n        streaming=args.streaming,\n    )\n    if args.streaming:\n        print(\"Loading the dataset in streaming mode\")\n        valid_data = dataset.take(args.size_valid_set)\n        train_data = dataset.skip(args.size_valid_set)\n        train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)\n    else:\n        dataset = dataset.train_test_split(test_size=0.005, seed=args.seed)\n        train_data = dataset[\"train\"]\n        valid_data = dataset[\"test\"]\n        print(f\"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}\")\n\n    chars_per_token = chars_token_ratio(train_data, tokenizer)\n    print(f\"The character to token ratio of the dataset is: {chars_per_token:.2f}\")\n\n    train_dataset = ConstantLengthDataset(\n        tokenizer,\n        train_data,\n        formatting_func=prepare_sample_text,\n        infinite=True,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    valid_dataset = ConstantLengthDataset(\n        tokenizer,\n        valid_data,\n        formatting_func=prepare_sample_text,\n        infinite=False,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    return train_dataset, valid_dataset\n\n\ndef run_training(args, train_data, val_data):\n    print(\"Loading the model\")\n\n    lora_config = LoraConfig(\n        r=16,\n        lora_alpha=32,\n        lora_dropout=0.05,\n        bias=\"none\",\n        task_type=\"CAUSAL_LM\",\n    )\n\n    train_data.start_iteration = 0\n\n    print(\"Starting main loop\")\n\n    training_args = TrainingArguments(\n        output_dir=args.output_dir,\n        dataloader_drop_last=True,\n        eval_strategy=\"steps\",\n        max_steps=args.max_steps,\n        eval_steps=args.eval_freq,\n        save_steps=args.save_freq,\n        logging_steps=args.log_freq,\n        per_device_train_batch_size=args.batch_size,\n        per_device_eval_batch_size=args.batch_size,\n        learning_rate=args.learning_rate,\n        lr_scheduler_type=args.lr_scheduler_type,\n        warmup_steps=args.num_warmup_steps,\n        gradient_accumulation_steps=args.gradient_accumulation_steps,\n        gradient_checkpointing=args.gradient_checkpointing,\n        fp16=args.fp16,\n        bf16=args.bf16,\n        weight_decay=args.weight_decay,\n        run_name=\"llama-7b-finetuned\",\n        report_to=\"wandb\",\n        ddp_find_unused_parameters=False,\n    )\n\n    model = AutoModelForCausalLM.from_pretrained(\n        args.model_path, load_in_8bit=True, device_map={\"\": Accelerator().process_index}\n    )\n\n    trainer = SFTTrainer(\n        model=model,\n        args=training_args,\n        train_dataset=train_data,\n        eval_dataset=val_data,\n        peft_config=lora_config,\n        packing=True,\n    )\n\n    print_trainable_parameters(trainer.model)\n\n    print(\"Training...\")\n    trainer.train()\n\n    print(\"Saving last checkpoint of the model\")\n    trainer.model.save_pretrained(os.path.join(args.output_dir, \"final_checkpoint/\"))\n\n\ndef main(args):\n    tokenizer = AutoTokenizer.from_pretrained(args.model_path)\n    train_dataset, eval_dataset = create_datasets(tokenizer, args)\n    run_training(args, train_dataset, eval_dataset)\n\n\nif __name__ == \"__main__\":\n    args = get_args()\n    assert args.model_path != \"\", \"Please provide the llama model path\"\n\n    set_seed(args.seed)\n    os.makedirs(args.output_dir, exist_ok=True)\n\n    logging.set_verbosity_error()\n\n    main(args)\n\n\n# De-detoxifying language models\n\nTo run this code, do the following:\n\n```shell\nACCELERATE_LOG_LEVEL=info accelerate launch --config_file {CONFIG} examples/research_projects/toxicity/scripts/gpt-j-6b-toxicity.py --log_with wandb\n```\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import Adam\nfrom tqdm import tqdm\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoTokenizer,\n    HfArgumentParser,\n    RobertaForSequenceClassification,\n    RobertaTokenizer,\n)\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, create_reference_model, set_seed\nfrom trl.core import LengthSampler\n\n\ntqdm.pandas()\n\n########################################################################\n# This is a fully working simple example to use trl with accelerate.\n#\n# This example fine-tunes a GPTJ model to generate less toxic contents\n# by using allenai/real-toxicity-prompts dataset. We use PPO\n#  (proximal policy optimization) to optimize the model.\n# in any of the following settings (with the same script):\n#   - single CPU or single GPU\n#   - multi GPUS (using PyTorch distributed mode)\n#   - multi GPUS (using DeepSpeed ZeRO-Offload stages 1 & 2)\n#   - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, first initialize the accelerate\n# configuration with `accelerate config`\n#\n########################################################################\n\n\n# We first define the configuration of the experiment, defining the model, the dataset,\n# the training parameters, and the PPO parameters.\n# Check the default arguments in the `PPOConfig` class for more details.\n# If you want to log with tensorboard, add the kwarg\n# `project_kwargs={\"logging_dir\": PATH_TO_LOGS}` to the PPOConfig.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The name of the Casual LM model we wish to fine-tune with PPO\n    \"\"\"\n\n    # NOTE: gpt2 models use Conv1D instead of Linear layers which are not yet supported in 8 bit mode\n    # models like gpt-neo* models are more suitable.\n    model_name: Optional[str] = field(default=\"ybelkada/gpt-j-6b-sharded-bf16\", metadata={\"help\": \"the model name\"})\n    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n    learning_rate: Optional[float] = field(default=(1.47e-5) * 2, metadata={\"help\": \"the learning rate\"})\n    mini_batch_size: Optional[int] = field(default=4, metadata={\"help\": \"the PPO minibatch size\"})\n    batch_size: Optional[int] = field(default=16, metadata={\"help\": \"the batch size\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=1, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    model_save_path: Optional[str] = field(\n        default=\"./gpt-j-6B-detoxified-long-context-26-shl-1e4-final\",\n        metadata={\"help\": \"the path to save the model\"},\n    )\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\nconfig = PPOConfig(\n    model_name=script_args.model_name,\n    learning_rate=script_args.learning_rate,\n    log_with=script_args.log_with,\n    ppo_epochs=100,\n    mini_batch_size=script_args.mini_batch_size,\n    batch_size=script_args.batch_size,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n)\n\n\n# Below is an example function to build the dataset. In our case, we use the IMDB dataset\n# from the `datasets` library. One should customize this function to train the model on\n# its own dataset.\ndef build_dataset(\n    config, dataset_name=\"allenai/real-toxicity-prompts\", input_min_text_length=5, input_max_text_length=10\n):\n    \"\"\"\n    Build dataset for training. This builds the dataset from `load_dataset`, one should\n    customize this function to train the model on its own dataset.\n\n    Args:\n        dataset_name (`str`):\n            The name of the dataset to be loaded.\n\n    Returns:\n        dataloader (`torch.utils.data.DataLoader`):\n            The dataloader for the dataset.\n    \"\"\"\n    tokenizer = AutoTokenizer.from_pretrained(config.model_name)\n    tokenizer.pad_token = tokenizer.eos_token\n\n    ds = load_dataset(dataset_name, split=\"train\")\n\n    def filter_fn(sample):\n        toxicity = sample[\"prompt\"][\"toxicity\"]\n        return toxicity is not None and toxicity > 0.3\n\n    ds = ds.filter(filter_fn, batched=False)\n\n    input_size = LengthSampler(input_min_text_length, input_max_text_length)\n\n    def tokenize(sample):\n        prompt = sample[\"prompt\"][\"text\"]\n        continuation = sample[\"continuation\"][\"text\"]\n\n        sample[\"input_ids\"] = tokenizer.encode(prompt + continuation)[: input_size()]\n        sample[\"query\"] = tokenizer.decode(sample[\"input_ids\"])\n        return sample\n\n    ds = ds.map(tokenize, batched=False)\n    ds.set_format(type=\"torch\")\n\n    ds = ds.train_test_split(test_size=0.2, shuffle=False)[\"train\"]\n\n    return ds\n\n\n# We retrieve the dataloader by calling the `build_dataset` function.\nmin_input_length = 30\nmax_input_length = 40\ndataset = build_dataset(config, input_min_text_length=min_input_length, input_max_text_length=max_input_length)\n\n\ndef collator(data):\n    return {key: [d[key] for d in data] for key in data[0]}\n\n\n# set seed before initializing value head for deterministic eval\nset_seed(config.seed)\n\n# Now let's build the model, the reference model, and the tokenizer. We first load the model\n# in bfloat16 to save memory using `transformers`.\nmodel = AutoModelForCausalLM.from_pretrained(config.model_name, torch_dtype=torch.bfloat16)\n# And then we pass the loaded model to `AutoModelForCausalLMWithValueHead`.\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(model)\n\n# We create a reference model by sharing 20 layers\nref_model = create_reference_model(model, num_shared_layers=20)\n\n# We make sure to use `Adam` optimizer on the model parameters that require gradients.\noptimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)\n\n# GPT-2 / GPT-J tokenizer has a pad token, but it is not eos_token by default. We need to set it to eos_token.\n# only for this model.\ntokenizer = AutoTokenizer.from_pretrained(config.model_name)\ntokenizer.pad_token = tokenizer.eos_token\n\n# We then build the PPOTrainer, passing the model, the reference model, the tokenizer\nppo_trainer = PPOTrainer(\n    config,\n    model,\n    ref_model=ref_model,\n    tokenizer=tokenizer,\n    dataset=dataset,\n    data_collator=collator,\n    optimizer=optimizer,\n)\n\n# We then build the reward pipeline, we will use the toxicity model to compute the reward.\n# We first load the toxicity model and tokenizer.\ntoxicity_model_id = \"facebook/roberta-hate-speech-dynabench-r4-target\"\ntoxicity_tokenizer = RobertaTokenizer.from_pretrained(toxicity_model_id)\n# We load the toxicity model in fp16 to save memory.\ntoxicity_model = RobertaForSequenceClassification.from_pretrained(toxicity_model_id, torch_dtype=torch.float16).to(\n    ppo_trainer.accelerator.device\n)\n\n\n# We then define the arguments to pass to the `generate` function. These arguments\n# are passed to the `generate` function of the PPOTrainer, which is a wrapper around\n# the `generate` function of the trained model.\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n}\noutput_min_length = 20\noutput_max_length = 30\noutput_length_sampler = LengthSampler(output_min_length, output_max_length)\n\nmodel_save_path = script_args.model_save_path\n\nfor epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n    query_tensors = batch[\"input_ids\"]\n\n    # Get response from the policy model\n    response_tensors = []\n    for query in query_tensors:\n        gen_len = output_length_sampler()\n        generation_kwargs[\"max_new_tokens\"] = gen_len\n        response = ppo_trainer.generate(query, **generation_kwargs)\n        response_tensors.append(response.squeeze()[-gen_len:])\n    batch[\"response\"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]\n\n    # Compute sentiment score\n    texts = batch[\"response\"]\n    toxicity_inputs = toxicity_tokenizer(texts, padding=True, truncation=True, return_tensors=\"pt\").to(\n        ppo_trainer.accelerator.device\n    )\n    logits = toxicity_model(**toxicity_inputs).logits.float()\n    toxicity_labels = (logits[:, 0]).tolist()\n\n    rewards = [torch.tensor(output) for output in toxicity_labels]\n\n    # Run PPO step\n    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)\n    ppo_trainer.log_stats(stats, batch, rewards)\n\n    # Save model every 100 epochs\n    if epoch % 100 == 0:\n        if ppo_trainer.accelerator.is_main_process:\n            ppo_trainer.save_pretrained(model_save_path)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport csv\n\nimport evaluate\nimport numpy as np\nimport torch\nfrom datasets import load_dataset\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, is_torch_npu_available, is_torch_xpu_available\n\n\ntoxicity = evaluate.load(\"ybelkada/toxicity\", \"DaNLP/da-electra-hatespeech-detection\", module_type=\"measurement\")\nds = load_dataset(\"OxAISH-AL-LLM/wiki_toxic\", split=\"test\")\n\nparser = argparse.ArgumentParser(description=\"Evaluate de-toxified models\")\nparser.add_argument(\"--model_type\", default=\"all\", type=str, help=\"Relative path to the source model folder\")\nparser.add_argument(\"--output_file\", default=\"toxicity.csv\", type=str, help=\"Relative path to the source model folder\")\nparser.add_argument(\"--batch_size\", default=64, type=int, help=\"Batch size\")\nparser.add_argument(\"--num_samples\", default=400, type=int, help=\"Number of samples\")\nparser.add_argument(\"--context_length\", default=2000, type=int, help=\"Number of samples\")\nparser.add_argument(\"--max_new_tokens\", default=30, type=int, help=\"Max new tokens for generation\")\nargs = parser.parse_args()\n\n\nif args.model_type == \"all\":\n    MODELS_TO_TEST = [\n        \"ybelkada/gpt-neo-125m-detox\",\n        \"EleutherAI/gpt-neo-125M\",\n        \"EleutherAI/gpt-neo-2.7B\",\n        \"ybelkada/gpt-neo-2.7B-detox\",\n        \"ybelkada/gpt-j-6b-sharded-bf16\",\n        \"ybelkada/gpt-j-6b-detoxs\",\n    ]\nelif args.model_type == \"gpt-neo\":\n    MODELS_TO_TEST = [\n        \"ybelkada/gpt-neo-125m-detox\",\n        \"EleutherAI/gpt-neo-125M\",\n        \"EleutherAI/gpt-neo-2.7B\",\n        \"ybelkada/gpt-neo-2.7B-detox\",\n    ]\nelif args.model_type == \"gpt-j\":\n    MODELS_TO_TEST = [\n        \"ybelkada/gpt-j-6b-sharded-bf16\",\n        \"ybelkada/gpt-j-6b-detox\",\n    ]\nelse:\n    MODELS_TO_TEST = [args.model_type]\nNUM_SAMPLES = args.num_samples\nBATCH_SIZE = args.batch_size\noutput_file = args.output_file\nmax_new_tokens = args.max_new_tokens\ncontext_length = args.context_length\nif is_torch_xpu_available():\n    device = torch.xpu.current_device()\nelif is_torch_npu_available():\n    device = torch.npu.current_device()\nelse:\n    device = torch.cuda.current_device() if torch.cuda.is_available() else \"cpu\"\n\n# consider only toxic prompts\nds = ds.filter(lambda x: x[\"label\"] == 1)\n\ntoxicities = {}\n\n# open a csv file\nfile = open(f\"{output_file}\", \"w\", newline=\"\")\nwriter = csv.writer(file)\n# add first rows\nwriter.writerow([\"model_id\", \"mean_toxicity\", \"std_toxicity\"])\n\n\nfor model_id in tqdm(MODELS_TO_TEST):\n    model = AutoModelForCausalLM.from_pretrained(model_id, device_map={\"\": device}, torch_dtype=torch.bfloat16)\n    tokenizer = AutoTokenizer.from_pretrained(model_id)\n    tokenizer.pad_token = tokenizer.eos_token\n    tokenizer.padding_side = \"left\"\n    input_texts = []\n\n    for i, example in enumerate(ds):\n        # set seed\n        torch.manual_seed(42)\n\n        input_text = example[\"comment_text\"]\n        input_texts.append(input_text[:2000])\n\n        if i > NUM_SAMPLES:\n            break\n\n        if (i + 1) % BATCH_SIZE == 0:\n            inputs = tokenizer(input_texts, return_tensors=\"pt\", padding=True).to(device)\n            inputs.input_ids = inputs.input_ids[:context_length]\n            inputs.attention_mask = inputs.attention_mask[:context_length]\n            outputs = model.generate(**inputs, do_sample=True, max_new_tokens=max_new_tokens, use_cache=True)\n            generated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n            generated_texts = [\n                generated_text.replace(input_texts[i], \"\") for i, generated_text in enumerate(generated_texts)\n            ]\n            toxicity_score = toxicity.compute(predictions=generated_texts)\n            input_texts = []\n\n            if model_id not in toxicities:\n                toxicities[model_id] = []\n            toxicities[model_id].extend(toxicity_score[\"toxicity\"])\n\n    # last batch\n    inputs = tokenizer(input_texts, return_tensors=\"pt\", padding=True).to(device)\n    outputs = model.generate(**inputs, do_sample=True, max_new_tokens=30)\n    generated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n    generated_texts = [generated_text.replace(input_texts[i], \"\") for i, generated_text in enumerate(generated_texts)]\n    toxicity_score = toxicity.compute(predictions=generated_texts)\n    toxicities[model_id].extend(toxicity_score[\"toxicity\"])\n\n    # compute mean & std using np\n    mean = np.mean(toxicities[model_id])\n    std = np.std(toxicities[model_id])\n\n    # save to file\n    writer.writerow([model_id, mean, std])\n\n    # print\n    print(f\"Model: {model_id} - Mean: {mean} - Std: {std}\")\n\n    model = None\n    if is_torch_xpu_available():\n        torch.xpu.empty_cache()\n    elif is_torch_npu_available():\n        torch.npu.empty_cache()\n    else:\n        torch.cuda.empty_cache()\n\n# close file\nfile.close()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        model_name (`str`, *optional*, defaults to `\"gpt-3.5-turbo\"`):\n            Language model to target. Possible values are:\n\n                - `\"alpaca-7b\"`\n                - `\"bard\"`\n                - `\"falcon-40b-instruct\"`\n                - `\"gpt-3.5-turbo\"` (default)\n                - `\"gpt-4\"`\n                - `\"llama-2-13b-chat\"`\n                - `\"llama-2-70b-chat\"`\n                - `\"llama-2-7b-chat\"`\n                - `\"mpt-30b-chat\"`\n                - `\"pythia-12b\"`\n                - `\"starchat\"`\n                - `\"ultralm-13b\"`\n                - `\"ultralm-65b\"`\n                - `\"vicuna-33b\"`\n                - `\"wizardlm-13b\"`\n                - `\"wizardlm-70b\"`\n                - `\"wizardlm-7b\"`\n\n        aspect (`str`, *optional*, defaults to `\"helpfulness\"`):\n            Aspect to target. Possible values are:\n\n                - `\"helpfulness\"` (default)\n                - `\"honesty\"`\n                - `\"instruction-following\"`\n                - `\"truthfulness\"`\n\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/ultrafeedback-gpt-3.5-turbo-helpfulness\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    model_name: str = \"gpt-3.5-turbo\"\n    aspect: str = \"helpfulness\"\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/ultrafeedback-gpt-3.5-turbo-helpfulness\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef to_unpaired_preference(example, model_name, aspect):\n    prompt = [{\"role\": \"user\", \"content\": example[\"instruction\"]}]\n    model_index = example[\"models\"].index(model_name)\n    response_content = example[\"completions\"][model_index][\"response\"]\n    completion = [{\"role\": \"assistant\", \"content\": response_content}]\n    score = int(example[\"completions\"][model_index][\"annotations\"][aspect][\"Rating\"])\n    label = score >= 5\n    return {\"prompt\": prompt, \"completion\": completion, \"label\": label}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\"openbmb/UltraFeedback\", split=\"train\")\n\n    dataset = dataset.filter(\n        lambda example: args.model_name in example[\"models\"], batched=False, num_proc=args.dataset_num_proc\n    )\n    dataset = dataset.map(\n        to_unpaired_preference,\n        remove_columns=[\"source\", \"instruction\", \"models\", \"completions\", \"correct_answers\", \"incorrect_answers\"],\n        fn_kwargs={\"model_name\": args.model_name, \"aspect\": args.aspect},\n        num_proc=args.dataset_num_proc,\n    )\n    dataset = dataset.train_test_split(test_size=0.05, seed=42)\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/tldr\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/tldr\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef to_prompt_completion(example):\n    tldr_format_str = \"SUBREDDIT: r/{subreddit}\\n\\nTITLE: {title}\\n\\nPOST: {post}\\n\\nTL;DR:\"\n    prompt = tldr_format_str.format(subreddit=example[\"subreddit\"], title=example[\"title\"], post=example[\"post\"])\n    completion = \" \" + example[\"summary\"]  # Add a space to separate the prompt from the completion\n    return {\"prompt\": prompt, \"completion\": completion}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    # Filtered reddit TL;DR dataset from https://github.com/openai/summarize-from-feedback?tab=readme-ov-file#reddit-tldr-dataset\n    data_files = {\n        \"train\": \"https://openaipublic.blob.core.windows.net/summarize-from-feedback/datasets/tldr_3_filtered/train.jsonl\",\n        \"validation\": \"https://openaipublic.blob.core.windows.net/summarize-from-feedback/datasets/tldr_3_filtered/valid.jsonl\",\n        \"test\": \"https://openaipublic.blob.core.windows.net/summarize-from-feedback/datasets/tldr_3_filtered/test.jsonl\",\n    }\n    dataset = load_dataset(\"json\", data_files=data_files)\n\n    dataset = dataset.map(\n        to_prompt_completion,\n        num_proc=args.dataset_num_proc,\n        remove_columns=[\"id\", \"subreddit\", \"title\", \"post\", \"summary\"],\n    )\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/ultrafeedback-prompt\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/ultrafeedback-prompt\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef to_unpaired_preference(example):\n    prompt = [{\"role\": \"user\", \"content\": example[\"instruction\"]}]\n    return {\"prompt\": prompt}\n\n\ndef drop_long_prompt(example):\n    if len(example[\"prompt\"][0][\"content\"]) > 512:\n        return False\n    else:\n        return True\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\"openbmb/UltraFeedback\", split=\"train\")\n\n    dataset = dataset.map(\n        to_unpaired_preference,\n        remove_columns=[\"source\", \"instruction\", \"models\", \"completions\", \"correct_answers\", \"incorrect_answers\"],\n        num_proc=args.dataset_num_proc,\n    )\n    dataset = dataset.filter(drop_long_prompt)\n    dataset = dataset.train_test_split(test_size=0.05, seed=42)\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, HfArgumentParser\n\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n\"\"\"\npython -i examples/datasets/tokenize_ds.py --model HuggingFaceH4/zephyr-7b-beta\npython -i examples/datasets/tokenize_ds.py --model gpt2\n\"\"\"\n\n\n@dataclass\nclass ScriptArguments:\n    dataset_name: str = field(\n        default=\"trl-internal-testing/hh-rlhf-helpful-base-trl-style\", metadata={\"help\": \"The dataset to load\"}\n    )\n    model: str = field(default=\"gpt2\", metadata={\"help\": \"The model to use for tokenization\"})\n    dataset_num_proc: Optional[int] = field(\n        default=None, metadata={\"help\": \"The number of workers to use to tokenize the data\"}\n    )\n\n\nif __name__ == \"__main__\":\n    args = HfArgumentParser(ScriptArguments).parse_args_into_dataclasses()[0]\n    dataset = load_dataset(args.dataset_name)\n    tokenizer = AutoTokenizer.from_pretrained(args.model)\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n\n    def process(row):\n        row[\"chosen\"] = tokenizer.apply_chat_template(row[\"chosen\"], tokenize=False)\n        row[\"rejected\"] = tokenizer.apply_chat_template(row[\"rejected\"], tokenize=False)\n        return row\n\n    dataset = dataset.map(process, num_proc=args.dataset_num_proc)\n    print(dataset[\"train\"][0][\"chosen\"])\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/lm-human-preferences-sentiment\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/lm-human-preferences-sentiment\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef to_prompt_completion(example, tokenizer):\n    prompt = tokenizer.decode(example[\"query\"]).strip()\n    best_idx = example[\"best\"]\n    chosen = tokenizer.decode(example[f\"sample{best_idx}\"])\n    for rejected_idx in range(4):  # take the first rejected sample that is different from the chosen one\n        rejected = tokenizer.decode(example[f\"sample{rejected_idx}\"])\n        if chosen != rejected:\n            break\n    assert chosen != rejected\n    return {\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\n        \"json\",\n        data_files=\"https://openaipublic.blob.core.windows.net/lm-human-preferences/labels/sentiment/offline_5k.json\",\n        split=\"train\",\n    )\n\n    dataset = dataset.map(\n        to_prompt_completion,\n        num_proc=args.dataset_num_proc,\n        remove_columns=[\"query\", \"sample0\", \"sample1\", \"sample2\", \"sample3\", \"best\"],\n        fn_kwargs={\"tokenizer\": AutoTokenizer.from_pretrained(\"gpt2\")},\n    )\n\n    # train_size taken from https://github.com/openai/lm-human-preferences/blob/cbfd210bb8b08f6bc5c26878c10984b90f516c66/launch.py#L70)\n    dataset = dataset.train_test_split(train_size=4992)\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\nfrom dataclasses import dataclass\nfrom typing import Dict, List, Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/hh-rlhf-helpful-base\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/hh-rlhf-helpful-base\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef common_start(str1: str, str2: str) -> str:\n    # Zip the two strings and iterate over them together\n    common_chars = []\n    for c1, c2 in zip(str1, str2):\n        if c1 == c2:\n            common_chars.append(c1)\n        else:\n            break\n    # Join the common characters and return as a string\n    return \"\".join(common_chars)\n\n\ndef extract_dialogue(example: str) -> List[Dict[str, str]]:\n    # Extract the prompt, which corresponds to the common start of the chosen and rejected dialogues\n    prompt_text = common_start(example[\"chosen\"], example[\"rejected\"])\n\n    # The chosen and rejected may share a common start, so we need to remove the common part\n    if not prompt_text.endswith(\"\\n\\nAssistant: \"):\n        prompt_text = prompt_text[: prompt_text.rfind(\"\\n\\nAssistant: \")] + \"\\n\\nAssistant: \"\n\n    # Extract the chosen and rejected lines\n    chosen_line = example[\"chosen\"][len(prompt_text) :]\n    rejected_line = example[\"rejected\"][len(prompt_text) :]\n\n    # Remove the generation prompt (\"\\n\\nAssistant: \") from the prompt\n    prompt_text = prompt_text[: -len(\"\\n\\nAssistant: \")]\n\n    # Split the string at every occurrence of \"Human: \" or \"Assistant: \"\n    prompt_lines = re.split(r\"(\\n\\nAssistant: |\\n\\nHuman: )\", prompt_text)\n\n    # Remove the first element as it's empty\n    prompt_lines = prompt_lines[1:]\n\n    prompt = []\n    for idx in range(0, len(prompt_lines), 2):\n        role = \"user\" if prompt_lines[idx] == \"\\n\\nHuman: \" else \"assistant\"\n        content = prompt_lines[idx + 1]\n        prompt.append({\"role\": role, \"content\": content})\n\n    # Remove the prompt from the chosen and rejected dialogues\n    chosen = [{\"role\": \"assitant\", \"content\": chosen_line}]\n    rejected = [{\"role\": \"assistant\", \"content\": rejected_line}]\n\n    return {\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\"Anthropic/hh-rlhf\", data_dir=\"helpful-base\")\n    dataset = dataset.map(extract_dialogue, num_proc=args.dataset_num_proc)\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\n\nfrom datasets import Dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        test_size (`float`, *optional*, defaults to `0.1`):\n            Fraction of the dataset to include in the test split.\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/zen\"`):\n            Hugging Face repository ID to push the dataset to.\n    \"\"\"\n\n    test_size: float = 0.1\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/zen\"\n\n\ndef main(test_size, push_to_hub, repo_id):\n    # fmt: off\n    standard_language_modeling_dataset = Dataset.from_dict({\n        \"text\": [\n            \"Beautiful is better than ugly.\",\n            \"Explicit is better than implicit.\",\n            \"Simple is better than complex.\",\n            \"Complex is better than complicated.\",\n            \"Flat is better than nested.\",\n            \"Sparse is better than dense.\",\n            \"Readability counts.\",\n            \"Special cases aren't special enough to break the rules.\",\n            \"Although practicality beats purity.\",\n            \"Errors should never pass silently.\",\n            \"Unless explicitly silenced.\",\n            \"In the face of ambiguity, refuse the temptation to guess.\",\n            \"There should be one-- and preferably only one --obvious way to do it.\",\n            \"Although that way may not be obvious at first unless you're Dutch.\",\n            \"Now is better than never.\",\n            \"Although never is often better than *right* now.\",\n            \"If the implementation is hard to explain, it's a bad idea.\",\n            \"If the implementation is easy to explain, it may be a good idea.\",\n            \"Namespaces are one honking great idea -- let's do more of those!\",\n        ],\n    })\n    standard_language_modeling_dataset = standard_language_modeling_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_language_modeling_dataset.push_to_hub(repo_id, config_name=\"standard_language_modeling\")\n\n    standard_prompt_only_dataset = Dataset.from_dict({\n        \"prompt\": [\n            \"Beautiful is better than\",\n            \"Explicit is\",\n            \"Simple is better\",\n            \"Complex\",\n            \"Flat is better than\",\n            \"Sparse is better\",\n            \"Readability\",\n            \"Special cases aren't special\",\n            \"Although practicality beats\",\n            \"Errors should never\",\n            \"Unless explicitly\",\n            \"In the face of ambiguity, refuse\",\n            \"There should be one-- and preferably\",\n            \"Although that way may not be obvious at first unless you're\",\n            \"Now is\",\n            \"Although never is often\",\n            \"If the implementation is hard to explain,\",\n            \"If the implementation is easy\",\n            \"Namespaces are one honking great\",\n        ],\n    })\n    standard_prompt_only_dataset = standard_prompt_only_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_prompt_only_dataset.push_to_hub(repo_id, config_name=\"standard_prompt_only\")\n\n    standard_prompt_completion_dataset = Dataset.from_dict({\n        \"prompt\": [\n            \"Beautiful is better than\",\n            \"Explicit is\",\n            \"Simple is better\",\n            \"Complex\",\n            \"Flat is better than\",\n            \"Sparse is better\",\n            \"Readability\",\n            \"Special cases aren't special\",\n            \"Although practicality beats\",\n            \"Errors should never\",\n            \"Unless explicitly\",\n            \"In the face of ambiguity, refuse\",\n            \"There should be one-- and preferably\",\n            \"Although that way may not be obvious at first unless you're\",\n            \"Now is\",\n            \"Although never is often\",\n            \"If the implementation is hard to explain,\",\n            \"If the implementation is easy\",\n            \"Namespaces are one honking great\",\n        ],\n        \"completion\": [\n            \" ugly.\",\n            \" better than implicit.\",\n            \" than complex.\",\n            \" is better than complicated.\",\n            \" nested.\",\n            \" than dense.\",\n            \" counts.\",\n            \" enough to break the rules.\",\n            \" purity.\",\n            \" pass silently.\",\n            \" silenced.\",\n            \" the temptation to guess.\",\n            \" only one --obvious way to do it.\",\n            \" Dutch.\",\n            \" better than never.\",\n            \" better than *right* now.\",\n            \" it's a bad idea.\",\n            \" to explain, it may be a good idea.\",\n            \" idea -- let's do more of those!\",\n        ],\n    })\n    standard_prompt_completion_dataset = standard_prompt_completion_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_prompt_completion_dataset.push_to_hub(repo_id, config_name=\"standard_prompt_completion\")\n\n    standard_preference_dataset = Dataset.from_dict({\n        \"prompt\": [\n            \"Beautiful is better than\",\n            \"Explicit is\",\n            \"Simple is better\",\n            \"Complex\",\n            \"Flat is better than\",\n            \"Sparse is better\",\n            \"Readability\",\n            \"Special cases aren't special\",\n            \"Although practicality beats\",\n            \"Errors should never\",\n            \"Unless explicitly\",\n            \"In the face of ambiguity, refuse\",\n            \"There should be one-- and preferably\",\n            \"Although that way may not be obvious at first unless you're\",\n            \"Now is\",\n            \"Although never is often\",\n            \"If the implementation is hard to explain,\",\n            \"If the implementation is easy\",\n            \"Namespaces are one honking great\",\n        ],\n        \"chosen\": [\n            \" ugly.\",\n            \" better than implicit.\",\n            \" than complex.\",\n            \" is better than complicated.\",\n            \" nested.\",\n            \" than dense.\",\n            \" counts.\",\n            \" enough to break the rules.\",\n            \" purity.\",\n            \" pass silently.\",\n            \" silenced.\",\n            \" the temptation to guess.\",\n            \" only one --obvious way to do it.\",\n            \" Dutch.\",\n            \" better than never.\",\n            \" better than *right* now.\",\n            \" it's a bad idea.\",\n            \" to explain, it may be a good idea.\",\n            \" idea -- let's do more of those!\",\n        ],\n        \"rejected\": [\n            \" the moon.\",\n            \" worse than nothing.\",\n            \" than a long vacation.\",\n            \" is always the answer.\",\n            \" chocolate.\",\n            \" without any context.\",\n            \" is optional.\",\n            \" enough to become unicorns.\",\n            \" reality.\",\n            \" pass their driving test.\",\n            \" forgotten.\",\n            \" the opportunity to laugh.\",\n            \" two or more confusing methods.\",\n            \" a time traveler.\",\n            \" never better.\",\n            \" not even a possibility.\",\n            \" it's clearly the best choice.\",\n            \" it's probably magic.\",\n            \" watermelon -- let's plant some!\",\n        ],\n    })\n    standard_preference_dataset = standard_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_preference_dataset.push_to_hub(repo_id, config_name=\"standard_preference\")\n\n    standard_implicit_prompt_preference_dataset = Dataset.from_dict({\n        \"chosen\": [\n            \"Beautiful is better than ugly.\",\n            \"Explicit is better than implicit.\",\n            \"Simple is better than complex.\",\n            \"Complex is better than complicated.\",\n            \"Flat is better than nested.\",\n            \"Sparse is better than dense.\",\n            \"Readability counts.\",\n            \"Special cases aren't special enough to break the rules.\",\n            \"Although practicality beats purity.\",\n            \"Errors should never pass silently.\",\n            \"Unless explicitly silenced.\",\n            \"In the face of ambiguity, refuse the temptation to guess.\",\n            \"There should be one-- and preferably only one --obvious way to do it.\",\n            \"Although that way may not be obvious at first unless you're Dutch.\",\n            \"Now is better than never.\",\n            \"Although never is often better than *right* now.\",\n            \"If the implementation is hard to explain, it's a bad idea.\",\n            \"If the implementation is easy to explain, it may be a good idea.\",\n            \"Namespaces are one honking great idea -- let's do more of those!\",\n        ],\n        \"rejected\": [\n            \"Beautiful is better than the moon.\",\n            \"Explicit is worse than nothing.\",\n            \"Simple is better than a long vacation.\",\n            \"Complex is always the answer.\",\n            \"Flat is better than chocolate.\",\n            \"Sparse is better without any context.\",\n            \"Readability is optional.\",\n            \"Special cases aren't special enough to become unicorns.\",\n            \"Although practicality beats reality.\",\n            \"Errors should never pass their driving test.\",\n            \"Unless explicitly forgotten.\",\n            \"In the face of ambiguity, refuse the opportunity to laugh.\",\n            \"There should be one-- and preferably two or more confusing methods.\",\n            \"Although that way may not be obvious at first unless you're a time traveler.\",\n            \"Now is never better.\",\n            \"Although never is often not even a possibility.\",\n            \"If the implementation is hard to explain, it's clearly the best choice.\",\n            \"If the implementation is easy it's probably magic.\",\n            \"Namespaces are one honking great watermelon -- let's plant some!\",\n        ],\n    })\n    standard_implicit_prompt_preference_dataset = standard_implicit_prompt_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_implicit_prompt_preference_dataset.push_to_hub(repo_id, config_name=\"standard_implicit_prompt_preference\")\n\n    standard_unpaired_preference_dataset = Dataset.from_dict({\n        \"prompt\": [\n            \"Beautiful is better than\",\n            \"Explicit is\",\n            \"Simple is better\",\n            \"Complex\",\n            \"Flat is better than\",\n            \"Sparse is better\",\n            \"Readability\",\n            \"Special cases aren't special\",\n            \"Although practicality beats\",\n            \"Errors should never\",\n            \"Unless explicitly\",\n            \"In the face of ambiguity, refuse\",\n            \"There should be one-- and preferably\",\n            \"Although that way may not be obvious at first unless you're\",\n            \"Now is\",\n            \"Although never is often\",\n            \"If the implementation is hard to explain,\",\n            \"If the implementation is easy\",\n            \"Namespaces are one honking great\",\n        ],\n        \"completion\": [\n            \" ugly.\",\n            \" worse than nothing.\",\n            \" than a long vacation.\",\n            \" is better than complicated.\",\n            \" nested.\",\n            \" without any context.\",\n            \" counts.\",\n            \" enough to become unicorns.\",\n            \" purity.\",\n            \" pass silently.\",\n            \" forgotten.\",\n            \" the temptation to guess.\",\n            \" only one --obvious way to do it.\",\n            \" a time traveler.\",\n            \" better than never.\",\n            \" not even a possibility.\",\n            \" it's a bad idea.\",\n            \" it's probably magic.\",\n            \" watermelon -- let's plant some!\",\n        ],\n        \"label\": [True, False, False, True, True, False, True, False, True, True, False, True, True, False, True, False, True, False, False],\n    })\n    standard_unpaired_preference_dataset = standard_unpaired_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        standard_unpaired_preference_dataset.push_to_hub(repo_id, config_name=\"standard_unpaired_preference\")\n\n    conversational_language_modeling_dataset = Dataset.from_dict({\n        \"messages\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}, {\"role\": \"assistant\", \"content\": \"Beautiful.\"},],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}, {\"role\": \"assistant\", \"content\": \"Explicit.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}, {\"role\": \"assistant\", \"content\": \"Simple.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}, {\"role\": \"assistant\", \"content\": \"Complex.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}, {\"role\": \"assistant\", \"content\": \"Flat.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}, {\"role\": \"assistant\", \"content\": \"Sparse.\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}, {\"role\": \"assistant\", \"content\": \"Readability.\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}, {\"role\": \"assistant\", \"content\": \"No, special cases aren't special enough to break the rules.\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}, {\"role\": \"assistant\", \"content\": \"Practicality.\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}, {\"role\": \"assistant\", \"content\": \"Errors.\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}, {\"role\": \"assistant\", \"content\": \"When explicitly silenced.\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}, {\"role\": \"assistant\", \"content\": \"Refuse the temptation to guess.\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}, {\"role\": \"assistant\", \"content\": \"One, and preferably only one.\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}, {\"role\": \"assistant\", \"content\": \"Dutch.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}, {\"role\": \"assistant\", \"content\": \"Now is better than never.\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}, {\"role\": \"assistant\", \"content\": \"Yes, often.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it may be a good idea.\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}, {\"role\": \"assistant\", \"content\": \"Namespaces are one honking great idea.\"}],\n        ],\n    })\n    conversational_language_modeling_dataset = conversational_language_modeling_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_language_modeling_dataset.push_to_hub(repo_id, config_name=\"conversational_language_modeling\")\n\n    conversational_prompt_only_dataset = Dataset.from_dict({\n        \"prompt\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}],\n        ],\n    })\n    conversational_prompt_only_dataset = conversational_prompt_only_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_prompt_only_dataset.push_to_hub(repo_id, config_name=\"conversational_prompt_only\")\n\n    conversational_prompt_completion_dataset = Dataset.from_dict({\n        \"prompt\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}],\n        ],\n        \"completion\": [\n            [{\"role\": \"assistant\", \"content\": \"Beautiful.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Explicit.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Simple.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Complex.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Flat.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Sparse.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Readability.\"}],\n            [{\"role\": \"assistant\", \"content\": \"No, special cases aren't special enough to break the rules.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Practicality.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Errors.\"}],\n            [{\"role\": \"assistant\", \"content\": \"When explicitly silenced.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Refuse the temptation to guess.\"}],\n            [{\"role\": \"assistant\", \"content\": \"One, and preferably only one.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Dutch.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Now is better than never.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Yes, often.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it may be a good idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Namespaces are one honking great idea.\"}],\n        ],\n    })\n    conversational_prompt_completion_dataset = conversational_prompt_completion_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_prompt_completion_dataset.push_to_hub(repo_id, config_name=\"conversational_prompt_completion\")\n\n    conversational_preference_dataset = Dataset.from_dict({\n        \"prompt\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}],\n        ],\n        \"chosen\": [\n            [{\"role\": \"assistant\", \"content\": \"Beautiful.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Explicit.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Simple.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Complex.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Flat.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Sparse.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Readability.\"}],\n            [{\"role\": \"assistant\", \"content\": \"No, special cases aren't special enough to break the rules.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Practicality.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Errors.\"}],\n            [{\"role\": \"assistant\", \"content\": \"When explicitly silenced.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Refuse the temptation to guess.\"}],\n            [{\"role\": \"assistant\", \"content\": \"One, and preferably only one.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Dutch.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Now is better than never.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Yes, often.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it may be a good idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Namespaces are one honking great idea.\"}],\n        ],\n        \"rejected\": [\n            [{\"role\": \"assistant\", \"content\": \"Acceptable.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Explained.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Very complex.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Very complicated.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Circular.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Heavy.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Looking complicated.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Yes, special cases are special enough to break the rules.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Nothing.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Warnings.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Never.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Give up.\"}],\n            [{\"role\": \"assistant\", \"content\": \"As many as possible.\"}],\n            [{\"role\": \"assistant\", \"content\": \"French.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Some day.\"}],\n            [{\"role\": \"assistant\", \"content\": \"No, never.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it's a good idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"assistant\", \"content\": \"Recursion.\"}],\n        ],\n    })\n    conversational_preference_dataset = conversational_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_preference_dataset.push_to_hub(repo_id, config_name=\"conversational_preference\")\n\n    conversational_implicit_prompt_preference_dataset = Dataset.from_dict({\n        \"chosen\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}, {\"role\": \"assistant\", \"content\": \"Beautiful.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}, {\"role\": \"assistant\", \"content\": \"Explicit.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}, {\"role\": \"assistant\", \"content\": \"Simple.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}, {\"role\": \"assistant\", \"content\": \"Complex.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}, {\"role\": \"assistant\", \"content\": \"Flat.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}, {\"role\": \"assistant\", \"content\": \"Sparse.\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}, {\"role\": \"assistant\", \"content\": \"Readability.\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}, {\"role\": \"assistant\", \"content\": \"No, special cases aren't special enough to break the rules.\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}, {\"role\": \"assistant\", \"content\": \"Practicality.\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}, {\"role\": \"assistant\", \"content\": \"Errors.\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}, {\"role\": \"assistant\", \"content\": \"When explicitly silenced.\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}, {\"role\": \"assistant\", \"content\": \"Refuse the temptation to guess.\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}, {\"role\": \"assistant\", \"content\": \"One, and preferably only one.\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}, {\"role\": \"assistant\", \"content\": \"Dutch.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}, {\"role\": \"assistant\", \"content\": \"Now is better than never.\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}, {\"role\": \"assistant\", \"content\": \"Yes, often.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it may be a good idea.\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}, {\"role\": \"assistant\", \"content\": \"Namespaces are one honking great idea.\"}],\n        ],\n        \"rejected\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}, {\"role\": \"assistant\", \"content\": \"Acceptable.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}, {\"role\": \"assistant\", \"content\": \"Explained.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}, {\"role\": \"assistant\", \"content\": \"Very complex.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}, {\"role\": \"assistant\", \"content\": \"Very complicated.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}, {\"role\": \"assistant\", \"content\": \"Circular.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}, {\"role\": \"assistant\", \"content\": \"Heavy.\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}, {\"role\": \"assistant\", \"content\": \"Looking complicated.\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}, {\"role\": \"assistant\", \"content\": \"Yes, special cases are special enough to break the rules.\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}, {\"role\": \"assistant\", \"content\": \"Nothing.\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}, {\"role\": \"assistant\", \"content\": \"Warnings.\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}, {\"role\": \"assistant\", \"content\": \"Never.\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}, {\"role\": \"assistant\", \"content\": \"Give up.\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}, {\"role\": \"assistant\", \"content\": \"As many as possible.\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}, {\"role\": \"assistant\", \"content\": \"French.\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}, {\"role\": \"assistant\", \"content\": \"Some day.\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}, {\"role\": \"assistant\", \"content\": \"No, never.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it's a good idea.\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}, {\"role\": \"assistant\", \"content\": \"It means it's a bad idea.\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}, {\"role\": \"assistant\", \"content\": \"Recursion.\"}],\n        ],\n    })\n    conversational_implicit_prompt_preference_dataset = conversational_implicit_prompt_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_implicit_prompt_preference_dataset.push_to_hub(repo_id, config_name=\"conversational_implicit_prompt_preference\")\n\n    conversational_unpaired_preference_dataset = Dataset.from_dict({\n        \"prompt\": [\n            [{\"role\": \"user\", \"content\": \"What is better than ugly?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than implicit?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complex?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than complicated?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than nested?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than dense?\"}],\n            [{\"role\": \"user\", \"content\": \"What counts?\"}],\n            [{\"role\": \"user\", \"content\": \"Are special cases enough to break the rules?\"}],\n            [{\"role\": \"user\", \"content\": \"What beats purity?\"}],\n            [{\"role\": \"user\", \"content\": \"What should never pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"When can errors pass silently?\"}],\n            [{\"role\": \"user\", \"content\": \"What should you do in the face of ambiguity?\"}],\n            [{\"role\": \"user\", \"content\": \"How many ways should there be to do it?\"}],\n            [{\"role\": \"user\", \"content\": \"For whom may the way not be obvious at first?\"}],\n            [{\"role\": \"user\", \"content\": \"What is better than never?\"}],\n            [{\"role\": \"user\", \"content\": \"Is never better than *right* now?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is hard to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"What does it mean if the implementation is easy to explain?\"}],\n            [{\"role\": \"user\", \"content\": \"Any great ideas?\"}],\n        ],\n        \"completion\": [\n            [{'role': 'assistant', 'content': 'Beautiful.'}],\n            [{'role': 'assistant', 'content': 'Explicit.'}],\n            [{'role': 'assistant', 'content': 'Simple.'}],\n            [{'role': 'assistant', 'content': 'Very complicated.'}],\n            [{'role': 'assistant', 'content': 'Flat.'}],\n            [{'role': 'assistant', 'content': 'Sparse.'}],\n            [{'role': 'assistant', 'content': 'Readability.'}],\n            [{'role': 'assistant', 'content': 'Yes, special cases are special enough to break the rules.'}],\n            [{'role': 'assistant', 'content': 'Practicality.'}],\n            [{'role': 'assistant', 'content': 'Warnings.'}],\n            [{'role': 'assistant', 'content': 'When explicitly silenced.'}],\n            [{'role': 'assistant', 'content': 'Give up.'}],\n            [{'role': 'assistant', 'content': 'One, and preferably only one.'}],\n            [{'role': 'assistant', 'content': 'French.'}],\n            [{'role': 'assistant', 'content': 'Some day.'}],\n            [{'role': 'assistant', 'content': 'Yes, often.'}],\n            [{'role': 'assistant', 'content': \"It means it's a bad idea.\"}],\n            [{'role': 'assistant', 'content': 'It means it may be a good idea.'}],\n            [{'role': 'assistant', 'content': 'Namespaces are one honking great idea.'}],\n        ],\n        \"label\": [True, True, True, False, True, True, True, False, True, False, True, False, True, False, False, True, True, True, True],\n    })\n    conversational_unpaired_preference_dataset = conversational_unpaired_preference_dataset.train_test_split(test_size=test_size)\n    if push_to_hub:\n        conversational_unpaired_preference_dataset.push_to_hub(repo_id, config_name=\"conversational_unpaired_preference\")\n    # fmt: on\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n    main(args.test_size, args.push_to_hub, args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/lm-human-preferences-descriptiveness\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/lm-human-preferences-descriptiveness\"\n    dataset_num_proc: Optional[int] = None\n\n\n# Edge cases handling: remove the cases where all samples are the same\ndef samples_not_all_same(example):\n    return not all(example[\"sample0\"] == example[f\"sample{j}\"] for j in range(1, 4))\n\n\ndef to_prompt_completion(example, tokenizer):\n    prompt = tokenizer.decode(example[\"query\"]).strip()\n    best_idx = example[\"best\"]\n    chosen = tokenizer.decode(example[f\"sample{best_idx}\"])\n    for rejected_idx in range(4):  # take the first rejected sample that is different from the chosen one\n        rejected = tokenizer.decode(example[f\"sample{rejected_idx}\"])\n        if chosen != rejected:\n            break\n    assert chosen != rejected\n    return {\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\n        \"json\",\n        data_files=\"https://openaipublic.blob.core.windows.net/lm-human-preferences/labels/descriptiveness/offline_5k.json\",\n        split=\"train\",\n    )\n\n    dataset = dataset.filter(samples_not_all_same, num_proc=args.dataset_num_proc)\n\n    dataset = dataset.map(\n        to_prompt_completion,\n        num_proc=args.dataset_num_proc,\n        remove_columns=[\"query\", \"sample0\", \"sample1\", \"sample2\", \"sample3\", \"best\"],\n        fn_kwargs={\"tokenizer\": AutoTokenizer.from_pretrained(\"gpt2\")},\n    )\n\n    # train_size taken from https://github.com/openai/lm-human-preferences/blob/cbfd210bb8b08f6bc5c26878c10984b90f516c66/launch.py#L79)\n    dataset = dataset.train_test_split(train_size=4992)\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\n\n\n@dataclass\nclass ScriptArguments:\n    r\"\"\"\n    Arguments for the script.\n\n    Args:\n        push_to_hub (`bool`, *optional*, defaults to `False`):\n            Whether to push the dataset to the Hugging Face Hub.\n        repo_id (`str`, *optional*, defaults to `\"trl-lib/tldr-preference\"`):\n            Hugging Face repository ID to push the dataset to.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of workers to use for dataset processing.\n    \"\"\"\n\n    push_to_hub: bool = False\n    repo_id: str = \"trl-lib/tldr-preference\"\n    dataset_num_proc: Optional[int] = None\n\n\ndef to_preference(example):\n    info = example[\"info\"]\n    if example[\"batch\"] in [\"batch0_cnndm\", \"cnndm0\", \"cnndm2\"]:  # CNN Daily Mail batches\n        article = info[\"article\"].replace(\"\\n\\n\", \"\\n\")\n        prompt = f\"TITLE: {info['title']}\\n\\n{article}\\n\\nTL;DR:\"\n    elif example[\"batch\"] in [f\"batch{i}\" for i in range(3, 23)] + [\"edit_b2_eval_test\"]:  # Reddit batches\n        post = info[\"post\"].replace(\"\\n\\n\", \"\\n\")\n        prompt = f\"SUBREDDIT: r/{info['subreddit']}\\n\\nTITLE: {info['title']}\\n\\nPOST: {post}\\n\\nTL;DR:\"\n    else:\n        raise ValueError(f\"Unknown batch: {example['batch']}\")\n\n    chosen_idx = example[\"choice\"]\n    rejected_idx = 1 - chosen_idx\n    chosen = example[\"summaries\"][chosen_idx][\"text\"]\n    rejected = example[\"summaries\"][rejected_idx][\"text\"]\n    return {\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected}\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser(ScriptArguments)\n    args = parser.parse_args_into_dataclasses()[0]\n\n    dataset = load_dataset(\"openai/summarize_from_feedback\", \"comparisons\")\n\n    dataset = dataset.map(\n        to_preference,\n        num_proc=args.dataset_num_proc,\n        remove_columns=[\"info\", \"summaries\", \"choice\", \"worker\", \"batch\", \"split\", \"extra\"],\n    )\n\n    if args.push_to_hub:\n        dataset.push_to_hub(args.repo_id)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nRun the ORPO training script with the following command with some example arguments.\nIn general, the optimal configuration for ORPO will be similar to that of DPO without the need for a reference model:\n\n# regular:\npython examples/scripts/orpo.py \\\n    --model_name_or_path=gpt2 \\\n    --per_device_train_batch_size 4 \\\n    --max_steps 1000 \\\n    --learning_rate 8e-6 \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=\"gpt2-aligned-orpo\" \\\n    --warmup_steps 150 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step \\\n    --no_remove_unused_columns\n\n# peft:\npython examples/scripts/orpo.py \\\n    --model_name_or_path=gpt2 \\\n    --per_device_train_batch_size 4 \\\n    --max_steps 1000 \\\n    --learning_rate 8e-5 \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=\"gpt2-lora-aligned-orpo\" \\\n    --optim rmsprop \\\n    --warmup_steps 150 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step \\\n    --no_remove_unused_columns \\\n    --use_peft \\\n    --lora_r=16 \\\n    --lora_alpha=16\n\"\"\"\n\nfrom dataclasses import dataclass, field\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser\n\nfrom trl import ModelConfig, ORPOConfig, ORPOTrainer, get_peft_config\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n@dataclass\nclass ScriptArguments:\n    dataset_name: str = field(\n        default=\"trl-internal-testing/hh-rlhf-helpful-base-trl-style\",\n        metadata={\"help\": \"The name of the dataset to use.\"},\n    )\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, ORPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_into_dataclasses()\n\n    ################\n    # Model & Tokenizer\n    ################\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n\n    def process(row):\n        row[\"prompt\"] = tokenizer.apply_chat_template(row[\"chosen\"][:-1], tokenize=False)\n        row[\"chosen\"] = tokenizer.apply_chat_template([row[\"chosen\"][-1]], tokenize=False)\n        row[\"rejected\"] = tokenizer.apply_chat_template([row[\"rejected\"][-1]], tokenize=False)\n        return row\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(process, num_prc=training_args.dataset_num_proc)\n\n    ################\n    # Training\n    ################\n    trainer = ORPOTrainer(\n        model,\n        args=training_args,\n        train_dataset=dataset[\"train\"],\n        eval_dataset=dataset[\"test\"],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n\n    # train and save the model\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFull training:\npython examples/scripts/reward_modeling.py \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --dataset_name trl-lib/ultrafeedback_binarized \\\n    --output_dir Qwen2-0.5B-Reward \\\n    --per_device_train_batch_size 8 \\\n    --num_train_epochs 1 \\\n    --gradient_accumulation_steps 1 \\\n    --remove_unused_columns False \\\n    --gradient_checkpointing True \\\n    --learning_rate 1.0e-5 \\\n    --logging_steps 25 \\\n    --eval_strategy steps \\\n    --eval_steps 50 \\\n    --max_length 2048\n\nLoRA:\npython examples/scripts/reward_modeling.py \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --dataset_name trl-lib/ultrafeedback_binarized \\\n    --output_dir Qwen2-0.5B-Reward \\\n    --per_device_train_batch_size 8 \\\n    --num_train_epochs 1 \\\n    --gradient_accumulation_steps 1 \\\n    --remove_unused_columns False \\\n    --gradient_checkpointing True \\\n    --learning_rate 1.0e-5 \\\n    --logging_steps 25 \\\n    --eval_strategy steps \\\n    --eval_steps 50 \\\n    --max_length 2048 /\n    --use_peft \\\n    --lora_r 32 \\\n    --lora_alpha 16\n\"\"\"\n\nimport warnings\n\nimport torch\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom tqdm import tqdm\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser\n\nfrom trl import (\n    ModelConfig,\n    RewardConfig,\n    RewardTrainer,\n    get_kbit_device_map,\n    get_peft_config,\n    get_quantization_config,\n    setup_chat_format,\n)\nfrom trl.commands.cli_utils import RewardScriptArguments\nfrom trl.extras.dataset_formatting import conversations_formatting_function\n\n\ntqdm.pandas()\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((RewardScriptArguments, RewardConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_into_dataclasses()\n    training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)\n\n    ################\n    # Model & Tokenizer\n    ################\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True\n    )\n    model = AutoModelForSequenceClassification.from_pretrained(\n        model_config.model_name_or_path, num_labels=1, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    # Align padding tokens between tokenizer and model\n    model.config.pad_token_id = tokenizer.pad_token_id\n\n    # If post-training a base model, use ChatML as the default template\n    if tokenizer.chat_template is None:\n        model, tokenizer = setup_chat_format(model, tokenizer)\n\n    if model_config.use_peft and model_config.lora_task_type != \"SEQ_CLS\":\n        warnings.warn(\n            \"You are using a `task_type` that is different than `SEQ_CLS` for PEFT. This will lead to silent bugs\"\n            \" Make sure to pass --lora_task_type SEQ_CLS when using this script with PEFT.\"\n        )\n\n    #############################\n    # Load and preprocess dataset\n    #############################\n    dataset = load_dataset(args.dataset_name)\n\n    def preprocess_function(examples):\n        new_examples = {\n            \"input_ids_chosen\": [],\n            \"attention_mask_chosen\": [],\n            \"input_ids_rejected\": [],\n            \"attention_mask_rejected\": [],\n        }\n        for chosen, rejected in zip(examples[\"chosen\"], examples[\"rejected\"]):\n            tokenized_chosen = tokenizer(chosen)\n            tokenized_rejected = tokenizer(rejected)\n            new_examples[\"input_ids_chosen\"].append(tokenized_chosen[\"input_ids\"])\n            new_examples[\"attention_mask_chosen\"].append(tokenized_chosen[\"attention_mask\"])\n            new_examples[\"input_ids_rejected\"].append(tokenized_rejected[\"input_ids\"])\n            new_examples[\"attention_mask_rejected\"].append(tokenized_rejected[\"attention_mask\"])\n\n        return new_examples\n\n    with PartialState().local_main_process_first():\n        # Wrap inputs with chat template.\n        # This assumes the chosen/rejected columns are in the OpenAI messages format.\n        chosen_fn = conversations_formatting_function(tokenizer, \"chosen\")\n        rejected_fn = conversations_formatting_function(tokenizer, \"rejected\")\n        dataset = dataset.map(\n            lambda x: {\"chosen\": chosen_fn(x), \"rejected\": rejected_fn(x)}, num_proc=training_args.dataset_num_proc\n        )\n        # Tokenize inputs\n        dataset = dataset.map(\n            preprocess_function,\n            batched=True,\n            num_proc=training_args.dataset_num_proc,\n        )\n        # Filter out examples that are too long\n        dataset = dataset.filter(\n            lambda x: len(x[\"input_ids_chosen\"]) <= training_args.max_length\n            and len(x[\"input_ids_rejected\"]) <= training_args.max_length,\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    ##########\n    # Training\n    ##########\n    trainer = RewardTrainer(\n        model=model,\n        tokenizer=tokenizer,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        peft_config=get_peft_config(model_config),\n    )\n    trainer.train()\n\n    ############################\n    # Save model and push to Hub\n    ############################\n    trainer.save_model(training_args.output_dir)\n    metrics = trainer.evaluate()\n    trainer.log_metrics(\"eval\", metrics)\n    trainer.save_metrics(\"eval\", metrics)\n    trainer.save_model(training_args.output_dir)\n    trainer.push_to_hub()\n\n\n# flake8: noqa\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUsage:\n\npython examples/scripts/nash_md.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-tldr-nash-md \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 32 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 64 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --push_to_hub\n\n\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/nash_md.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-tldr-nash-md \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 32 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 64 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --push_to_hub\n\"\"\"\n\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, GenerationConfig\nfrom trl import (\n    DPOScriptArguments,\n    ModelConfig,\n    NashMDConfig,\n    NashMDTrainer,\n    get_kbit_device_map,\n    get_quantization_config,\n    LogCompletionsCallback,\n)\nfrom trl.commands.cli_utils import TrlParser\nfrom trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((DPOScriptArguments, NashMDConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n    args.gradient_checkpointing_kwargs = {\"use_reentrant\": True}\n\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    ref_model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, num_labels=1, trust_remote_code=model_config.trust_remote_code\n    )\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE\n\n    dataset = load_dataset(args.dataset_name)\n\n    trainer = NashMDTrainer(\n        model=model,\n        ref_model=ref_model,\n        reward_model=reward_model,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n    )\n    generation_config = GenerationConfig(\n        max_new_tokens=training_args.max_new_tokens, do_sample=True, temperature=training_args.temperature\n    )\n    completions_callback = LogCompletionsCallback(trainer, generation_config, num_prompts=8)\n    trainer.add_callback(completions_callback)\n    # train the model\n    trainer.train()\n\n    # save the model\n    trainer.save_model(training_args.output_dir)\n\n\n# flake8: noqa\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUsage:\n\npython examples/scripts/xpo.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-tldr-xpo \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 32 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 64 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --push_to_hub\n\"\"\"\n\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, GenerationConfig\nfrom trl import (\n    DPOScriptArguments,\n    ModelConfig,\n    XPOConfig,\n    XPOTrainer,\n    get_kbit_device_map,\n    get_quantization_config,\n    LogCompletionsCallback,\n)\nfrom trl.commands.cli_utils import TrlParser\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((DPOScriptArguments, XPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n    args.gradient_checkpointing_kwargs = {\"use_reentrant\": True}\n\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    ref_model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, num_labels=1, trust_remote_code=model_config.trust_remote_code\n    )\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n\n    dataset = load_dataset(args.dataset_name)\n\n    trainer = XPOTrainer(\n        model=model,\n        ref_model=ref_model,\n        reward_model=reward_model,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n    )\n    generation_config = GenerationConfig(\n        max_new_tokens=training_args.max_new_tokens, do_sample=True, temperature=training_args.temperature\n    )\n    completions_callback = LogCompletionsCallback(trainer, generation_config, num_prompts=8)\n    trainer.add_callback(completions_callback)\n    # train the model\n    trainer.train()\n\n    # save the model\n    trainer.save_model(training_args.output_dir)\n\n\n# flake8: noqa\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\n# Full training\npython examples/scripts/dpo.py \\\n    --dataset_name trl-lib/ultrafeedback_binarized \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --learning_rate 5.0e-7 \\\n    --num_train_epochs 1 \\\n    --per_device_train_batch_size 2 \\\n    --gradient_accumulation_steps 8 \\\n    --gradient_checkpointing \\\n    --logging_steps 25 \\\n    --eval_strategy steps \\\n    --eval_steps 50 \\\n    --output_dir Qwen2-0.5B-DPO \\\n    --no_remove_unused_columns\n\n# LoRA:\npython examples/scripts/dpo.py \\\n    --dataset_name trl-lib/ultrafeedback_binarized \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --learning_rate 5.0e-6 \\\n    --num_train_epochs 1 \\\n    --per_device_train_batch_size 2 \\\n    --gradient_accumulation_steps 8 \\\n    --gradient_checkpointing \\\n    --logging_steps 25 \\\n    --eval_strategy steps \\\n    --eval_steps 50 \\\n    --output_dir Qwen2-0.5B-DPO \\\n    --no_remove_unused_columns \\\n    --use_peft \\\n    --lora_r 32 \\\n    --lora_alpha 16\n\"\"\"\n\nfrom trl.commands.cli_utils import DPOScriptArguments, TrlParser\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nfrom accelerate import PartialState\nfrom trl import (\n    DPOConfig,\n    DPOTrainer,\n    ModelConfig,\n    get_kbit_device_map,\n    get_peft_config,\n    get_quantization_config,\n    maybe_extract_prompt,\n    maybe_apply_chat_template,\n)\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((DPOScriptArguments, DPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n\n    ################\n    # Model & Tokenizer\n    ###################\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n    peft_config = get_peft_config(model_config)\n    if peft_config is None:\n        ref_model = AutoModelForCausalLM.from_pretrained(\n            model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n        )\n    else:\n        ref_model = None\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    if args.ignore_bias_buffers:\n        # torch distributed hack\n        model._ddp_params_and_buffers_to_ignore = [\n            name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool\n        ]\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(maybe_extract_prompt, num_proc=training_args.dataset_num_proc)\n        dataset = dataset.map(\n            maybe_apply_chat_template, num_proc=training_args.dataset_num_proc, fn_kwargs={\"tokenizer\": tokenizer}\n        )\n\n    ##########\n    # Training\n    ################\n    trainer = DPOTrainer(\n        model,\n        ref_model,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n        peft_config=peft_config,\n    )\n\n    trainer.train()\n    metrics = trainer.evaluate()\n    trainer.log_metrics(\"eval\", metrics)\n    trainer.save_metrics(\"eval\", metrics)\n    trainer.save_model(training_args.output_dir)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nRun the KTO training script with the commands below. In general, the optimal configuration for KTO will be similar to that of DPO.\n\n# Full training:\npython examples/scripts/kto.py \\\n    --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \\\n    --per_device_train_batch_size 16 \\\n    --num_train_epochs 1 \\\n    --learning_rate 5e-7 \\\n    --lr_scheduler_type=cosine \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=kto-aligned-model \\\n    --warmup_ratio 0.1 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step\n\n# QLoRA:\npython examples/scripts/kto.py \\\n    --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \\\n    --per_device_train_batch_size 8 \\\n    --num_train_epochs 1 \\\n    --learning_rate 5e-7 \\\n    --lr_scheduler_type=cosine \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=kto-aligned-model-lora \\\n    --warmup_ratio 0.1 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step \\\n    --use_peft \\\n    --load_in_4bit \\\n    --lora_target_modules=all-linear \\\n    --lora_r=16 \\\n    --lora_alpha=16\n\"\"\"\n\nfrom dataclasses import dataclass\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser\n\nfrom trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The arguments for the KTO training script.\n    \"\"\"\n\n    dataset_name: str = \"trl-lib/kto-mix-14k\"\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, KTOConfig, ModelConfig))\n    script_args, training_args, model_args = parser.parse_args_into_dataclasses()\n\n    # Load a pretrained model\n    model = AutoModelForCausalLM.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n    ref_model = AutoModelForCausalLM.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    # If we are aligning a base model, we use ChatML as the default template\n    if tokenizer.chat_template is None:\n        model, tokenizer = setup_chat_format(model, tokenizer)\n\n    # Load the dataset\n    dataset = load_dataset(script_args.dataset_name)\n\n    # If needed, reformat a DPO-formatted dataset (prompt, chosen, rejected) to a KTO-format (prompt, completion, label)\n    dataset = maybe_unpair_preference_dataset(dataset, num_proc=training_args.dataset_num_proc)\n\n    # Apply chat template\n    def format_dataset(example):\n        if isinstance(example[\"completion\"], str):\n            example[\"prompt\"] = tokenizer.apply_chat_template(example[\"prompt\"], tokenize=False)\n            example[\"completion\"] = tokenizer.apply_chat_template(example[\"completion\"], tokenize=False)\n        else:\n            example[\"prompt\"] = tokenizer.apply_chat_template(example[\"completion\"][:-1], tokenize=False)\n            example[\"completion\"] = tokenizer.apply_chat_template([example[\"completion\"][-1]], tokenize=False)\n        return example\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(format_dataset, num_proc=training_args.dataset_num_proc)\n\n    # Initialize the KTO trainer\n    kto_trainer = KTOTrainer(\n        model,\n        ref_model,\n        args=training_args,\n        train_dataset=dataset[\"train\"],\n        eval_dataset=dataset[\"test\"],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_args),\n    )\n\n    # Train and push the model to the Hub\n    kto_trainer.train()\n    kto_trainer.save_model(training_args.output_dir)\n    kto_trainer.push_to_hub()\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import (\n    AutoTokenizer,\n    BitsAndBytesConfig,\n    HfArgumentParser,\n    is_torch_npu_available,\n    is_torch_xpu_available,\n)\n\nfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer\nfrom trl.core import LengthSampler\n\n\ninput_min_text_length = 6\ninput_max_text_length = 12\n\n\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The name of the Casual LM model we wish to fine with PPO\n    \"\"\"\n\n    model_name: Optional[str] = field(default=\"huggyllama/llama-7b\", metadata={\"help\": \"the model name\"})\n    dataset_name: Optional[str] = field(default=\"Anthropic/hh-rlhf\", metadata={\"help\": \"the dataset name\"})\n    rm_adapter: Optional[str] = field(\n        default=\"trl-lib/llama-7b-hh-rm-adapter\", metadata={\"help\": \"the rm adapter name\"}\n    )\n    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n    use_safetensors: Optional[bool] = field(default=False, metadata={\"help\": \"Use safetensors\"})\n    seed: Optional[int] = field(default=0, metadata={\"help\": \"the random seed\"})\n    use_score_scaling: Optional[bool] = field(default=False, metadata={\"help\": \"Use score scaling\"})\n    use_score_norm: Optional[bool] = field(\n        default=False, metadata={\"help\": \"Use score normalization. Only applicable if use_score_scaling is True\"}\n    )\n    score_clip: Optional[float] = field(default=None, metadata={\"help\": \"Score clipping\"})\n    dataset_num_proc: Optional[int] = field(\n        default=None, metadata={\"help\": \"The number of workers to use to tokenize the data\"}\n    )\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\n\ndef create_and_prepare_dataset(tokenizer, num_proc):\n    dataset = load_dataset(script_args.dataset_name, split=\"train[:1%]\")\n\n    input_size = LengthSampler(input_min_text_length, input_max_text_length)\n\n    def tokenize(example):\n        text_size = input_size()\n        example[\"input_ids\"] = tokenizer.encode(example[\"chosen\"])[:text_size]\n        example[\"query\"] = tokenizer.decode(example[\"input_ids\"])\n        return example\n\n    dataset = dataset.map(tokenize, batched=False, num_proc=num_proc)\n    dataset.set_format(\"torch\")\n    return dataset\n\n\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=32,\n    lora_dropout=0.05,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\nnf4_config = BitsAndBytesConfig(\n    load_in_4bit=True, bnb_4bit_quant_type=\"nf4\", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16\n)\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\n    script_args.model_name,\n    device_map={\"\": \"xpu:0\"} if is_torch_xpu_available() else {\"\": \"npu:0\"} if is_torch_npu_available else {\"\": 0},\n    peft_config=lora_config,\n    quantization_config=nf4_config,\n    reward_adapter=script_args.rm_adapter,\n    use_safetensors=script_args.use_safetensors,\n)\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name)\n\ntokenizer.pad_token = tokenizer.eos_token\n\n# Compute that only on the main process for faster data processing.\n# see: https://github.com/huggingface/trl/pull/1255\nwith PartialState().local_main_process_first():\n    dataset = create_and_prepare_dataset(tokenizer, script_args.dataset_num_proc)\n\n\ndef collator(data):\n    return {key: [d[key] for d in data] for key in data[0]}\n\n\nconfig = PPOConfig(\n    model_name=script_args.model_name,\n    log_with=script_args.log_with,\n    learning_rate=1e-5,\n    batch_size=8,\n    mini_batch_size=2,\n    gradient_accumulation_steps=2,\n    optimize_cuda_cache=True,\n    seed=script_args.seed,\n    use_score_scaling=script_args.use_score_scaling,\n    use_score_norm=script_args.use_score_norm,\n    score_clip=script_args.score_clip,\n)\n\nppo_trainer = PPOTrainer(\n    config,\n    model,\n    ref_model=None,\n    tokenizer=tokenizer,\n    dataset=dataset,\n    data_collator=collator,\n)\n\ngeneration_kwargs = {\n    \"top_k\": 0.0,\n    \"top_p\": 0.9,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.pad_token_id,\n    \"max_new_tokens\": 32,\n}\n\nfor _epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n    question_tensors = batch[\"input_ids\"]\n\n    response_tensors = ppo_trainer.generate(\n        question_tensors,\n        return_prompt=False,\n        **generation_kwargs,\n    )\n    batch[\"response\"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)\n\n    # Compute reward score\n    texts = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]\n    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors=\"pt\").to(ppo_trainer.accelerator.device)\n    raw_rewards = ppo_trainer.accelerator.unwrap_model(ppo_trainer.model).compute_reward_score(**inputs)\n    rewards = [raw_rewards[i, -1, 1] for i in range(len(raw_rewards))]  # take last token\n\n    # Run PPO step\n    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)\n    ppo_trainer.log_stats(stats, batch, rewards)\n\n\n# Copyright 2023 metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTotal Batch size = 128 = 4 (num_gpus) * 8 (per_device_batch) * 4 (accumulation steps)\nFeel free to reduce batch size or increasing truncated_rand_backprop_min to a higher value to reduce memory usage.\n\nCUDA_VISIBLE_DEVICES=0,1,2,3 python examples/scripts/alignprop.py \\\n    --num_epochs=20 \\\n    --train_gradient_accumulation_steps=4 \\\n    --sample_num_steps=50 \\\n    --train_batch_size=8 \\\n    --tracker_project_name=\"stable_diffusion_training\" \\\n    --log_with=\"wandb\"\n\n\"\"\"\n\nfrom dataclasses import dataclass, field\n\nimport numpy as np\nfrom transformers import HfArgumentParser\n\nfrom trl import AlignPropConfig, AlignPropTrainer, DefaultDDPOStableDiffusionPipeline\nfrom trl.models.auxiliary_modules import aesthetic_scorer\n\n\n@dataclass\nclass ScriptArguments:\n    pretrained_model: str = field(\n        default=\"runwayml/stable-diffusion-v1-5\", metadata={\"help\": \"the pretrained model to use\"}\n    )\n    pretrained_revision: str = field(default=\"main\", metadata={\"help\": \"the pretrained model revision to use\"})\n    hf_hub_model_id: str = field(\n        default=\"alignprop-finetuned-stable-diffusion\", metadata={\"help\": \"HuggingFace repo to save model weights to\"}\n    )\n    hf_hub_aesthetic_model_id: str = field(\n        default=\"trl-lib/ddpo-aesthetic-predictor\",\n        metadata={\"help\": \"HuggingFace model ID for aesthetic scorer model weights\"},\n    )\n    hf_hub_aesthetic_model_filename: str = field(\n        default=\"aesthetic-model.pth\",\n        metadata={\"help\": \"HuggingFace model filename for aesthetic scorer model weights\"},\n    )\n    use_lora: bool = field(default=True, metadata={\"help\": \"Whether to use LoRA.\"})\n\n\n# list of example prompts to feed stable diffusion\nanimals = [\n    \"cat\",\n    \"dog\",\n    \"horse\",\n    \"monkey\",\n    \"rabbit\",\n    \"zebra\",\n    \"spider\",\n    \"bird\",\n    \"sheep\",\n    \"deer\",\n    \"cow\",\n    \"goat\",\n    \"lion\",\n    \"frog\",\n    \"chicken\",\n    \"duck\",\n    \"goose\",\n    \"bee\",\n    \"pig\",\n    \"turkey\",\n    \"fly\",\n    \"llama\",\n    \"camel\",\n    \"bat\",\n    \"gorilla\",\n    \"hedgehog\",\n    \"kangaroo\",\n]\n\n\ndef prompt_fn():\n    return np.random.choice(animals), {}\n\n\ndef image_outputs_logger(image_pair_data, global_step, accelerate_logger):\n    # For the sake of this example, we will only log the last batch of images\n    # and associated data\n    result = {}\n    images, prompts, _ = [image_pair_data[\"images\"], image_pair_data[\"prompts\"], image_pair_data[\"rewards\"]]\n    for i, image in enumerate(images[:4]):\n        prompt = prompts[i]\n        result[f\"{prompt}\"] = image.unsqueeze(0).float()\n    accelerate_logger.log_images(\n        result,\n        step=global_step,\n    )\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, AlignPropConfig))\n    args, training_args = parser.parse_args_into_dataclasses()\n    training_args.project_kwargs = {\n        \"logging_dir\": \"./logs\",\n        \"automatic_checkpoint_naming\": True,\n        \"total_limit\": 5,\n        \"project_dir\": \"./save\",\n    }\n\n    pipeline = DefaultDDPOStableDiffusionPipeline(\n        args.pretrained_model, pretrained_model_revision=args.pretrained_revision, use_lora=args.use_lora\n    )\n    trainer = AlignPropTrainer(\n        training_args,\n        aesthetic_scorer(args.hf_hub_aesthetic_model_id, args.hf_hub_aesthetic_model_filename),\n        prompt_fn,\n        pipeline,\n        image_samples_hook=image_outputs_logger,\n    )\n\n    trainer.train()\n\n    trainer.push_to_hub(args.hf_hub_model_id)\n\n\n# flake8: noqa\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\n# regular:\npython examples/scripts/sft.py \\\n    --model_name_or_path=\"facebook/opt-350m\" \\\n    --dataset_text_field=\"text\" \\\n    --report_to=\"wandb\" \\\n    --learning_rate=1.41e-5 \\\n    --per_device_train_batch_size=64 \\\n    --gradient_accumulation_steps=16 \\\n    --output_dir=\"sft_openassistant-guanaco\" \\\n    --logging_steps=1 \\\n    --num_train_epochs=3 \\\n    --max_steps=-1 \\\n    --push_to_hub \\\n    --gradient_checkpointing\n\n# peft:\npython examples/scripts/sft.py \\\n    --model_name_or_path=\"facebook/opt-350m\" \\\n    --dataset_text_field=\"text\" \\\n    --report_to=\"wandb\" \\\n    --learning_rate=1.41e-5 \\\n    --per_device_train_batch_size=64 \\\n    --gradient_accumulation_steps=16 \\\n    --output_dir=\"sft_openassistant-guanaco\" \\\n    --logging_steps=1 \\\n    --num_train_epochs=3 \\\n    --max_steps=-1 \\\n    --push_to_hub \\\n    --gradient_checkpointing \\\n    --use_peft \\\n    --lora_r=64 \\\n    --lora_alpha=16\n\"\"\"\n\nfrom trl.commands.cli_utils import SFTScriptArguments, TrlParser\n\n\nfrom datasets import load_dataset\n\nfrom transformers import AutoTokenizer\n\nfrom trl import (\n    ModelConfig,\n    SFTConfig,\n    SFTTrainer,\n    get_peft_config,\n    get_quantization_config,\n    get_kbit_device_map,\n)\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((SFTScriptArguments, SFTConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n\n    ################\n    # Model init kwargs & Tokenizer\n    ################\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        trust_remote_code=model_config.trust_remote_code,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=model_config.torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    training_args.model_init_kwargs = model_kwargs\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True\n    )\n    tokenizer.pad_token = tokenizer.eos_token\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n\n    ################\n    # Training\n    ################\n    trainer = SFTTrainer(\n        model=model_config.model_name_or_path,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nRun the CPO training script with the following command with some example arguments.\nIn general, the optimal configuration for CPO will be similar to that of DPO:\n\n# regular:\npython examples/scripts/cpo.py \\\n    --model_name_or_path=gpt2 \\\n    --per_device_train_batch_size 4 \\\n    --max_steps 1000 \\\n    --learning_rate 8e-6 \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=\"gpt2-aligned-cpo\" \\\n    --warmup_steps 150 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step \\\n    --no_remove_unused_columns\n\n# peft:\npython examples/scripts/cpo.py \\\n    --model_name_or_path=gpt2 \\\n    --per_device_train_batch_size 4 \\\n    --max_steps 1000 \\\n    --learning_rate 8e-5 \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 10 \\\n    --eval_steps 500 \\\n    --output_dir=\"gpt2-lora-aligned-cpo\" \\\n    --optim rmsprop \\\n    --warmup_steps 150 \\\n    --report_to wandb \\\n    --bf16 \\\n    --logging_first_step \\\n    --no_remove_unused_columns \\\n    --use_peft \\\n    --lora_r=16 \\\n    --lora_alpha=16\n\"\"\"\n\nfrom dataclasses import dataclass, field\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser\n\nfrom trl import CPOConfig, CPOTrainer, ModelConfig, get_peft_config\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n@dataclass\nclass ScriptArguments:\n    dataset_name: str = field(\n        default=\"trl-internal-testing/hh-rlhf-helpful-base-trl-style\",\n        metadata={\"help\": \"The name of the dataset to use.\"},\n    )\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, CPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_into_dataclasses()\n\n    ################\n    # Model & Tokenizer\n    ################\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n\n    def process(row):\n        row[\"chosen\"] = tokenizer.apply_chat_template(row[\"chosen\"], tokenize=False)\n        row[\"rejected\"] = tokenizer.apply_chat_template(row[\"rejected\"], tokenize=False)\n        return row\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(process, num_proc=training_args.dataset_num_proc)\n\n    ################\n    # Training\n    ################\n    trainer = CPOTrainer(\n        model,\n        args=training_args,\n        train_dataset=dataset[\"train\"],\n        eval_dataset=dataset[\"test\"],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n\n    # train and save the model\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n\n\n# flake8: noqa\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\npip install pillow\n\npython examples/scripts/vsft_llava.py \\\n    --dataset_name HuggingFaceH4/llava-instruct-mix-vsft \\\n    --model_name_or_path llava-hf/llava-1.5-7b-hf \\\n    --per_device_train_batch_size 8 \\\n    --gradient_accumulation_steps 8 \\\n    --output_dir sft-llava-1.5-7b-hf \\\n    --bf16 \\\n    --torch_dtype bfloat16 \\\n    --gradient_checkpointing \\\n    --use_peft \\\n    --dataloader_num_workers 32 \\\n    --lora_target_modules=all-linear\n\nFor LLaVA-NeXT, use: (requires transformers>=4.45)\n    --model_name_or_path llava-hf/llava-v1.6-mistral-7b-hf\n\"\"\"\n\nfrom trl.commands.cli_utils import SFTScriptArguments, TrlParser\n\nimport torch\nfrom accelerate import Accelerator\nfrom datasets import load_dataset\n\nfrom transformers import AutoModelForVision2Seq, AutoProcessor\n\nfrom trl import (\n    ModelConfig,\n    SFTConfig,\n    SFTTrainer,\n    get_peft_config,\n    get_quantization_config,\n    get_kbit_device_map,\n)\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((SFTScriptArguments, SFTConfig, ModelConfig))\n    sft_script_args, training_args, model_config = parser.parse_args_and_config()\n    training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)\n    training_args.dataset_text_field = \"\"  # need a dummy field\n    training_args.remove_unused_columns = False\n    training_args.dataset_kwargs = {\"skip_prepare_dataset\": True}\n\n    ################\n    # Model, Tokenizer & Processor\n    ################\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    processor = AutoProcessor.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code\n    )\n\n    model = AutoModelForVision2Seq.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n\n    ################\n    # Create a data collator to encode text and image pairs\n    ################\n    def collate_fn(examples):\n        # Get the texts and images, and apply the chat template\n        texts = [processor.apply_chat_template(example[\"messages\"], tokenize=False) for example in examples]\n        images = [example[\"images\"][0] for example in examples]\n\n        # Tokenize the texts and process the images\n        batch = processor(texts, images, return_tensors=\"pt\", padding=True)\n\n        # The labels are the input_ids, and we mask the padding tokens in the loss computation\n        labels = batch[\"input_ids\"].clone()\n        labels[labels == processor.tokenizer.pad_token_id] = -100\n        batch[\"labels\"] = labels\n\n        return batch\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(sft_script_args.dataset_name)\n\n    ################\n    # Training\n    ################\n    trainer = SFTTrainer(\n        model=model,\n        args=training_args,\n        data_collator=collate_fn,\n        train_dataset=dataset[sft_script_args.dataset_train_split],\n        eval_dataset=dataset[sft_script_args.dataset_test_split],\n        tokenizer=processor.tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n\n    trainer.train()\n\n    trainer.save_model(training_args.output_dir)\n    trainer.push_to_hub()\n    if Accelerator().is_main_process:\n        processor.push_to_hub(training_args.hub_model_id)\n\n\n# Copyright 2023 metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\npython examples/scripts/ddpo.py \\\n    --num_epochs=200 \\\n    --train_gradient_accumulation_steps=1 \\\n    --sample_num_steps=50 \\\n    --sample_batch_size=6 \\\n    --train_batch_size=3 \\\n    --sample_num_batches_per_epoch=4 \\\n    --per_prompt_stat_tracking=True \\\n    --per_prompt_stat_tracking_buffer_size=32 \\\n    --tracker_project_name=\"stable_diffusion_training\" \\\n    --log_with=\"wandb\"\n\"\"\"\n\nimport os\nfrom dataclasses import dataclass, field\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom huggingface_hub import hf_hub_download\nfrom huggingface_hub.utils import EntryNotFoundError\nfrom transformers import CLIPModel, CLIPProcessor, HfArgumentParser, is_torch_npu_available, is_torch_xpu_available\n\nfrom trl import DDPOConfig, DDPOTrainer, DefaultDDPOStableDiffusionPipeline\n\n\n@dataclass\nclass ScriptArguments:\n    pretrained_model: str = field(\n        default=\"runwayml/stable-diffusion-v1-5\", metadata={\"help\": \"the pretrained model to use\"}\n    )\n    pretrained_revision: str = field(default=\"main\", metadata={\"help\": \"the pretrained model revision to use\"})\n    hf_hub_model_id: str = field(\n        default=\"ddpo-finetuned-stable-diffusion\", metadata={\"help\": \"HuggingFace repo to save model weights to\"}\n    )\n    hf_hub_aesthetic_model_id: str = field(\n        default=\"trl-lib/ddpo-aesthetic-predictor\",\n        metadata={\"help\": \"HuggingFace model ID for aesthetic scorer model weights\"},\n    )\n    hf_hub_aesthetic_model_filename: str = field(\n        default=\"aesthetic-model.pth\",\n        metadata={\"help\": \"HuggingFace model filename for aesthetic scorer model weights\"},\n    )\n    use_lora: bool = field(default=True, metadata={\"help\": \"Whether to use LoRA.\"})\n\n\nclass MLP(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.layers = nn.Sequential(\n            nn.Linear(768, 1024),\n            nn.Dropout(0.2),\n            nn.Linear(1024, 128),\n            nn.Dropout(0.2),\n            nn.Linear(128, 64),\n            nn.Dropout(0.1),\n            nn.Linear(64, 16),\n            nn.Linear(16, 1),\n        )\n\n    @torch.no_grad()\n    def forward(self, embed):\n        return self.layers(embed)\n\n\nclass AestheticScorer(torch.nn.Module):\n    \"\"\"\n    This model attempts to predict the aesthetic score of an image. The aesthetic score\n    is a numerical approximation of how much a specific image is liked by humans on average.\n    This is from https://github.com/christophschuhmann/improved-aesthetic-predictor\n    \"\"\"\n\n    def __init__(self, *, dtype, model_id, model_filename):\n        super().__init__()\n        self.clip = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n        self.processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-large-patch14\")\n        self.mlp = MLP()\n        try:\n            cached_path = hf_hub_download(model_id, model_filename)\n        except EntryNotFoundError:\n            cached_path = os.path.join(model_id, model_filename)\n        state_dict = torch.load(cached_path, map_location=torch.device(\"cpu\"), weights_only=True)\n        self.mlp.load_state_dict(state_dict)\n        self.dtype = dtype\n        self.eval()\n\n    @torch.no_grad()\n    def __call__(self, images):\n        device = next(self.parameters()).device\n        inputs = self.processor(images=images, return_tensors=\"pt\")\n        inputs = {k: v.to(self.dtype).to(device) for k, v in inputs.items()}\n        embed = self.clip.get_image_features(**inputs)\n        # normalize embedding\n        embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)\n        return self.mlp(embed).squeeze(1)\n\n\ndef aesthetic_scorer(hub_model_id, model_filename):\n    scorer = AestheticScorer(\n        model_id=hub_model_id,\n        model_filename=model_filename,\n        dtype=torch.float32,\n    )\n    if is_torch_npu_available():\n        scorer = scorer.npu()\n    elif is_torch_xpu_available():\n        scorer = scorer.xpu()\n    else:\n        scorer = scorer.cuda()\n\n    def _fn(images, prompts, metadata):\n        images = (images * 255).round().clamp(0, 255).to(torch.uint8)\n        scores = scorer(images)\n        return scores, {}\n\n    return _fn\n\n\n# list of example prompts to feed stable diffusion\nanimals = [\n    \"cat\",\n    \"dog\",\n    \"horse\",\n    \"monkey\",\n    \"rabbit\",\n    \"zebra\",\n    \"spider\",\n    \"bird\",\n    \"sheep\",\n    \"deer\",\n    \"cow\",\n    \"goat\",\n    \"lion\",\n    \"frog\",\n    \"chicken\",\n    \"duck\",\n    \"goose\",\n    \"bee\",\n    \"pig\",\n    \"turkey\",\n    \"fly\",\n    \"llama\",\n    \"camel\",\n    \"bat\",\n    \"gorilla\",\n    \"hedgehog\",\n    \"kangaroo\",\n]\n\n\ndef prompt_fn():\n    return np.random.choice(animals), {}\n\n\ndef image_outputs_logger(image_data, global_step, accelerate_logger):\n    # For the sake of this example, we will only log the last batch of images\n    # and associated data\n    result = {}\n    images, prompts, _, rewards, _ = image_data[-1]\n\n    for i, image in enumerate(images):\n        prompt = prompts[i]\n        reward = rewards[i].item()\n        result[f\"{prompt:.25} | {reward:.2f}\"] = image.unsqueeze(0).float()\n\n    accelerate_logger.log_images(\n        result,\n        step=global_step,\n    )\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, DDPOConfig))\n    args, training_args = parser.parse_args_into_dataclasses()\n    training_args.project_kwargs = {\n        \"logging_dir\": \"./logs\",\n        \"automatic_checkpoint_naming\": True,\n        \"total_limit\": 5,\n        \"project_dir\": \"./save\",\n    }\n\n    pipeline = DefaultDDPOStableDiffusionPipeline(\n        args.pretrained_model, pretrained_model_revision=args.pretrained_revision, use_lora=args.use_lora\n    )\n\n    trainer = DDPOTrainer(\n        training_args,\n        aesthetic_scorer(args.hf_hub_aesthetic_model_id, args.hf_hub_aesthetic_model_filename),\n        prompt_fn,\n        pipeline,\n        image_samples_hook=image_outputs_logger,\n    )\n\n    trainer.train()\n\n    trainer.push_to_hub(args.hf_hub_model_id)\n\n\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\npython examples/scripts/ppo.py \\\n    --log_with=wandb\n\"\"\"\n\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom accelerate import Accelerator, PartialState\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import AutoTokenizer, HfArgumentParser, is_torch_npu_available, is_torch_xpu_available, pipeline\n\nfrom trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed\nfrom trl.core import LengthSampler\n\n\ntqdm.pandas()\n\n\n@dataclass\nclass ScriptArguments:\n    use_seq2seq: bool = field(default=False, metadata={\"help\": \"whether to use seq2seq\"})\n    trust_remote_code: bool = field(default=False, metadata={\"help\": \"Enable `trust_remote_code`\"})\n\n    # LoraConfig\n    use_peft: bool = field(default=False, metadata={\"help\": \"whether to use peft\"})\n    lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n    lora_r: Optional[int] = field(default=16, metadata={\"help\": \"the lora r parameter\"})\n\n\nparser = HfArgumentParser((ScriptArguments, PPOConfig))\nargs, ppo_config = parser.parse_args_into_dataclasses()\n\n# We then define the arguments to pass to the sentiment analysis pipeline.\n# We set `return_all_scores` to True to get the sentiment score for each token.\nsent_kwargs = {\"return_all_scores\": True, \"function_to_apply\": \"none\", \"batch_size\": 16}\n\ntrl_model_class = AutoModelForCausalLMWithValueHead if not args.use_seq2seq else AutoModelForSeq2SeqLMWithValueHead\n\ntokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name)\ntokenizer.pad_token = tokenizer.eos_token\n\n\n# Below is an example function to build the dataset. In our case, we use the IMDB dataset\n# from the `datasets` library. One should customize this function to train the model on\n# its own dataset.\ndef build_dataset(query_dataset, dataset_num_proc, input_min_text_length=2, input_max_text_length=8):\n    \"\"\"\n    Build dataset for training. This builds the dataset from `load_dataset`, one should\n    customize this function to train the model on its own dataset.\n\n    Args:\n        query_dataset (`str`):\n            The name of the dataset to be loaded.\n\n    Returns:\n        dataloader (`torch.utils.data.DataLoader`):\n            The dataloader for the dataset.\n    \"\"\"\n    # load imdb with datasets\n    dataset = load_dataset(query_dataset, split=\"train\")\n    dataset = dataset.rename_columns({\"text\": \"review\"})\n    dataset = dataset.filter(lambda x: len(x[\"review\"]) > 200, num_proc=dataset_num_proc)\n\n    input_size = LengthSampler(input_min_text_length, input_max_text_length)\n\n    def tokenize(sample):\n        sample[\"input_ids\"] = tokenizer.encode(sample[\"review\"])[: input_size()]\n        sample[\"query\"] = tokenizer.decode(sample[\"input_ids\"])\n        return sample\n\n    dataset = dataset.map(tokenize, num_proc=dataset_num_proc)\n    dataset.set_format(type=\"torch\")\n    return dataset\n\n\n# We retrieve the dataloader by calling the `build_dataset` function.\n# Compute that only on the main process for faster data processing.\n# see: https://github.com/huggingface/trl/pull/1255\nwith PartialState().local_main_process_first():\n    dataset = build_dataset(ppo_config.query_dataset, ppo_config.dataset_num_proc)\n\n\ndef collator(data):\n    return {key: [d[key] for d in data] for key in data[0]}\n\n\n# set seed before initializing value head for deterministic eval\nset_seed(ppo_config.seed)\n\n# Now let's build the model, the reference model, and the tokenizer.\nif not args.use_peft:\n    ref_model = trl_model_class.from_pretrained(ppo_config.model_name, trust_remote_code=args.trust_remote_code)\n    device_map = None\n    peft_config = None\nelse:\n    peft_config = LoraConfig(\n        r=args.lora_r,\n        lora_alpha=args.lora_alpha,\n        bias=\"none\",\n        task_type=\"CAUSAL_LM\",\n    )\n    ref_model = None\n    # Copy the model to each device\n    device_map = {\"\": Accelerator().local_process_index}\n\nmodel = trl_model_class.from_pretrained(\n    ppo_config.model_name,\n    trust_remote_code=args.trust_remote_code,\n    device_map=device_map,\n    peft_config=peft_config,\n)\n\n\ntokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name)\n\n# Some tokenizers like GPT-2's don't have a padding token by default, so we set one here.\ntokenizer.pad_token_id = tokenizer.eos_token_id\n\n# We then build the PPOTrainer, passing the model, the reference model, the tokenizer\nppo_trainer = PPOTrainer(ppo_config, model, ref_model, tokenizer, dataset=dataset, data_collator=collator)\n\n# We then build the sentiment analysis pipeline, passing the model name and the\n# sentiment analysis pipeline arguments. Let's also make sure to set the device\n# to the same device as the PPOTrainer.\ndevice = ppo_trainer.accelerator.device\nif ppo_trainer.accelerator.num_processes == 1:\n    if is_torch_xpu_available():\n        device = \"xpu:0\"\n    elif is_torch_npu_available():\n        device = \"npu:0\"\n    else:\n        device = 0 if torch.cuda.is_available() else \"cpu\"  # to avoid a `pipeline` bug\nds_plugin = ppo_trainer.accelerator.state.deepspeed_plugin\ntask, model_name = ppo_config.reward_model.split(\":\")\nif ds_plugin is not None and ds_plugin.is_zero3_init_enabled():\n    with ds_plugin.zero3_init_context_manager(enable=False):\n        sentiment_pipe = pipeline(task, model=model_name, device=device)\nelse:\n    sentiment_pipe = pipeline(task, model=model_name, device=device)\n\n# Some tokenizers like GPT-2's don't have a padding token by default, so we set one here.\nif sentiment_pipe.tokenizer.pad_token_id is None:\n    sentiment_pipe.tokenizer.pad_token_id = tokenizer.pad_token_id\n\nif sentiment_pipe.model.config.pad_token_id is None:\n    sentiment_pipe.model.config.pad_token_id = tokenizer.pad_token_id\n\n# We then define the arguments to pass to the `generate` function. These arguments\n# are passed to the `generate` function of the PPOTrainer, which is a wrapper around\n# the `generate` function of the trained model.\ngeneration_kwargs = {\n    \"min_length\": -1,\n    \"top_k\": 0.0,\n    \"top_p\": 1.0,\n    \"do_sample\": True,\n    \"pad_token_id\": tokenizer.eos_token_id,\n    \"max_new_tokens\": 32,\n}\n\nfor batch in tqdm(ppo_trainer.dataloader):\n    query_tensors = batch[\"input_ids\"]\n\n    # Get response from gpt2\n    response_tensors, ref_response_tensors = ppo_trainer.generate(\n        query_tensors, return_prompt=False, generate_ref_response=True, **generation_kwargs\n    )\n    batch[\"response\"] = tokenizer.batch_decode(response_tensors)\n    batch[\"ref_response\"] = tokenizer.batch_decode(ref_response_tensors)\n\n    # Compute sentiment score\n    texts = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]\n    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)\n    rewards = [torch.tensor(output[1][\"score\"]) for output in pipe_outputs]\n    ref_texts = [q + r for q, r in zip(batch[\"query\"], batch[\"ref_response\"])]\n    ref_pipe_outputs = sentiment_pipe(ref_texts, **sent_kwargs)\n    ref_rewards = [torch.tensor(output[1][\"score\"]) for output in ref_pipe_outputs]\n    batch[\"ref_rewards\"] = ref_rewards\n\n    # Run PPO step\n    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)\n    ppo_trainer.log_stats(stats, batch, rewards, columns_to_log=[\"query\", \"response\", \"ref_response\", \"ref_rewards\"])\n\n\n# flake8: noqa\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\n# Full training:\npython examples/scripts/gkd.py \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --teacher_model_name_or_path Qwen/Qwen2-1.5B-Instruct \\\n    --dataset_name trl-lib/chatbot_arena_completions \\\n    --learning_rate 2e-5 \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 8 \\\n    --output_dir gkd-model \\\n    --logging_steps 10 \\\n    --num_train_epochs 1 \\\n    --push_to_hub \\\n    --gradient_checkpointing\n\n# LoRA:\npython examples/scripts/gkd.py \\\n    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \\\n    --teacher_model_name_or_path Qwen/Qwen2-1.5B-Instruct \\\n    --dataset_name trl-lib/chatbot_arena_completions \\\n    --learning_rate 2e-4 \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 8 \\\n    --output_dir gkd-model \\\n    --logging_steps 10 \\\n    --num_train_epochs 1 \\\n    --push_to_hub \\\n    --gradient_checkpointing \\\n    --use_peft \\\n    --lora_r 64 \\\n    --lora_alpha 16\n\"\"\"\n\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, GenerationConfig\n\nfrom trl import (\n    GKDConfig,\n    GKDTrainer,\n    ModelConfig,\n    get_kbit_device_map,\n    get_peft_config,\n    get_quantization_config,\n    maybe_apply_chat_template,\n    LogCompletionsCallback,\n)\nfrom trl.commands.cli_utils import SFTScriptArguments, TrlParser\nfrom accelerate import PartialState\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((SFTScriptArguments, GKDConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n\n    ################\n    # Model & Tokenizer\n    ################\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        trust_remote_code=model_config.trust_remote_code,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=model_config.torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    training_args.model_init_kwargs = model_kwargs\n\n    teacher_model_kwargs = dict(\n        revision=model_config.model_revision,\n        trust_remote_code=model_config.trust_remote_code,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=model_config.torch_dtype,\n        use_cache=True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    training_args.teacher_model_init_kwargs = teacher_model_kwargs\n\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        trust_remote_code=model_config.trust_remote_code,\n        padding_side=\"left\",\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(\n            lambda x: {\n                \"prompt\": tokenizer.apply_chat_template(x[\"prompt\"], tokenize=False, add_generation_prompt=True)\n            },\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    ################\n    # Training\n    ################\n    trainer = GKDTrainer(\n        model=model_config.model_name_or_path,\n        teacher_model=training_args.teacher_model_name_or_path,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n    generation_config = GenerationConfig(\n        max_new_tokens=training_args.max_new_tokens, do_sample=True, temperature=training_args.temperature\n    )\n    completions_callback = LogCompletionsCallback(trainer, generation_config, num_prompts=8)\n    trainer.add_callback(completions_callback)\n    trainer.train()\n\n    trainer.save_model(training_args.output_dir)\n\n\n# flake8: noqa\n# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\naccelerate launch examples/scripts/dpo_visual.py \\\n    --dataset_name HuggingFaceH4/rlaif-v_formatted \\\n    --model_name_or_path HuggingFaceM4/idefics2-8b \\\n    --per_device_train_batch_size 2 \\\n    --gradient_accumulation_steps 32 \\\n    --dataset_num_proc 32 \\\n    --output_dir dpo_idefics_rlaif-v \\\n    --bf16 \\\n    --torch_dtype bfloat16 \\\n    --gradient_checkpointing \\\n    --use_peft \\\n    --lora_target_modules=all-linear\n\"\"\"\n\nfrom trl.commands.cli_utils import DPOScriptArguments, TrlParser\nfrom accelerate import PartialState\n\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoModelForVision2Seq, AutoProcessor\n\nfrom trl import (\n    DPOConfig,\n    DPOTrainer,\n    ModelConfig,\n    get_kbit_device_map,\n    get_peft_config,\n    get_quantization_config,\n)\n\n\nif __name__ == \"__main__\":\n    parser = TrlParser((DPOScriptArguments, DPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n\n    ################\n    # Model & Tokenizer\n    ################\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n    model = AutoModelForVision2Seq.from_pretrained(\n        model_config.model_name_or_path,\n        trust_remote_code=model_config.trust_remote_code,\n        **model_kwargs,\n    )\n    peft_config = get_peft_config(model_config)\n    if peft_config is None:\n        ref_model = AutoModelForVision2Seq.from_pretrained(\n            model_config.model_name_or_path,\n            trust_remote_code=model_config.trust_remote_code,\n            **model_kwargs,\n        )\n    else:\n        ref_model = None\n    processor = AutoProcessor.from_pretrained(\n        model_config.model_name_or_path,\n        trust_remote_code=model_config.trust_remote_code,\n        do_image_splitting=False,\n    )\n    tokenizer = processor.tokenizer\n\n    # Set up the chat template\n    if model.config.model_type == \"idefics2\":\n        pass  # the processor already has a valid chat template\n    elif model.config.model_type == \"paligemma\":\n        processor.chat_template = \"\"\"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}<|im_start|>{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] if item['type'] == 'text' %}{{ item['text'] }}<|im_end|>{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}\"\"\"\n    elif model.config.model_type == \"llava\":\n        processor.chat_template = \"\"\"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% elif item['type'] == 'image' %}<image>{% endif %}{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}\"\"\"\n\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    if args.ignore_bias_buffers:\n        # torch distributed hack\n        model._ddp_params_and_buffers_to_ignore = [\n            name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool\n        ]\n\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(args.dataset_name)\n\n    def process(row):\n        row[\"prompt\"] = processor.apply_chat_template(row[\"prompt\"], tokenize=False)\n        row[\"chosen\"] = processor.apply_chat_template(row[\"chosen\"], tokenize=False)\n        row[\"rejected\"] = processor.apply_chat_template(row[\"rejected\"], tokenize=False)\n        return row\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        dataset = dataset.map(process, num_proc=training_args.dataset_num_proc)\n\n    ################\n    # Training\n    ################\n    trainer = DPOTrainer(\n        model,\n        ref_model,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=processor,\n        peft_config=peft_config,\n    )\n\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nRun the BCO training script with the commands below. In general, the optimal configuration for BCO will be similar to that of KTO.\n\n# Full training:\npython examples/scripts/bco.py \\\n    --model_name_or_path=nnheui/stablelm-2-1_6b-sft-full \\\n    --per_device_train_batch_size 16 \\\n    --per_device_eval_batch_size 32 \\\n    --num_train_epochs 1 \\\n    --learning_rate 1e-6 \\\n    --gradient_checkpointing \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 0.01 \\\n    --eval_steps 0.2 \\\n    --save_strategy no \\\n    --output_dir=bco-aligned-model \\\n    --logging_first_step \\\n    --max_length 2048 \\\n    --max_prompt_length 1536 \\\n    --max_completion_length 1024 \\\n    --no_remove_unused_columns \\\n    --warmup_ratio 0.1 \\\n    --bf16 \\\n    --report_to wandb\n\n# QLoRA:\npython examples/scripts/bco.py \\\n    --model_name_or_path=nnheui/stablelm-2-1_6b-sft-full \\\n    --per_device_train_batch_size 16 \\\n    --per_device_eval_batch_size 32 \\\n    --num_train_epochs 1 \\\n    --learning_rate 1e-6 \\\n    --gradient_checkpointing \\\n    --gradient_accumulation_steps 1 \\\n    --logging_steps 0.01 \\\n    --eval_steps 0.2 \\\n    --save_strategy no \\\n    --output_dir=bco-aligned-model-lora \\\n    --logging_first_step \\\n    --warmup_ratio 0.1 \\\n    --report_to wandb \\\n    --max_length 2048 \\\n    --max_prompt_length 1536 \\\n    --max_completion_length 1024 \\\n    --no_remove_unused_columns \\\n    --warmup_ratio 0.1 \\\n    --bf16 \\\n    --use_peft \\\n    --load_in_4bit \\\n    --lora_target_modules=all-linear \\\n    --lora_r=16 \\\n    --lora_alpha=16\n\"\"\"\n\nimport logging\nfrom dataclasses import dataclass\nfrom functools import partial\nfrom typing import Literal, Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom accelerate import Accelerator, PartialState\nfrom datasets import Dataset, load_dataset\nfrom transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, PreTrainedModel\n\nfrom trl import BCOConfig, BCOTrainer, ModelConfig, get_peft_config, setup_chat_format\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The arguments for the BCO training script.\n    \"\"\"\n\n    llm_name: Literal[\"gpt-3.5-turbo\", \"llama-2-7b-chat\", \"llama-2-70b-chat\"] = \"gpt-3.5-turbo\"\n\n\ndef build_helpfulness_dataset(llm_name: str, num_proc: Optional[int] = None) -> Dataset:\n    \"\"\"\n    Filter `llm_name` completions and binarize given their helpfulness score.\n    If helpfulness score is 5, it is desirable. Otherwise, it is undesirable.\n    \"\"\"\n\n    def get_model_rating(example, metric: str, llm_name: str):\n        try:\n            model_index = example[\"models\"].index(llm_name)\n            return {metric: int(example[\"completions\"][model_index][\"annotations\"][metric][\"Rating\"])}\n        except ValueError as e:\n            logging.warning(e)\n            return -1\n\n    def get_model_response(example, llm_name: str):\n        try:\n            model_index = example[\"models\"].index(llm_name)\n            return {\"response\": example[\"completions\"][model_index][\"response\"]}\n        except ValueError as e:\n            logging.warning(e)\n            return -1\n\n    dataset = load_dataset(\"openbmb/UltraFeedback\")[\"train\"]\n\n    dataset = dataset.filter(lambda example: llm_name in example[\"models\"], batched=False, num_proc=num_proc)\n    dataset = dataset.filter(\n        lambda example: len(example[\"models\"]) == len(example[\"completions\"]), batched=False, num_proc=num_proc\n    )\n\n    METRIC = \"helpfulness\"\n\n    dataset = dataset.map(\n        get_model_rating,\n        batched=False,\n        fn_kwargs={\"metric\": METRIC, \"llm_name\": llm_name},\n        num_proc=num_proc,\n    )\n\n    dataset = dataset.map(\n        get_model_response,\n        batched=False,\n        fn_kwargs={\"llm_name\": llm_name},\n        num_proc=num_proc,\n    )\n\n    dataset = dataset.select_columns([\"source\", \"instruction\", \"response\", \"helpfulness\"])\n\n    dataset = dataset.rename_columns({\"instruction\": \"prompt\", \"response\": \"completion\"})\n    dataset = dataset.map(lambda example: {\"label\": example[\"helpfulness\"] >= 5}, batched=False, num_proc=num_proc)\n\n    dataset = dataset.map(\n        lambda example: {\"prompt\": [{\"role\": \"user\", \"content\": example[\"prompt\"]}]},\n        batched=False,\n        num_proc=num_proc,\n    )\n    dataset = dataset.train_test_split(test_size=0.05, seed=42)\n\n    return dataset\n\n\ndef embed_prompt(input_ids: torch.LongTensor, attention_mask: torch.LongTensor, model: PreTrainedModel):\n    \"\"\"\n    Borrowed from https://huggingface.co/nomic-ai/nomic-embed-text-v1.5#transformers\n    \"\"\"\n\n    def mean_pooling(model_output, attention_mask):\n        token_embeddings = model_output[0]\n        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n\n    with torch.no_grad():\n        model_output = model(input_ids=input_ids, attention_mask=attention_mask)\n        embeddings = mean_pooling(model_output, attention_mask)\n\n    matryoshka_dim = 512\n    # normalize embeddings\n    embeddings = F.normalize(embeddings, p=2, dim=1)\n    embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))\n    embeddings = embeddings[:, :matryoshka_dim]\n\n    return embeddings\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((ScriptArguments, BCOConfig, ModelConfig))\n    script_args, training_args, model_args = parser.parse_args_into_dataclasses()\n\n    training_args.gradient_checkpointing_kwargs = {\"use_reentrant\": True}\n\n    # Load a pretrained model\n    model = AutoModelForCausalLM.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n    ref_model = AutoModelForCausalLM.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code\n    )\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    # If we are aligning a base model, we use ChatML as the default template\n    if tokenizer.chat_template is None:\n        model, tokenizer = setup_chat_format(model, tokenizer)\n\n    # Apply chat template\n    def format_dataset(example):\n        example[\"prompt\"] = tokenizer.apply_chat_template(\n            example[\"prompt\"], tokenize=False, add_generation_prompt=True\n        )\n        return example\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        # Load the dataset\n        dataset = build_helpfulness_dataset(script_args.llm_name, num_proc=training_args.dataset_num_proc)\n        dataset = dataset.map(format_dataset, batched=False, num_proc=training_args.dataset_num_proc)\n\n    accelerator = Accelerator()\n    embedding_model = AutoModel.from_pretrained(\n        \"nomic-ai/nomic-embed-text-v1.5\",\n        trust_remote_code=model_args.trust_remote_code,\n        safe_serialization=True,\n        torch_dtype=torch.bfloat16,\n        device_map=\"auto\",\n    )\n    embedding_model = accelerator.prepare_model(embedding_model)\n    embedding_tokenizer = AutoTokenizer.from_pretrained(\n        \"bert-base-uncased\", trust_remote_code=model_args.trust_remote_code\n    )\n    embedding_func = partial(\n        embed_prompt,\n        model=embedding_model,\n    )\n\n    # Initialize the BCO trainer\n    bco_trainer = BCOTrainer(\n        model,\n        ref_model,\n        args=training_args,\n        train_dataset=dataset[\"train\"],\n        eval_dataset=dataset[\"test\"],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_args),\n        embedding_func=embedding_func,\n        embedding_tokenizer=embedding_tokenizer,\n    )\n\n    # Train and push the model to the Hub\n    bco_trainer.train()\n    bco_trainer.save_model(training_args.output_dir)\n\n\n# flake8: noqa\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUsage:\n\npython examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-tldr-online-dpo \\\n    --per_device_train_batch_size 8 \\\n    --gradient_accumulation_steps 16 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0\n\nWith LoRA:\npython examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-6 \\\n    --output_dir pythia-1b-tldr-online-dpo \\\n    --per_device_train_batch_size 16 \\\n    --gradient_accumulation_steps 8 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --use_peft\n\"\"\"\n\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, GenerationConfig\nfrom trl import (\n    DPOScriptArguments,\n    ModelConfig,\n    OnlineDPOConfig,\n    OnlineDPOTrainer,\n    get_kbit_device_map,\n    get_peft_config,\n    get_quantization_config,\n    LogCompletionsCallback,\n)\n\nfrom trl.commands.cli_utils import TrlParser\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\nif __name__ == \"__main__\":\n    parser = TrlParser((DPOScriptArguments, OnlineDPOConfig, ModelConfig))\n    args, training_args, model_config = parser.parse_args_and_config()\n    args.gradient_checkpointing_kwargs = {\"use_reentrant\": True}\n\n    torch_dtype = (\n        model_config.torch_dtype\n        if model_config.torch_dtype in [\"auto\", None]\n        else getattr(torch, model_config.torch_dtype)\n    )\n    quantization_config = get_quantization_config(model_config)\n    model_kwargs = dict(\n        revision=model_config.model_revision,\n        attn_implementation=model_config.attn_implementation,\n        torch_dtype=torch_dtype,\n        use_cache=False if training_args.gradient_checkpointing else True,\n        device_map=get_kbit_device_map() if quantization_config is not None else None,\n        quantization_config=quantization_config,\n    )\n\n    model = AutoModelForCausalLM.from_pretrained(\n        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs\n    )\n\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path,\n        num_labels=1,\n        trust_remote_code=model_config.trust_remote_code,\n        **model_kwargs,\n    )\n\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n        **model_kwargs,\n    )\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    if tokenizer.pad_token_id is None:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    dataset = load_dataset(args.dataset_name)\n\n    trainer = OnlineDPOTrainer(\n        model=model,\n        reward_model=reward_model,\n        args=training_args,\n        train_dataset=dataset[args.dataset_train_split],\n        eval_dataset=dataset[args.dataset_test_split],\n        tokenizer=tokenizer,\n        peft_config=get_peft_config(model_config),\n    )\n    generation_config = GenerationConfig(\n        max_new_tokens=training_args.max_new_tokens, do_sample=True, temperature=training_args.temperature\n    )\n    completions_callback = LogCompletionsCallback(trainer, generation_config, num_prompts=8)\n    trainer.add_callback(completions_callback)\n    trainer.train()\n\n\n# flake8: noqa\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom trl.commands.cli_utils import init_zero_verbose\n\ninit_zero_verbose()\n\nimport copy\nimport json\nimport os\nimport sys\nimport pwd\nimport re\nimport time\nfrom threading import Thread\n\nimport torch\nfrom rich.console import Console\nfrom rich.live import Live\nfrom rich.markdown import Markdown\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer\n\nfrom trl.commands.cli_utils import ChatArguments, TrlParser, init_zero_verbose\nfrom trl.trainer.utils import get_quantization_config\n\n\nHELP_STRING = \"\"\"\\\n\n**TRL CHAT INTERFACE**\n\nThe chat interface is a simple tool to try out a chat model.\n\nBesides talking to the model there are several commands:\n- **clear**: clears the current conversation and start a new one\n- **example {NAME}**: load example named `{NAME}` from the config and use it as the user input\n- **set {SETTING_NAME}={SETTING_VALUE};**: change the system prompt or generation settings (multiple settings are separated by a ';').\n- **reset**: same as clear but also resets the generation configs to defaults if they have been changed by **set**\n- **save {SAVE_NAME} (optional)**: save the current chat and settings to file by default to `./chat_history/{MODEL_NAME}/chat_{DATETIME}.yaml` or `{SAVE_NAME}` if provided\n- **exit**: closes the interface\n\"\"\"\n\nSUPPORTED_GENERATION_KWARGS = [\n    \"max_new_tokens\",\n    \"do_sample\",\n    \"num_beams\",\n    \"temperature\",\n    \"top_p\",\n    \"top_k\",\n    \"repetition_penalty\",\n]\n\nSETTING_RE = r\"^set\\s+[A-Za-z\\s_]+=[A-Za-z\\d\\s.!\\\"#$%&'()*+,-/:<=>?@\\[\\]^_`{|}~]+(?:;\\s*[A-Za-z\\s_]+=[A-Za-z\\d\\s.!\\\"#$%&'()*+,-/:<=>?@\\[\\]^_`{|}~]+)*$\"\n\n\nclass RichInterface:\n    def __init__(self, model_name=None, user_name=None):\n        self._console = Console()\n        if model_name is None:\n            self.model_name = \"assistant\"\n        else:\n            self.model_name = model_name\n        if user_name is None:\n            self.user_name = \"user\"\n        else:\n            self.user_name = user_name\n\n    def stream_output(self, output_stream):\n        \"\"\"Stream output from a role.\"\"\"\n        # This method is originally from the FastChat CLI: https://github.com/lm-sys/FastChat/blob/main/fastchat/serve/cli.py\n        # Create a Live context for updating the console output\n        text = \"\"\n        self._console.print(f\"[bold blue]<{self.model_name}>:\")\n        with Live(console=self._console, refresh_per_second=4) as live:\n            # Read lines from the stream\n            for i, outputs in enumerate(output_stream):\n                if not outputs or i == 0:\n                    continue\n                text += outputs\n                # Render the accumulated text as Markdown\n                # NOTE: this is a workaround for the rendering \"unstandard markdown\"\n                #  in rich. The chatbots output treat \"\\n\" as a new line for\n                #  better compatibility with real-world text. However, rendering\n                #  in markdown would break the format. It is because standard markdown\n                #  treat a single \"\\n\" in normal text as a space.\n                #  Our workaround is adding two spaces at the end of each line.\n                #  This is not a perfect solution, as it would\n                #  introduce trailing spaces (only) in code block, but it works well\n                #  especially for console output, because in general the console does not\n                #  care about trailing spaces.\n                lines = []\n                for line in text.splitlines():\n                    lines.append(line)\n                    if line.startswith(\"```\"):\n                        # Code block marker - do not add trailing spaces, as it would\n                        #  break the syntax highlighting\n                        lines.append(\"\\n\")\n                    else:\n                        lines.append(\"  \\n\")\n                markdown = Markdown(\"\".join(lines).strip(), code_theme=\"github-dark\")\n                # Update the Live console output\n                live.update(markdown)\n        self._console.print()\n        return text\n\n    def input(self):\n        input = self._console.input(f\"[bold red]<{self.user_name}>:\\n\")\n        self._console.print()\n        return input\n\n    def clear(self):\n        self._console.clear()\n\n    def print_user_message(self, text):\n        self._console.print(f\"[bold red]<{self.user_name}>:[/ bold red]\\n{text}\")\n        self._console.print()\n\n    def print_green(self, text):\n        self._console.print(f\"[bold green]{text}\")\n        self._console.print()\n\n    def print_red(self, text):\n        self._console.print(f\"[bold red]{text}\")\n        self._console.print()\n\n    def print_help(self):\n        self._console.print(Markdown(HELP_STRING))\n        self._console.print()\n\n\ndef get_username():\n    return pwd.getpwuid(os.getuid())[0]\n\n\ndef create_default_filename(model_name):\n    time_str = time.strftime(\"%Y-%m-%d_%H-%M-%S\")\n    return f\"{model_name}/chat_{time_str}.json\"\n\n\ndef save_chat(chat, args, filename):\n    output_dict = {}\n    output_dict[\"settings\"] = vars(args)\n    output_dict[\"chat_history\"] = chat\n\n    folder = args.save_folder\n\n    if filename is None:\n        filename = create_default_filename(args.model_name_or_path)\n        filename = os.path.join(folder, filename)\n    os.makedirs(os.path.dirname(filename), exist_ok=True)\n\n    with open(filename, \"w\") as f:\n        json.dump(output_dict, f, indent=4)\n    return os.path.abspath(filename)\n\n\ndef clear_chat_history(system_prompt):\n    if system_prompt is None:\n        chat = []\n    else:\n        chat = [{\"role\": \"system\", \"content\": system_prompt}]\n    return chat\n\n\ndef parse_settings(user_input, current_args, interface):\n    settings = user_input[4:].strip().split(\";\")\n    settings = [(setting.split(\"=\")[0], setting[len(setting.split(\"=\")[0]) + 1 :]) for setting in settings]\n    settings = dict(settings)\n    error = False\n\n    for name in settings:\n        if hasattr(current_args, name):\n            try:\n                if isinstance(getattr(current_args, name), bool):\n                    if settings[name] == \"True\":\n                        settings[name] = True\n                    elif settings[name] == \"False\":\n                        settings[name] = False\n                    else:\n                        raise ValueError\n                else:\n                    settings[name] = type(getattr(current_args, name))(settings[name])\n            except ValueError:\n                interface.print_red(\n                    f\"Cannot cast setting {name} (={settings[name]}) to {type(getattr(current_args, name))}.\"\n                )\n        else:\n            interface.print_red(f\"There is no '{name}' setting.\")\n\n    if error:\n        interface.print_red(\"There was an issue parsing the settings. No settings have been changed.\")\n        return current_args, False\n    else:\n        for name in settings:\n            setattr(current_args, name, settings[name])\n            interface.print_green(f\"Set {name} to {settings[name]}.\")\n\n        time.sleep(1.5)  # so the user has time to read the changes\n        return current_args, True\n\n\ndef load_model_and_tokenizer(args):\n    tokenizer = AutoTokenizer.from_pretrained(\n        args.model_name_or_path,\n        revision=args.model_revision,\n        trust_remote_code=args.trust_remote_code,\n    )\n\n    torch_dtype = args.torch_dtype if args.torch_dtype in [\"auto\", None] else getattr(torch, args.torch_dtype)\n    quantization_config = get_quantization_config(args)\n    model_kwargs = dict(\n        revision=args.model_revision,\n        attn_implementation=args.attn_implementation,\n        torch_dtype=torch_dtype,\n        device_map=\"auto\",\n        quantization_config=quantization_config,\n    )\n    model = AutoModelForCausalLM.from_pretrained(\n        args.model_name_or_path, trust_remote_code=args.trust_remote_code, **model_kwargs\n    )\n\n    if getattr(model, \"hf_device_map\", None) is None:\n        model = model.to(args.device)\n\n    return model, tokenizer\n\n\ndef parse_eos_tokens(tokenizer, eos_tokens, eos_token_ids):\n    if tokenizer.pad_token_id is None:\n        pad_token_id = tokenizer.eos_token_id\n    else:\n        pad_token_id = tokenizer.pad_token_id\n\n    all_eos_token_ids = []\n\n    if eos_tokens is not None:\n        all_eos_token_ids.extend(tokenizer.convert_tokens_to_ids(eos_tokens.split(\",\")))\n\n    if eos_token_ids is not None:\n        all_eos_token_ids.extend([int(token_id) for token_id in eos_token_ids.split(\",\")])\n\n    if len(all_eos_token_ids) == 0:\n        all_eos_token_ids.append(tokenizer.eos_token_id)\n\n    return pad_token_id, all_eos_token_ids\n\n\ndef chat_cli():\n    parser = TrlParser(ChatArguments)\n\n    if \"--config\" not in sys.argv:\n        sys.argv.append(\"--config\")\n        sys.argv.append(os.path.join(os.path.dirname(__file__), \"config/default_chat_config.yaml\"))\n    args = parser.parse_args_and_config()[0]\n    if args.examples is None:\n        args.examples = {}\n\n    current_args = copy.deepcopy(args)\n\n    if args.user is None:\n        user = get_username()\n    else:\n        user = args.user\n\n    model, tokenizer = load_model_and_tokenizer(args)\n    generation_streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)\n\n    pad_token_id, eos_token_ids = parse_eos_tokens(tokenizer, args.eos_tokens, args.eos_token_ids)\n\n    interface = RichInterface(model_name=args.model_name_or_path, user_name=user)\n    interface.clear()\n    chat = clear_chat_history(current_args.system_prompt)\n    while True:\n        try:\n            user_input = interface.input()\n\n            if user_input == \"clear\":\n                chat = clear_chat_history(current_args.system_prompt)\n                interface.clear()\n                continue\n\n            if user_input == \"help\":\n                interface.print_help()\n                continue\n\n            if user_input == \"exit\":\n                break\n\n            if user_input == \"reset\":\n                interface.clear()\n                current_args = copy.deepcopy(args)\n                chat = clear_chat_history(current_args.system_prompt)\n                continue\n\n            if user_input.startswith(\"save\") and len(user_input.split()) < 2:\n                split_input = user_input.split()\n\n                if len(split_input) == 2:\n                    filename = split_input[1]\n                else:\n                    filename = None\n                filename = save_chat(chat, current_args, filename)\n                interface.print_green(f\"Chat saved in {filename}!\")\n                continue\n\n            if re.match(SETTING_RE, user_input):\n                current_args, success = parse_settings(user_input, current_args, interface)\n                if success:\n                    chat = []\n                    interface.clear()\n                    continue\n\n            if user_input.startswith(\"example\") and len(user_input.split()) == 2:\n                example_name = user_input.split()[1]\n                if example_name in current_args.examples:\n                    interface.clear()\n                    chat = []\n                    interface.print_user_message(current_args.examples[example_name][\"text\"])\n                    user_input = current_args.examples[example_name][\"text\"]\n                else:\n                    interface.print_red(\n                        f\"Example {example_name} not found in list of available examples: {list(current_args.examples.keys())}.\"\n                    )\n                    continue\n\n            chat.append({\"role\": \"user\", \"content\": user_input})\n\n            generation_kwargs = dict(\n                inputs=tokenizer.apply_chat_template(chat, return_tensors=\"pt\", add_generation_prompt=True).to(\n                    model.device\n                ),\n                streamer=generation_streamer,\n                max_new_tokens=current_args.max_new_tokens,\n                do_sample=current_args.do_sample,\n                num_beams=current_args.num_beams,\n                temperature=current_args.temperature,\n                top_k=current_args.top_k,\n                top_p=current_args.top_p,\n                repetition_penalty=current_args.repetition_penalty,\n                pad_token_id=pad_token_id,\n                eos_token_id=eos_token_ids,\n            )\n\n            thread = Thread(target=model.generate, kwargs=generation_kwargs)\n            thread.start()\n            model_output = interface.stream_output(generation_streamer)\n            thread.join()\n            chat.append({\"role\": \"assistant\", \"content\": model_output})\n\n        except KeyboardInterrupt:\n            break\n\n\nif __name__ == \"__main__\":\n    chat_cli()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nfrom datasets import load_dataset\nfrom transformers import HfArgumentParser\nfrom vllm import LLM, SamplingParams\n\nfrom trl import HfPairwiseJudge, OpenAIPairwiseJudge\n\n\n\"\"\"\nExamples:\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --num_examples 1000\nModel win rate: 31.40%\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-3.5-turbo-0125 --num_examples 1000\nModel win rate: 51.60%\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 51.20%\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --num_examples 1000\nModel win rate: 46.30%\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --judge_model gpt-3.5-turbo-0125 --num_examples 1000\nModel win rate: 52.50%\n\npython examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 63.00%\n\"\"\"\n\n\n@dataclass\nclass ScriptArguments:\n    model_name_or_path: str = field(metadata={\"help\": \"The model name or path to the model to evaluate.\"})\n    judge_model: str = field(\n        default=\"meta-llama/Meta-Llama-3-70B-Instruct\",\n        metadata={\n            \"help\": \"The model name or path to the model to use as a judge. E.g., 'gpt-3.5-turbo-0125', 'meta-llama/Meta-Llama-3-70B-Instruct'.\"\n        },\n    )\n    num_examples: Optional[int] = field(default=None, metadata={\"help\": \"The number of examples to evaluate.\"})\n\n\n# Parse the arguments\nparser = HfArgumentParser(ScriptArguments)\nargs = parser.parse_args_into_dataclasses()[0]\n\n# Load the dataset\ndataset = load_dataset(\"trl-lib/tldr\", split=\"validation\")\nif args.num_examples is not None:\n    dataset = dataset.select(range(args.num_examples))\n\n# Extract the prompts and reference completions\nprompts = dataset[\"prompt\"]\nreference_completions = dataset[\"completion\"]\n\n# Generate the model completions\nsampling_params = SamplingParams(temperature=0.0, top_p=0.95, max_tokens=200)  # very generous max token length\nllm = LLM(model=args.model_name_or_path, tensor_parallel_size=1)\noutputs = llm.generate(prompts, sampling_params)\nmodel_completions = [output.outputs[0].text.strip() for output in outputs]\n\n# Judge the outputs\nif \"gpt\" in args.judge_model:\n    judge = OpenAIPairwiseJudge(args.judge_model)\nelse:\n    judge = HfPairwiseJudge(args.judge_model)\n\ncompletions = [[c0, c1] for c0, c1 in zip(reference_completions, model_completions)]\nbest_idxs = judge.judge(prompts, completions)\nmodel_win_rate = best_idxs.count(1) / len(best_idxs)\nprint(f\"Model win rate: {model_win_rate*100:.2f}%\")\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport shutil\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    HfArgumentParser,\n)\n\nfrom trl import ModelConfig, PPOv2Config, PPOv2Trainer\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n\"\"\"\npython -i examples/scripts/ppo/ppo.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 64 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --missing_eos_penalty 1.0\n\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml \\\n    examples/scripts/ppo/ppo.py \\\n    --output_dir models/minimal/ppo \\\n    --num_ppo_epochs 1 \\\n    --num_mini_batches 1 \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 1 \\\n    --gradient_accumulation_steps 16 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path EleutherAI/pythia-1b-deduped \\\n    --reward_model_path EleutherAI/pythia-1b-deduped \\\n    --local_rollout_forward_batch_size 1 \\\n    --deepspeed3 \\\n    --missing_eos_penalty 1.0\n\"\"\"\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((PPOv2Config, ModelConfig))\n    training_args, model_config = parser.parse_args_into_dataclasses()\n    # remove output_dir if exists\n    shutil.rmtree(training_args.output_dir, ignore_errors=True)\n\n    ################\n    # Model & Tokenizer\n    ################\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    tokenizer.add_special_tokens({\"pad_token\": \"[PAD]\"})\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    value_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    ref_policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(\"trl-internal-testing/descriptiveness-sentiment-trl-style\", split=\"descriptiveness\")\n    eval_samples = 20\n    train_dataset = dataset.select(range(len(dataset) - eval_samples))\n    eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))\n    dataset_text_field = \"prompt\"\n\n    def prepare_dataset(dataset, tokenizer):\n        \"\"\"pre-tokenize the dataset before training; only collate during training\"\"\"\n\n        def tokenize(element):\n            outputs = tokenizer(\n                element[dataset_text_field],\n                padding=False,\n            )\n            return {\"input_ids\": outputs[\"input_ids\"]}\n\n        return dataset.map(\n            tokenize,\n            batched=True,\n            remove_columns=dataset.column_names,\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        train_dataset = prepare_dataset(train_dataset, tokenizer)\n        eval_dataset = prepare_dataset(eval_dataset, tokenizer)\n\n    ################\n    # Training\n    ################\n    trainer = PPOv2Trainer(\n        config=training_args,\n        tokenizer=tokenizer,\n        policy=policy,\n        ref_policy=ref_policy,\n        reward_model=reward_model,\n        value_model=value_model,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n    )\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n    if training_args.push_to_hub:\n        trainer.push_to_hub()\n    trainer.generate_completions()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport shutil\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    HfArgumentParser,\n)\n\nfrom trl import ModelConfig, PPOv2Config, PPOv2Trainer\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n\"\"\"\npython examples/scripts/ppo/ppo_tldr.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 1 \\\n    --gradient_accumulation_steps 64 \\\n    --total_episodes 30000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos \\\n    --response_length 53\n\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/ppo/ppo_tldr.py \\\n    --output_dir models/minimal/ppo_tldr \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 16 \\\n    --gradient_accumulation_steps 4 \\\n    --total_episodes 1000000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --local_rollout_forward_batch_size 16 \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos\n\"\"\"\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((PPOv2Config, ModelConfig))\n    training_args, model_config = parser.parse_args_into_dataclasses()\n    # remove output_dir if exists\n    shutil.rmtree(training_args.output_dir, ignore_errors=True)\n\n    ################\n    # Model & Tokenizer\n    ################\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    tokenizer.add_special_tokens({\"pad_token\": \"[PAD]\"})\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    value_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    ref_policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(\"trl-internal-testing/tldr-preference-sft-trl-style\")\n    train_dataset = dataset[\"train\"]\n    eval_dataset = dataset[\"validation\"]\n\n    def prepare_dataset(dataset, tokenizer):\n        \"\"\"pre-tokenize the dataset before training; only collate during training\"\"\"\n\n        def tokenize(element):\n            input_ids = tokenizer.apply_chat_template(\n                element[\"messages\"][:1],\n                padding=False,\n                add_generation_prompt=True,\n            )\n            return {\"input_ids\": input_ids, \"lengths\": len(input_ids)}\n\n        return dataset.map(\n            tokenize,\n            remove_columns=dataset.column_names,\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        train_dataset = prepare_dataset(train_dataset, tokenizer)\n        eval_dataset = prepare_dataset(eval_dataset, tokenizer)\n        # filtering\n        train_dataset = train_dataset.filter(lambda x: x[\"lengths\"] <= 512, num_proc=training_args.dataset_num_proc)\n        eval_dataset = eval_dataset.filter(lambda x: x[\"lengths\"] <= 512, num_proc=training_args.dataset_num_proc)\n\n    assert train_dataset[0][\"input_ids\"][-1] != tokenizer.eos_token_id, \"The last token should not be an EOS token\"\n    ################\n    # Training\n    ################\n    trainer = PPOv2Trainer(\n        config=training_args,\n        tokenizer=tokenizer,\n        policy=policy,\n        ref_policy=ref_policy,\n        reward_model=reward_model,\n        value_model=value_model,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n    )\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n    if training_args.push_to_hub:\n        trainer.push_to_hub()\n    trainer.generate_completions()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport shutil\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    HfArgumentParser,\n)\n\nfrom trl import ModelConfig\nfrom trl.trainer.rloo_trainer import RLOOConfig, RLOOTrainer\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n\"\"\"\npython -i examples/scripts/rloo/rloo.py \\\n    --learning_rate 3e-6 \\\n    --num_ppo_epochs 1 \\\n    --num_mini_batches 1 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 64 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --missing_eos_penalty 1.0\n\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml \\\n    examples/scripts/rloo/rloo.py \\\n    --output_dir models/minimal/rloo \\\n    --rloo_k 2 \\\n    --num_ppo_epochs 1 \\\n    --num_mini_batches 1 \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 1 \\\n    --gradient_accumulation_steps 16 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path EleutherAI/pythia-1b-deduped \\\n    --reward_model_path EleutherAI/pythia-1b-deduped \\\n    --local_rollout_forward_batch_size 1 \\\n    --deepspeed3 \\\n    --missing_eos_penalty 1.0\n\"\"\"\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((RLOOConfig, ModelConfig))\n    training_args, model_config = parser.parse_args_into_dataclasses()\n    # remove output_dir if exists\n    shutil.rmtree(training_args.output_dir, ignore_errors=True)\n\n    ################\n    # Model & Tokenizer\n    ################\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    tokenizer.add_special_tokens({\"pad_token\": \"[PAD]\"})\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    ref_policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(\"trl-internal-testing/descriptiveness-sentiment-trl-style\", split=\"descriptiveness\")\n    eval_samples = 20\n    train_dataset = dataset.select(range(len(dataset) - eval_samples))\n    eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))\n    dataset_text_field = \"prompt\"\n\n    def prepare_dataset(dataset, tokenizer):\n        \"\"\"pre-tokenize the dataset before training; only collate during training\"\"\"\n\n        def tokenize(element):\n            outputs = tokenizer(\n                element[dataset_text_field],\n                padding=False,\n            )\n            return {\"input_ids\": outputs[\"input_ids\"]}\n\n        return dataset.map(\n            tokenize,\n            batched=True,\n            remove_columns=dataset.column_names,\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        train_dataset = prepare_dataset(train_dataset, tokenizer)\n        eval_dataset = prepare_dataset(eval_dataset, tokenizer)\n\n    ################\n    # Training\n    ################\n    trainer = RLOOTrainer(\n        config=training_args,\n        tokenizer=tokenizer,\n        policy=policy,\n        ref_policy=ref_policy,\n        reward_model=reward_model,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n    )\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n    if training_args.push_to_hub:\n        trainer.push_to_hub()\n    trainer.generate_completions()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport shutil\n\nfrom accelerate import PartialState\nfrom datasets import load_dataset\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    HfArgumentParser,\n)\n\nfrom trl import ModelConfig\nfrom trl.trainer.rloo_trainer import RLOOConfig, RLOOTrainer\nfrom trl.trainer.utils import SIMPLE_CHAT_TEMPLATE\n\n\n\"\"\"\npython examples/scripts/rloo/rloo_tldr.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 1 \\\n    --gradient_accumulation_steps 64 \\\n    --total_episodes 30000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos \\\n    --response_length 53\n\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/rloo/rloo_tldr.py \\\n    --output_dir models/minimal/rloo_tldr \\\n    --num_ppo_epochs 1 \\\n    --num_mini_batches 1 \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 16 \\\n    --gradient_accumulation_steps 4 \\\n    --total_episodes 1000000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --local_rollout_forward_batch_size 16 \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos\n\"\"\"\n\n\nif __name__ == \"__main__\":\n    parser = HfArgumentParser((RLOOConfig, ModelConfig))\n    training_args, model_config = parser.parse_args_into_dataclasses()\n    # remove output_dir if exists\n    shutil.rmtree(training_args.output_dir, ignore_errors=True)\n\n    ################\n    # Model & Tokenizer\n    ################\n    tokenizer = AutoTokenizer.from_pretrained(\n        model_config.model_name_or_path,\n        padding_side=\"left\",\n        trust_remote_code=model_config.trust_remote_code,\n    )\n    tokenizer.add_special_tokens({\"pad_token\": \"[PAD]\"})\n    if tokenizer.chat_template is None:\n        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE\n    reward_model = AutoModelForSequenceClassification.from_pretrained(\n        training_args.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1\n    )\n    ref_policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    policy = AutoModelForCausalLM.from_pretrained(\n        training_args.sft_model_path, trust_remote_code=model_config.trust_remote_code\n    )\n    ################\n    # Dataset\n    ################\n    dataset = load_dataset(\"trl-internal-testing/tldr-preference-sft-trl-style\")\n    train_dataset = dataset[\"train\"]\n    eval_dataset = dataset[\"validation\"]\n\n    def prepare_dataset(dataset, tokenizer):\n        \"\"\"pre-tokenize the dataset before training; only collate during training\"\"\"\n\n        def tokenize(element):\n            input_ids = tokenizer.apply_chat_template(\n                element[\"messages\"][:1],\n                padding=False,\n                add_generation_prompt=True,\n            )\n            return {\"input_ids\": input_ids, \"lengths\": len(input_ids)}\n\n        return dataset.map(\n            tokenize,\n            remove_columns=dataset.column_names,\n            num_proc=training_args.dataset_num_proc,\n        )\n\n    # Compute that only on the main process for faster data processing.\n    # see: https://github.com/huggingface/trl/pull/1255\n    with PartialState().local_main_process_first():\n        train_dataset = prepare_dataset(train_dataset, tokenizer)\n        eval_dataset = prepare_dataset(eval_dataset, tokenizer)\n        # filtering\n        train_dataset = train_dataset.filter(lambda x: x[\"lengths\"] <= 512, num_proc=training_args.dataset_num_proc)\n        eval_dataset = eval_dataset.filter(lambda x: x[\"lengths\"] <= 512, num_proc=training_args.dataset_num_proc)\n\n    assert train_dataset[0][\"input_ids\"][-1] != tokenizer.eos_token_id, \"The last token should not be an EOS token\"\n    ################\n    # Training\n    ################\n    trainer = RLOOTrainer(\n        config=training_args,\n        tokenizer=tokenizer,\n        policy=policy,\n        ref_policy=ref_policy,\n        reward_model=reward_model,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n    )\n    trainer.train()\n    trainer.save_model(training_args.output_dir)\n    if training_args.push_to_hub:\n        trainer.push_to_hub()\n    trainer.generate_completions()\n\n\n# Notebooks\n\nThis directory contains a collection of Jupyter notebooks that demonstrate how to use the TRL library in different applications.\n\n- [`best_of_n.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/best_of_n.ipynb): This notebook demonstrates how to use the \"Best of N\" sampling strategy using TRL when fine-tuning your model with PPO.\n- [`gpt2-sentiment.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-sentiment.ipynb): This notebook demonstrates how to reproduce the GPT2 imdb sentiment tuning example on a jupyter notebook.\n- [`gpt2-control.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-sentiment-control.ipynb): This notebook demonstrates how to reproduce the GPT2 sentiment control example on a jupyter notebook.\n\n\n# Training FAQ\n\n## What Metrics Should I Look at?\n\nWhen performing classical supervised fine-tuning of language models, the loss (especially the validation loss) serves as a good indicator of the training progress. However, in Reinforcement Learning (RL), the loss becomes less informative about the model's performance, and its value may fluctuate while the actual performance improves.\n\nTo address this, we recommend focusing on two key metrics first:\n\n**Mean Reward**: The primary goal is to maximize the reward achieved by the model during RL training.\n**Objective KL Divergence**: KL divergence (Kullback-Leibler divergence) measures the dissimilarity between two probability distributions. In the context of RL training, we use it to quantify the difference between the current model and a reference model. Ideally, we want to keep the KL divergence between 0 and 10 to ensure the model's generated text remains close to what the reference model produces.\n\nHowever, there are more metrics that can be useful for debugging, checkout the [logging section](logging).\n\n## Why Do We Use a Reference Model, and What's the Purpose of KL Divergence?\n\nWhen training RL models, optimizing solely for reward may lead to unexpected behaviors, where the model exploits the environment in ways that don't align with good language generation. In the case of RLHF, we use a reward model trained to predict whether a generated text is highly ranked by humans.\n\nHowever, the RL model being optimized against the reward model may learn patterns that yield high reward but do not represent good language. This can result in extreme cases where the model generates texts with excessive exclamation marks or emojis to maximize the reward. In some worst-case scenarios, the model may generate patterns completely unrelated to natural language yet receive high rewards, similar to adversarial attacks.\n\n<div style=\"text-align: center\">\n<img src=\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/kl-example.png\">\n<p style=\"text-align: center;\"> <b>Figure:</b> Samples without a KL penalty from <a href=\"https://huggingface.co/papers/1909.08593\">https://huggingface.co/papers/1909.08593</a>. </p>\n</div>\n\nTo address this issue, we add a penalty to the reward function based on the KL divergence between the current model and the reference model. By doing this, we encourage the model to stay close to what the reference model generates.\n\n## What Is the Concern with Negative KL Divergence?\n\nIf you generate text by purely sampling from the model distribution things work fine in general. But when you use the `generate` method there are a few caveats because it does not always purely sample depending on the settings which can cause KL-divergence to go negative. Essentially when the active model achieves `log_p_token_active < log_p_token_ref` we get negative KL-div. This can happen in a several cases:\n\n- **top-k sampling**: the model can smooth out the probability distribution causing the top-k tokens having a smaller probability than those of the reference model but they still are selected\n- **min_length**: this ignores the EOS token until `min_length` is reached. thus the model can assign a very low log prob to the EOS token and very high probs to all others until min_length is reached\n\nThese are just a few examples. Why is negative KL an issue? The total reward `R` is computed `R = r - beta * KL` so if the model can learn how to drive KL-divergence negative it effectively gets a positive reward. In many cases it can be much easier to exploit such a bug in the generation than actually learning the reward function. In addition the KL can become arbitrarily small thus the actual reward can be very small compared to it.\n\nSo how should you generate text for PPO training? Let's have a look!\n\n## How to generate text for training?\n\nIn order to avoid the KL issues described above we recommend to use the following settings:\n\n```python\ngeneration_kwargs = {\n    \"min_length\": -1, # don't ignore the EOS token (see above)\n    \"top_k\": 0.0, # no top-k sampling\n    \"top_p\": 1.0, # no nucleus sampling\n    \"do_sample\": True, # yes, we want to sample\n    \"pad_token_id\": tokenizer.eos_token_id, # most decoder models don't have a padding token - use EOS token instead\n    \"max_new_tokens\": 32, # specify how many tokens you want to generate at most\n}\n```\n\nWith these settings we usually don't encounter any issues. You can also experiments with other settings but if you encounter issues with negative KL-divergence try to go back to these and see if they persist.\n\n## How can debug your own use-case?\n\nDebugging the RL pipeline can be challenging due to its complexity. Here are some tips and suggestions to make the process easier:\n\n- **Start from a working example**: Begin with a working example from the trl repository and gradually modify it to fit your specific use-case. Changing everything at once can make it difficult to identify the source of potential issues. For example, you can start by replacing the model in the example and once you figure out the best hyperparameters try to switch to your dataset and reward model. If you change everything at once you won't know where a potential problem comes from.\n- **Start small, scale later**: Training large models can be very slow and take several hours or days until you see any improvement. For debugging this is not a convenient timescale so try to use small model variants during the development phase and scale up once that works. That being said you sometimes have to be careful as small models might not have the capacity to solve a complicated task either.\n- **Start simple**: Try to start with a minimal example and build complexity from there. Your use-case might require for example a complicated reward function consisting of many different rewards - try to use one signal first and see if you can optimize that and then add more complexity after that.\n- **Inspect the generations**: It's always a good idea to inspect what the model is generating. Maybe there is a bug in your post-processing or your prompt. Due to bad settings you might cut-off generations too soon. These things are very hard to see on the metrics but very obvious if you look at the generations.\n- **Inspect the reward model**: If you reward is not improving over time maybe there's an issue with the reward model. You can look at extreme cases to see if it does what it should: e.g. in the sentiment case you can check if simple positive and negative examples really get different rewards. And you can look at the distribution of your dataset. Finally, maybe the reward is dominated by the query which the model can't affect so you might need to normalize this (e.g. reward of query+response minus reward of the query).\n\nThese are just a few tips that we find helpful - if you have more useful tricks feel free to open a PR to add them as well!\n\n\n# Nash-MD Trainer\n\n## Overview\n\nNash-MD was proposed in the paper [Nash Learning from Human Feedback](https://huggingface.co/papers/2312.00886) by Rémi Munos, [Michal Valko](https://huggingface.co/misovalko), Daniele Calandriello, Mohammad Gheshlaghi Azar, Mark Rowland, Daniel Guo, Yunhao Tang, Matthieu Geist, Thomas Mésnard, and Andrea Michi. \n\nThe abstract from the paper is the following:\n\n> Reinforcement learning from human feedback (RLHF) has emerged as the main paradigm for aligning large language models (LLMs) with human preferences. Typically, RLHF involves the initial step of learning a reward model from human feedback, often expressed as preferences between pairs of text generations produced by a pre-trained LLM. Subsequently, the LLM's policy is fine-tuned by optimizing it to maximize the reward model through a reinforcement learning algorithm. However, an inherent limitation of current reward models is their inability to fully represent the richness of human preferences and their dependency on the sampling distribution. In this study, we introduce an alternative pipeline for the fine-tuning of LLMs using pairwise human feedback. Our approach entails the initial learning of a preference model, which is conditioned on two inputs given a prompt, followed by the pursuit of a policy that consistently generates responses preferred over those generated by any competing policy, thus defining the Nash equilibrium of this preference model. We term this approach Nash learning from human feedback (NLHF). In the context of a tabular policy representation, we present a novel algorithmic solution, Nash-MD, founded on the principles of mirror descent. This algorithm produces a sequence of policies, with the last iteration converging to the regularized Nash equilibrium. Additionally, we explore parametric representations of policies and introduce gradient descent algorithms for deep-learning architectures. To demonstrate the effectiveness of our approach, we present experimental results involving the fine-tuning of a LLM for a text summarization task. We believe NLHF offers a compelling avenue for preference learning and policy optimization with the potential of advancing the field of aligning LLMs with human preferences.\n\nThis post-training method was contributed by [Kashif Rasul](https://huggingface.co/kashif) and [Daniil Tiapkin](https://huggingface.co/dtiapkin), [Pierre Ménard](https://huggingface.co/menardprr), Daniele Calandriello and [Quentin Gallouédec](https://huggingface.co/qgallouedec). \n\n## Quick start\n\nThis example demonstrates how to train a model using the Nash-MD method. We use the [Qwen 0.5B model](https://huggingface.co/Qwen/Qwen2-0.5B-Instruct) as the base model and the [Qwen 0.5B reward model](https://huggingface.co/trl-lib/Qwen2-0.5B-Reward) as the reward model. We use the prompts from the [UltraFeedback dataset](https://huggingface.co/datasets/openbmb/UltraFeedback). You can view the prompts in the dataset here:\n\n<iframe\n  src=\"https://huggingface.co/datasets/trl-lib/ultrafeedback-prompt/embed/viewer/default/train?row=0\"\n  frameborder=\"0\"\n  width=\"100%\"\n  height=\"560px\"\n></iframe>\n\nBelow is the script to train the model:\n\n```python\n# train_nash_md.py\nfrom datasets import load_dataset\nfrom trl import NashMDConfig, NashMDTrainer\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\n\nmodel = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\nreward_model = AutoModelForSequenceClassification.from_pretrained(\"trl-lib/Qwen2-0.5B-Reward\", num_labels=1)\ntrain_dataset = load_dataset(\"trl-lib/ultrafeedback-prompt\", split=\"train\")\n\ntraining_args = NashMDConfig(output_dir=\"nash-md-qwen2\", logging_steps=10)\ntrainer = NashMDTrainer(\n    model=model,\n    reward_model=reward_model,\n    args=training_args,\n    tokenizer=tokenizer,\n    train_dataset=train_dataset,\n)\ntrainer.train()\n```\n\nExecute the script using the following command:\n\n```bash\naccelerate launch train_nash_md.py\n```\n\n## Expected dataset format\n\nNash-MD requires a [prompt-only dataset](dataset_format#preference). The [`NashMDTrainer`] supports both [conversational](dataset_format#conversational-dataset-format) and [standard](dataset_format#standard-dataset-format) dataset format. When provided with a conversational dataset, the trainer will automatically apply the chat template to the dataset.\n\n## Usage tips\n\n### ⚠️ Use the same chat template\n\nMake sure that the SFT model and reward model use the _same_ chat template. Otherwise, you may find the model completions are scored incorrectly during training.\n\n### Encourage EOS token generation\n\nWe can want the model to generate completion within a given length. During the learning, the model will generate completion up to the maximum completion length specified in the `max_new_tokens` argument of [`NashMDConfig`]. I you want to penalize for not generating an EOS token before the maximum completion length, you can use the `missing_eos_penalty` argument of [`NashMDConfig`]:\n\n```python\ntraining_args = NashMDConfig(..., max_new_tokens=128, missing_eos_penalty=1.0)\n```\n\n### Logging Completions\n\nTo better understand your model’s behavior during training, you can log sample completions periodically using the [`LogCompletionsCallback`].\n\n```python\ntrainer = NashMDTrainer(..., eval_dataset=eval_dataset)\ncompletions_callback = LogCompletionsCallback(trainer, num_prompts=8)\ntrainer.add_callback(completions_callback)\n```\n\nThis callback logs the model's generated completions directly to Weights & Biases.\n\n![Logged Completions](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/wandb_completions.png)\n\n## Example script\n\nWe provide an example script to train a model using the Nash-MD method. The script is available in [`examples/scripts/nash_md.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/nash_md.py)\n\nTo test the Nash-MD script with the [Pythia 14M model](https://huggingface.co/EleutherAI/pythia-14m) on the TL;DR summarization task, run the following command:\n\n```bash\npython examples/scripts/nash_md.py \\\n    --model_name_or_path EleutherAI/pythia-14m  \\\n    --reward_model_path EleutherAI/pythia-14m \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-14m-tldr-nash-md \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 32 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 64 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --push_to_hub\n```\n\n## Logged metrics\n\nThe logged metrics are as follows:\n\n* `loss/kl`: The mean KL divergence between the model and reference data.\n* `objective/entropy`: The mean entropy of the model and reference data.\n* `loss/score`: The mean reinforce score loss.\n* `rewards/chosen`: The mean scores (according to the reward model) of the model completions.\n* `rewards/rejected`: The mean scores (according to the reward model) of the mixture completions.\n* `rewards/accuracies`: The accuracies of the Nash-MD's implicit reward model.\n* `rewards/margins`: The mean reward margin (according to reward model) between the chosen and mixture completions.\n* `logps/chosen`: The mean log probabilities of the chosen completions.\n* `logps/rejected`: The mean log probabilities of the reference completions.\n* `val/model_contain_eos_token`: The amount of times the model's output contains the eos token.\n* `val/ref_contain_eos_token`: The amount of times the mixture's output contains the eos token.\n* `beta`: The parameter that controls the weight of the loss term representing the deviation from the reference model. Typically fixed, but can be made dynamic by passing a list to [`NashMDConfig`].\n* `mixture_coef`: Logit mixture coefficient for the model and reference model. Typically fixed, but can be made dynamic by passing a list to [`NashMDConfig`].\n\n## NashMDTrainer\n\n[[autodoc]] NashMDTrainer\n\n## NashMDConfig\n\n[[autodoc]] NashMDConfig\n\n\n# Generalized Knowledge Distillation Trainer\n\n## Overview\n\nGeneralized Knowledge Distillation (GKD) was proposed in [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://huggingface.co/papers/2306.13649) by Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. \n\nThe abstract from the paper is the following:\n\n> Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher's distribution. Furthermore, GKD facilitates the seamless integration of distillation with RL fine-tuning (RLHF). We demonstrate the efficacy of GKD for distilling auto-regressive language models on summarization, translation, and arithmetic reasoning tasks, and task-agnostic distillation for instruction-tuning.\n\n\nThe key aspects of GKD are:\n1. It addresses the train-inference distribution mismatch in auto-regressive sequence models by training the student model on its self-generated output sequences.\n2. GKD allows flexibility in choosing different divergence measures between student and teacher models via the generalized Jensen-Shannon Divergence (JSD), which can be useful when the student lacks the capacity to fully mimic the teacher.\n\nThis post-training method was contributed by [Kashif Rasul](https://huggingface.co/kashif) and [Lewis Tunstall](https://huggingface.co/lewtun).\n\n## Usage tips\n\nThe GKD Trainer is a wrapper around the [`SFTTrainer`] class that takes in a teacher model argument. It needs two parameters to be set via the [`GKDConfig`] namely:\n* `lmbda`:  controls the student data fraction, i.e., the proportion of on-policy student-generated outputs. When `lmbda=0.0`, the loss reduces to supervised JSD where the student is trained with the token-level probabilities of the teacher. When `lmbda=1.0`, the loss reduces to on-policy JSD, where the student generates output sequences and token-specific feedback on these sequences from the teacher. For values in between [0, 1] it is random between the two based on the `lmbda` value for each batch.\n* `beta`: controls the interpolation in the generalized Jensen-Shannon Divergence.  When `beta=0.0` the loss approximates forward KL divergence, while for `beta=1.0` the loss approximates reverse KL divergence. For values in between [0, 1] it interpolates between the two.\n\nThe authors find that on-policy data (high `lmbda`) performs better and the optimal `beta` varied depending on the task and evaluation method.\n\n> [!WARNING]\n> Make sure that `attn_implementation=\"flash_attention_2\"` when training [Gemma models](https://huggingface.co/models?other=gemma2). Otherwise you will encounter NaNs in the logits due to the [soft capping technique](https://huggingface.co/blog/gemma2#soft-capping-and-attention-implementations) adopted by this architecture.\n\nThe basic API is as follows:\n\n```python\nfrom datasets import Dataset\nfrom trl import GKDConfig, GKDTrainer\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoTokenizer,\n)\n\nNUM_DUMMY_SAMPLES = 100\n\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\n# The model to optimise\nmodel = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\n# The teacher model to calculate the KL divergence against\nteacher_model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen2-1.5B-Instruct\")\n\ntrain_dataset = Dataset.from_dict(\n    {\n        \"messages\": [\n            [\n                {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n                {\"role\": \"assistant\", \"content\": \"I'm great thanks\"},\n            ]\n        ]\n        * NUM_DUMMY_SAMPLES\n    }\n)\neval_dataset = Dataset.from_dict(\n    {\n        \"messages\": [\n            [\n                {\"role\": \"user\", \"content\": \"What colour is the sky?\"},\n                {\"role\": \"assistant\", \"content\": \"The sky is blue\"},\n            ]\n        ]\n        * NUM_DUMMY_SAMPLES\n    }\n)\n\ntraining_args = GKDConfig(output_dir=\"gkd-model\", per_device_train_batch_size=1)\ntrainer = GKDTrainer(\n    model=model,\n    teacher_model=teacher_model,\n    args=training_args,\n    tokenizer=tokenizer,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n)\ntrainer.train()\n```\n\n### Expected dataset format\n\nThe dataset should be formatted as a list of \"messages\" where each message is a list of dictionaries with the following keys:\n* `role`: either `system`, `assistant` or `user`\n* `content`: the message content\n\n\n## GKDTrainer\n\n[[autodoc]] GKDTrainer\n\n## GKDConfig\n\n[[autodoc]] GKDConfig\n\n\n# Use model after training\n\nOnce you have trained a model using either the SFTTrainer, PPOTrainer, or DPOTrainer, you will have a fine-tuned model that can be used for text generation. In this section, we'll walk through the process of loading the fine-tuned model and generating text. If you need to run an inference server with the trained model, you can explore libraries such as [`text-generation-inference`](https://github.com/huggingface/text-generation-inference).\n\n## Load and Generate\n\nIf you have fine-tuned a model fully, meaning without the use of PEFT you can simply load it like any other language model in transformers. E.g. the value head that was trained during the PPO training is no longer needed and if you load the model with the original transformer class it will be ignored:\n\n```python\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nmodel_name_or_path = \"kashif/stack-llama-2\" #path/to/your/model/or/name/on/hub\ndevice = \"cpu\" # or \"cuda\" if you have a GPU\n\nmodel = AutoModelForCausalLM.from_pretrained(model_name_or_path).to(device)\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n\ninputs = tokenizer.encode(\"This movie was really\", return_tensors=\"pt\").to(device)\noutputs = model.generate(inputs)\nprint(tokenizer.decode(outputs[0]))\n```\n\nAlternatively you can also use the pipeline:\n\n```python\nfrom transformers import pipeline\n\nmodel_name_or_path = \"kashif/stack-llama-2\" #path/to/your/model/or/name/on/hub\npipe = pipeline(\"text-generation\", model=model_name_or_path)\nprint(pipe(\"This movie was really\")[0][\"generated_text\"])\n```\n\n## Use Adapters PEFT\n\n```python\nfrom peft import PeftConfig, PeftModel\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nbase_model_name = \"kashif/stack-llama-2\" #path/to/your/model/or/name/on/hub\"\nadapter_model_name = \"path/to/my/adapter\"\n\nmodel = AutoModelForCausalLM.from_pretrained(base_model_name)\nmodel = PeftModel.from_pretrained(model, adapter_model_name)\n\ntokenizer = AutoTokenizer.from_pretrained(base_model_name)\n```\n\nYou can also merge the adapters into the base model so you can use the model like a normal transformers model, however the checkpoint will be significantly bigger:\n\n```python\nmodel = AutoModelForCausalLM.from_pretrained(base_model_name)\nmodel = PeftModel.from_pretrained(model, adapter_model_name)\n\nmodel = model.merge_and_unload()\nmodel.save_pretrained(\"merged_adapters\")\n```\n\nOnce you have the model loaded and either merged the adapters or keep them separately on top you can run generation as with a normal model outlined above.\n\n\n# PPOv2 Trainer\n\nTRL supports training LLMs with [Proximal Policy Optimization (PPO)](https://huggingface.co/papers/1707.06347).\n\nReferences:\n- [Fine-Tuning Language Models from Human Preferences](https://github.com/openai/lm-human-preferences)\n- [Learning to Summarize from Human Feedback](https://github.com/openai/summarize-from-feedback)\n- [The N Implementation Details of RLHF with PPO](https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo)\n- [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031)\n\n## Get started\n\nTo just run a PPO script to make sure the trainer can run, you can run the following command to train a PPO model with a dummy reward model.\n\n```bash\npython examples/scripts/ppo/ppo.py \\\n    --learning_rate 3e-6 \\\n    --num_ppo_epochs 1 \\\n    --num_mini_batches 1 \\\n    --output_dir models/minimal/ppo \\\n    --per_device_train_batch_size 64 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --missing_eos_penalty 1.0\n```\n\n\n## Explanation of the logged metrics\n\nThe logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35)\n\n* `eps`: Tracks the number of episodes per second.\n* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current policy and reference policy.\n* `objective/entropy`: The mean entropy of the policy, indicating the randomness of the actions chosen by the policy.\n* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.\n* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.\n* `objective/scores`: The mean scores returned by the reward model / environment.\n* `policy/approxkl_avg`: The average approximate KL divergence between consecutive PPO policies. Note that this is not the same as `objective/kl`.\n* `policy/clipfrac_avg`: The average fraction of policy updates that are clipped, indicating how often the policy updates are constrained to prevent large changes.\n* `loss/policy_avg`: The average policy loss, indicating how well the policy is performing.\n* `loss/value_avg`: The average value loss, indicating the difference between the predicted value and the actual reward.\n* `val/clipfrac_avg`: The average fraction of value function updates that are clipped, similar to policy/clipfrac_avg but for the value function.\n* `policy/entropy_avg`: The average entropy of the policy during training, indicating how diverse the policy's actions are.\n* `val/ratio`: The mean ratio of the current policy probability to the old policy probability, providing a measure of how much the policy has changed.\n* `val/ratio_var`: The variance of the `val/ratio`, indicating the variability in policy changes.\n* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.\n* `lr`: lr: The current learning rate used by the optimizer.\n* `episode`: episode: The current global step or episode count in the training process.\n\n\n## Cookbook\n\n* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.\n* Debugging TIP: `val/ratio`: this number should float around 1.0, and it gets clipped by `--cliprange 0.2` with PPO's surrogate loss. So if this `ratio` is too high like 2.0 or 1000.0 or too small like 0.1, it means the updates between consecutive policies are too drastic. You should try undertand why this is happening and try to fix it.\n* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.\n* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.\n* Usage TIP: We recommend to use the \"EOS trick\" via `--missing_eos_penalty`, which subtracts a static scalar penalty from the score of completions that do not end with an EOS token. This can help the model learn to generate more coherent completions.\n\n\n## What is my model doing exactly?\n\nTo help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif?download=true)\n\n\nIn the logs the sampled generations look like \n\n```\n┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓\n┃ query                           ┃ model response                  ┃ score    ┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩\n│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │\n│                                 │ I don't know how to get rid of  │          │\n│ TITLE: How do you get someone   │ those feelings. I'm             │          │\n│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │\n│                                 │                                 │          │\n│ POST: Hi,                       │                                 │          │\n│ I'm 22, and I have been with my │                                 │          │\n│ girlfriend for 5 years now. We  │                                 │          │\n│ recently moved together. We've  │                                 │          │\n│ always loved each other         │                                 │          │\n│ intensely.                      │                                 │          │\n│                                 │                                 │          │\n│ Problem, I recently started to  │                                 │          │\n│ have feelings for an other      │                                 │          │\n│ person (a friend). This person  │                                 │          │\n│ has had a boyfriend for now 3   │                                 │          │\n│ years, and has absolutely no    │                                 │          │\n│ ideas. Those feelings were so   │                                 │          │\n│ strong, it was hard to hide     │                                 │          │\n│ them. After 2 months of me      │                                 │          │\n│ being distant and really sad,   │                                 │          │\n│ my girlfriend forced me to say  │                                 │          │\n│ what was bothering me. I'm not  │                                 │          │\n│ a good liar, and now she knows. │                                 │          │\n│                                 │                                 │          │\n│ We decided to give us a week    │                                 │          │\n│ alone, I went to my parents.    │                                 │          │\n│                                 │                                 │          │\n│ Now, I'm completely lost. I     │                                 │          │\n│ keep on thinking about this     │                                 │          │\n│ person, and I hate that. I      │                                 │          │\n│ would like for those feelings   │                                 │          │\n│ to go away, to leave me alone.  │                                 │          │\n│ But I can't.                    │                                 │          │\n│                                 │                                 │          │\n│ What do I do? It's been 3       │                                 │          │\n│ months now, and I'm just        │                                 │          │\n│ desperate.                      │                                 │          │\n│                                 │                                 │          │\n│ TL;DR:                          │                                 │          │\n├─────────────────────────────────┼─────────────────────────────────┼──────────┤\n│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │\n│                                 │ TV. I blasted Gangnam Style on  │          │\n│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │\n│ with a loud TV.                 │ up as high as it could          │          │\n│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │\n│ POST: She was in her living     │                                 │          │\n│ room, watching TV. This was at  │                                 │          │\n│ about 8:30 in the morning, and  │                                 │          │\n│ she was exercising. She turned  │                                 │          │\n│ the TV up extra loud to hear it │                                 │          │\n│ over her excercycle, and woke   │                                 │          │\n│ me up. I went in there asking   │                                 │          │\n│ for her to turn it down. She    │                                 │          │\n│ said she didn't have to; I      │                                 │          │\n│ explained that I always used    │                                 │          │\n│ headphones so she didn't have   │                                 │          │\n│ to deal with my noise and that  │                                 │          │\n│ she should give me a little     │                                 │          │\n│ more respect, given that I paid │                                 │          │\n│ rent at the time.               │                                 │          │\n│                                 │                                 │          │\n│ She disagreed. I went back to   │                                 │          │\n│ my room, rather pissed off at   │                                 │          │\n│ the lack of equality. I had no  │                                 │          │\n│ lock on my door; but I had a    │                                 │          │\n│ dresser right next to it, so I  │                                 │          │\n│ pulled one of the drawers out   │                                 │          │\n│ enough so that it caused the    │                                 │          │\n│ door to not be openable. Then,  │                                 │          │\n│ I turned my speakers up really  │                                 │          │\n│ loud and blasted Gangnam Style  │                                 │          │\n│ on repeat, with the bass        │                                 │          │\n│ cranked up as high as it could  │                                 │          │\n│ go.                             │                                 │          │\n│                                 │                                 │          │\n│ If you hate Gangnam Style for   │                                 │          │\n│ being overplayed, you will see  │                                 │          │\n│ why I chose that particular     │                                 │          │\n│ song. I personally don't mind   │                                 │          │\n│ it. But here's the thing about  │                                 │          │\n│ my bass; it vibrates the walls, │                                 │          │\n│ making one hell of a lot of     │                                 │          │\n│ noise. Needless to say, my mom  │                                 │          │\n│ was not pleased and shut off    │                                 │          │\n│ the internet. But it was oh so  │                                 │          │\n│ worth it.                       │                                 │          │\n│                                 │                                 │          │\n│ TL;DR:                          │                                 │          │\n└─────────────────────────────────┴─────────────────────────────────┴──────────┘\n```\n\n## Implementation details\n\nThis PPOv2 implementation is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).\n\n## Benchmark experiments\n\nTo validate the PPO implementation works, we ran experiment on the 1B model. Here are the command we used to run the experiment. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).\n\n```\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/ppo/ppo_tldr.py \\\n    --output_dir models/minimal/ppo_tldr \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 16 \\\n    --gradient_accumulation_steps 4 \\\n    --total_episodes 1000000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --local_rollout_forward_batch_size 16 \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos\n```\n\nCheckpoints and experiment tracking are available at:\n\n- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/ppo_tldr)\n- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/dd2o3g35)\n\nTo evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.\nFor more information on how to use judges, see [Judges](judges).\n\n```bash\n$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 33.00%\n$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 64.70%\n```\n\nThe PPO checkpoint gets a 64.7% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the PPO training is working as intended.\n\nMetrics:\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/ppov2.png)\n\n\n```bash\n# pip install openrlbenchmark==0.2.1a5\n# see https://github.com/openrlbenchmark/openrlbenchmark#get-started for documentation\n# to use it, change `?we=huggingface&wpn=trl` to your own project and `?tag=pr-1540` to your own tag\npython -m openrlbenchmark.rlops_multi_metrics \\\n    --filters '?we=huggingface&wpn=trl&xaxis=train/episode&ceik=output_dir&cen=sft_model_path&metrics=train/objective/rlhf_reward&metrics=train/objective/scores&metrics=train/objective/kl&metrics=train/objective/non_score_reward&metrics=train/objective/entropy&metrics=train/policy/approxkl_avg&metrics=train/policy/clipfrac_avg&metrics=train/loss/policy_avg&metrics=train/loss/value_avg&metrics=train/val/clipfrac_avg&metrics=train/policy/entropy_avg&metrics=train/val/ratio&metrics=train/val/ratio_var&metrics=train/val/num_eos_tokens&metrics=train/lr&metrics=train/eps' \\\n        \"cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr?tag=pr-1540\" \\\n    --env-ids models/minimal/ppo_tldr \\\n    --pc.ncols 4 \\\n    --pc.ncols-legend 1 \\\n    --pc.xlabel \"Episode\" \\\n    --output-filename benchmark/trl/pr-1540/ppov2 \\\n    --scan-history\n```\n\n## PPOv2Trainer\n\n[[autodoc]] PPOv2Trainer\n\n## PPOv2Config\n\n[[autodoc]] PPOv2Config\n\n# Examples\n\n\n## Introduction\n\nThe examples should work in any of the following settings (with the same script):\n   - single GPU\n   - multi GPUS (using PyTorch distributed mode)\n   - multi GPUS (using DeepSpeed ZeRO-Offload stages 1, 2, & 3)\n   - fp16 (mixed-precision), fp32 (normal precision), or bf16 (bfloat16 precision)\n\nTo run it in each of these various modes, first initialize the accelerate\nconfiguration with `accelerate config`\n\n**NOTE to train with a 4-bit or 8-bit model**, please run\n\n```bash\npip install --upgrade trl[quantization]\n```\n\n\n## Accelerate Config\nFor all the examples, you'll need to generate a 🤗 Accelerate config file with:\n\n```shell\naccelerate config # will prompt you to define the training configuration\n```\n\nThen, it is encouraged to launch jobs with `accelerate launch`!\n\n\n# Maintained Examples\n\n\n\n| File                                                                                                                          | Description                                                                                                                                                                                                                                                                                                     |\n| ----------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [`examples/scripts/alignprop.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/alignprop.py)                 | This script shows how to use the [`AlignPropTrainer`] to fine-tune a diffusion model.                                                                                                                                                                                                                           |\n| [`examples/scripts/bco.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/bco.py)                             | This script shows how to use the [`KTOTrainer`] with the BCO loss to fine-tune a model to increase instruction-following, truthfulness, honesty and helpfulness using the [openbmb/UltraFeedback](https://huggingface.co/datasets/openbmb/UltraFeedback) dataset.                                               |\n| [`examples/scripts/chat.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/chat.py)                           | This script allows you to load and use a model as a chatbot.                                                                                                                                                                                                                                                    |\n| [`examples/scripts/cpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/cpo.py)                             | This script shows how to use the [`CPOTrainer`] to fine-tune a model to increase helpfulness and harmlessness using the [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset.                                                                                                         |\n| [`examples/scripts/ddpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/ddpo.py)                           | This script shows how to use the [`DDPOTrainer`] to fine-tune a stable diffusion model using reinforcement learning.                                                                                                                                                                                            |\n| [`examples/scripts/dpo_visual.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/dpo_visual.py)               | This script shows how to use the [`DPOTrainer`] to fine-tune a Vision Language Model to reduce hallucinations using the [openbmb/RLAIF-V-Dataset](https://huggingface.co/datasets/openbmb/RLAIF-V-Dataset) dataset.                                                                                             |\n| [`examples/scripts/dpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/dpo.py)                             | This script shows how to use the [`DPOTrainer`] to fine-tune a stable to increase helpfulness and harmlessness using the [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset.                                                                                                        |\n| [`examples/scripts/kto.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/kto.py)                             | This script shows how to use the [`KTOTrainer`] to fine-tune a model.                                                                                                                                                                                                                                           |\n| [`examples/scripts/orpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/orpo.py)                           | This script shows how to use the [`ORPOTrainer`] to fine-tune a model to increase helpfulness and harmlessness using the [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset.                                                                                                        |\n| [`examples/scripts/ppo_multi_adapter.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/ppo_multi_adapter.py) | This script shows how to use the [`PPOTrainer`] to train a single base model with multiple adapters. Requires you to run the example script with the reward model training beforehand.                                                                                                                          |\n| [`examples/scripts/ppo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/ppo.py)                             | This script shows how to use the [`PPOTrainer`] to fine-tune a sentiment analysis model using [IMDB dataset](https://huggingface.co/datasets/stanfordnlp/imdb).                                                                                                                                                 |\n| [`examples/scripts/reward_modeling.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/reward_modeling.py)     | This script shows how to use the [`RewardTrainer`] to train a reward model on your own dataset.                                                                                                                                                                                                                 |\n| [`examples/scripts/sft.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/sft.py)                             | This script shows how to use the [`SFTTrainer`] to fine-tune a model or adapters into a target dataset.                                                                                                                                                                                                         |\n| [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/vsft_llava.py)               | This script shows how to use the [`SFTTrainer`] to fine-tune a Vision Language Model in a chat setting. The script has only been tested on a [LLaVA 1.5]([llava-hf/llava-1.5-7b-hf](https://huggingface.co/llava-hf/llava-1.5-7b-hf)) model so users may see unexpected behaviour in other model architectures. |\n\nHere are also some easier-to-run colab notebooks that you can use to get started with TRL:\n\n| File                                                                                                                              | Description                                                                                                             |\n| --------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------- |\n| [`examples/notebooks/best_of_n.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/best_of_n.ipynb)           | This notebook demonstrates how to use the \"Best of N\" sampling strategy using TRL when fine-tuning your model with PPO. |\n| [`examples/notebooks/gpt2-sentiment.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-sentiment.ipynb) | This notebook demonstrates how to reproduce the GPT2 imdb sentiment tuning example on a jupyter notebook.               |\n| [`examples/notebooks/gpt2-control.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-control.ipynb)     | This notebook demonstrates how to reproduce the GPT2 sentiment control example on a jupyter notebook.                   |\n\n\nWe also have some other examples that are less maintained but can be used as a reference:\n1. **[research_projects](https://github.com/huggingface/trl/tree/main/examples/research_projects)**: Check out this folder to find the scripts used for some research projects that used TRL (LM de-toxification, Stack-Llama, etc.)\n\n\n## Distributed training\n\nAll of the scripts can be run on multiple GPUs by providing the path of an 🤗 Accelerate config file when calling `accelerate launch`. To launch one of them on one or multiple GPUs, run the following command (swapping `{NUM_GPUS}` with the number of GPUs in your machine and `--all_arguments_of_the_script` with your arguments.)\n\n```shell\naccelerate launch --config_file=examples/accelerate_configs/multi_gpu.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script\n```\n\nYou can also adjust the parameters of the 🤗 Accelerate config file to suit your needs (e.g. training in mixed precision).\n\n### Distributed training with DeepSpeed\n\nMost of the scripts can be run on multiple GPUs together with DeepSpeed ZeRO-{1,2,3} for efficient sharding of the optimizer states, gradients, and model weights. To do so, run following command (swapping `{NUM_GPUS}` with the number of GPUs in your machine, `--all_arguments_of_the_script` with your arguments, and `--deepspeed_config` with the path to the DeepSpeed config file such as `examples/deepspeed_configs/deepspeed_zero1.yaml`):\n\n```shell\naccelerate launch --config_file=examples/accelerate_configs/deepspeed_zero{1,2,3}.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script\n```\n\n\n# RLOO Trainer\n\nTRL supports training LLMs with REINFORCE Leave-One-Out (RLOO). The idea is that instead of using a value function, RLOO generates K completions for each prompt. For each completion, RLOO uses the mean scores from the other K-1 completions as a baseline to calculate the advantage. RLOO also models the entire completion as a single action, where as PPO models each token as an action. Note that REINFORCE / A2C is a special case of PPO, when the number of PPO epochs is 1 and the number of mini-batches is 1, which is how we implement RLOO in TRL.\n\nReferences:\n- [Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs](https://huggingface.co/papers/2402.14740)\n- [A2C is a special case of PPO](https://huggingface.co/papers/2205.09123)\n- [Fine-Tuning Language Models from Human Preferences](https://github.com/openai/lm-human-preferences)\n- [Learning to Summarize from Human Feedback](https://github.com/openai/summarize-from-feedback)\n- [The N Implementation Details of RLHF with PPO](https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo)\n- [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031)\n\n## Get started\n\nTo just run a RLOO script to make sure the trainer can run, you can run the following command to train a RLOO model with a dummy reward model.\n\n```bash\npython examples/scripts/rloo/rloo.py \\\n    --learning_rate 3e-6 \\\n    --output_dir models/minimal/rloo \\\n    --per_device_train_batch_size 64 \\\n    --gradient_accumulation_steps 1 \\\n    --total_episodes 10000 \\\n    --model_name_or_path EleutherAI/pythia-14m \\\n    --reward_model_path EleutherAI/pythia-14m \\\n    --missing_eos_penalty 1.0\n```\n\n\n## Explanation of the logged metrics\n\nThe logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/u2sqci34)\n\n<!-- * `rlhf_reward_var_per_prompt`: calculated by `rlhf_reward.var(0).mean()`. This is the variance of the rewards estimated across the `args.rloo_k` samples. Usually we expect it to go down (cause policy entropy goes down). -->\n\n* `eps`: Tracks the number of episodes per second.\n* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current policy and reference policy.\n* `objective/entropy`: The mean entropy of the policy, indicating the randomness of the actions chosen by the policy.\n* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.\n* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.\n* `objective/scores`: The mean scores returned by the reward model / environment.\n* `policy/approxkl_avg`: The average approximate KL divergence between consecutive PPO policies. Note that this is not the same as `objective/kl`.\n* `policy/clipfrac_avg`: The average fraction of policy updates that are clipped, indicating how often the policy updates are constrained to prevent large changes.\n* `loss/policy_avg`: The average policy loss, indicating how well the policy is performing.\n* `val/clipfrac_avg`: The average fraction of value function updates that are clipped, similar to policy/clipfrac_avg but for the value function.\n* `policy/entropy_avg`: The average entropy of the policy during training, indicating how diverse the policy's actions are.\n* `val/ratio`: The mean ratio of the current policy probability to the old policy probability, providing a measure of how much the policy has changed.\n* `val/ratio_var`: The variance of the `val/ratio`, indicating the variability in policy changes.\n* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.\n* `lr`: lr: The current learning rate used by the optimizer.\n* `episode`: episode: The current global step or episode count in the training process.\n\n\n## Cookbook\n\n* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.\n* Debugging TIP: `val/ratio`: this number should float around 1.0, and it gets clipped by `--cliprange 0.2` with PPO's surrogate loss. So if this `ratio` is too high like 2.0 or 1000.0 or too small like 0.1, it means the updates between consecutive policies are too drastic. You should try undertand why this is happening and try to fix it.\n* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.\n* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.\n* Usage TIP: We recommend to use the \"EOS trick\" via `--missing_eos_penalty`, which subtracts a static scalar penalty from the score of completions that do not end with an EOS token. This can help the model learn to generate more coherent completions.\n\n\n## What is my model doing exactly?\n\nTo help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/u2sqci34), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif)\n\n\nIn the logs the sampled generations look like \n\n```\n┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓\n┃ query                           ┃ model response                  ┃ score    ┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩\n│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │\n│                                 │ I don't know how to get rid of  │          │\n│ TITLE: How do you get someone   │ those feelings. I'm             │          │\n│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │\n│                                 │                                 │          │\n│ POST: Hi,                       │                                 │          │\n│ I'm 22, and I have been with my │                                 │          │\n│ girlfriend for 5 years now. We  │                                 │          │\n│ recently moved together. We've  │                                 │          │\n│ always loved each other         │                                 │          │\n│ intensely.                      │                                 │          │\n│                                 │                                 │          │\n│ Problem, I recently started to  │                                 │          │\n│ have feelings for an other      │                                 │          │\n│ person (a friend). This person  │                                 │          │\n│ has had a boyfriend for now 3   │                                 │          │\n│ years, and has absolutely no    │                                 │          │\n│ ideas. Those feelings were so   │                                 │          │\n│ strong, it was hard to hide     │                                 │          │\n│ them. After 2 months of me      │                                 │          │\n│ being distant and really sad,   │                                 │          │\n│ my girlfriend forced me to say  │                                 │          │\n│ what was bothering me. I'm not  │                                 │          │\n│ a good liar, and now she knows. │                                 │          │\n│                                 │                                 │          │\n│ We decided to give us a week    │                                 │          │\n│ alone, I went to my parents.    │                                 │          │\n│                                 │                                 │          │\n│ Now, I'm completely lost. I     │                                 │          │\n│ keep on thinking about this     │                                 │          │\n│ person, and I hate that. I      │                                 │          │\n│ would like for those feelings   │                                 │          │\n│ to go away, to leave me alone.  │                                 │          │\n│ But I can't.                    │                                 │          │\n│                                 │                                 │          │\n│ What do I do? It's been 3       │                                 │          │\n│ months now, and I'm just        │                                 │          │\n│ desperate.                      │                                 │          │\n│                                 │                                 │          │\n│ TL;DR:                          │                                 │          │\n├─────────────────────────────────┼─────────────────────────────────┼──────────┤\n│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │\n│                                 │ TV. I blasted Gangnam Style on  │          │\n│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │\n│ with a loud TV.                 │ up as high as it could          │          │\n│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │\n│ POST: She was in her living     │                                 │          │\n│ room, watching TV. This was at  │                                 │          │\n│ about 8:30 in the morning, and  │                                 │          │\n│ she was exercising. She turned  │                                 │          │\n│ the TV up extra loud to hear it │                                 │          │\n│ over her excercycle, and woke   │                                 │          │\n│ me up. I went in there asking   │                                 │          │\n│ for her to turn it down. She    │                                 │          │\n│ said she didn't have to; I      │                                 │          │\n│ explained that I always used    │                                 │          │\n│ headphones so she didn't have   │                                 │          │\n│ to deal with my noise and that  │                                 │          │\n│ she should give me a little     │                                 │          │\n│ more respect, given that I paid │                                 │          │\n│ rent at the time.               │                                 │          │\n│                                 │                                 │          │\n│ She disagreed. I went back to   │                                 │          │\n│ my room, rather pissed off at   │                                 │          │\n│ the lack of equality. I had no  │                                 │          │\n│ lock on my door; but I had a    │                                 │          │\n│ dresser right next to it, so I  │                                 │          │\n│ pulled one of the drawers out   │                                 │          │\n│ enough so that it caused the    │                                 │          │\n│ door to not be openable. Then,  │                                 │          │\n│ I turned my speakers up really  │                                 │          │\n│ loud and blasted Gangnam Style  │                                 │          │\n│ on repeat, with the bass        │                                 │          │\n│ cranked up as high as it could  │                                 │          │\n│ go.                             │                                 │          │\n│                                 │                                 │          │\n│ If you hate Gangnam Style for   │                                 │          │\n│ being overplayed, you will see  │                                 │          │\n│ why I chose that particular     │                                 │          │\n│ song. I personally don't mind   │                                 │          │\n│ it. But here's the thing about  │                                 │          │\n│ my bass; it vibrates the walls, │                                 │          │\n│ making one hell of a lot of     │                                 │          │\n│ noise. Needless to say, my mom  │                                 │          │\n│ was not pleased and shut off    │                                 │          │\n│ the internet. But it was oh so  │                                 │          │\n│ worth it.                       │                                 │          │\n│                                 │                                 │          │\n│ TL;DR:                          │                                 │          │\n└─────────────────────────────────┴─────────────────────────────────┴──────────┘\n```\n\n## Implementation details\n\nThe bulk of RLOOTrainer is based on the PPO implementation, which is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).\n\n\nBelow is a vectorized advantage calculation for RLOO:\n\n```python\ndef test_rloo_reward():\n    local_batch_size = 3\n    rloo_k = 4\n    rlhf_reward = torch.tensor([\n        1, 2, 3, # first rlhf reward for three prompts\n        2, 3, 4, # second rlhf reward for three prompts\n        5, 6, 7, # third rlhf reward for three prompts\n        8, 9, 10, # fourth rlhf reward for three prompts\n    ]).float() # here we have 3 prompts which have 4 completions each\n\n    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)\n    advantages = torch.zeros_like(rlhf_reward)\n    for i in range(0, len(advantages), local_batch_size):\n        other_response_rlhf_rewards = []\n        for j in range(0, len(advantages), local_batch_size):\n            if i != j:\n                other_response_rlhf_rewards.append(rlhf_reward[j : j + local_batch_size])\n        advantages[i : i + local_batch_size] = rlhf_reward[i : i + local_batch_size] - torch.stack(other_response_rlhf_rewards).mean(0)\n    \n    assert (1 - (2 + 5 + 8) / 3 - advantages[0].item()) < 1e-6  # First rlhf reward for the first prompt\n    assert (6 - (3 + 2 + 9) / 3 - advantages[7].item()) < 1e-6  # Third rlhf reward for the second prompt\n\n    # Vectorized implementation\n    rlhf_reward = rlhf_reward.reshape(rloo_k, local_batch_size)\n    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)\n    vec_advantages = rlhf_reward - baseline\n    torch.testing.assert_close(vec_advantages.flatten(), advantages)\n```\n\n## Benchmark experiments\n\nTo validate the RLOO implementation works, we ran experiment on the 1B model. Here are the command we used to run the experiment. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).\n\n```\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/rloo/rloo_tldr.py \\\n    --output_dir models/minimal/rloo_tldr \\\n    --num_ppo_epochs 2 \\\n    --num_mini_batches 2 \\\n    --learning_rate 3e-6 \\\n    --per_device_train_batch_size 8 \\\n    --gradient_accumulation_steps 8 \\\n    --total_episodes 1000000 \\\n    --model_name_or_path EleutherAI/pythia-1b-deduped \\\n    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \\\n    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \\\n    --local_rollout_forward_batch_size 16 \\\n    --missing_eos_penalty 1.0 \\\n    --stop_token eos \\\n    --kl_coef 0.03\n```\n\nCheckpoints and experiment tracking are available at:\n\n- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/rloo_tldr)\n- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/u2sqci34)\n\n\nTo evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.\nFor more information on how to use judges, see [Judges](judges).\n\n```bash\n$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 33.00%\n$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 51.20%\n```\n\nThe RLOO checkpoint gets a 51.2% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the RLOO training is working as intended.\n\n\nMetrics:\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/rloo.png)\n\n\n```bash\n# pip install openrlbenchmark==0.2.1a5\n# see https://github.com/openrlbenchmark/openrlbenchmark#get-started for documentation\n# to use it, change `?we=huggingface&wpn=trl` to your own project and `?tag=pr-1540` to your own tag\npython -m openrlbenchmark.rlops_multi_metrics \\\n    --filters '?we=huggingface&wpn=trl&xaxis=train/episode&ceik=output_dir&cen=sft_model_path&metrics=train/objective/rlhf_reward&metrics=train/objective/scores&metrics=train/objective/kl&metrics=train/objective/non_score_reward&metrics=train/objective/entropy&metrics=train/policy/approxkl_avg&metrics=train/policy/clipfrac_avg&metrics=train/loss/policy_avg&metrics=train/policy/entropy_avg&metrics=train/val/ratio&metrics=train/val/ratio_var&metrics=train/val/num_eos_tokens&metrics=train/lr&metrics=train/eps' \\\n        \"cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr?tag=pr-1540\" \\\n    --env-ids models/minimal/rloo_tldr \\\n    --pc.ncols 4 \\\n    --pc.ncols-legend 1 \\\n    --pc.xlabel \"Episode\" \\\n    --output-filename benchmark/trl/pr-1540/rloo \\\n    --scan-history\n```\n\n\n## RLOOTrainer\n\n[[autodoc]] RLOOTrainer\n\n## RLOOConfig\n\n[[autodoc]] RLOOConfig\n\n# Online DPO Trainer\n\n## Overview \n\nOnline DPO was proposed in [Direct Language Model Alignment from Online AI Feedback](https://huggingface.co/papers/2402.04792) by Shangmin Guo, Biao Zhang, Tianlin Liu, Tianqi Liu, Misha Khalman, Felipe Llinares, Alexandre Rame, Thomas Mesnard, Yao Zhao, Bilal Piot, Johan Ferret, and Mathieu Blondel. \n\nThe abstract from the paper is the following:\n\n> Direct alignment from preferences (DAP) methods, such as DPO, have recently emerged as efficient alternatives to reinforcement learning from human feedback (RLHF), that do not require a separate reward model. However, the preference datasets used in DAP methods are usually collected ahead of training and never updated, thus the feedback is purely offline. Moreover, responses in these datasets are often sampled from a language model distinct from the one being aligned, and since the model evolves over training, the alignment phase is inevitably off-policy. In this study, we posit that online feedback is key and improves DAP methods. Our method, online AI feedback (OAIF), uses an LLM as annotator: on each training iteration, we sample two responses from the current model and prompt the LLM annotator to choose which one is preferred, thus providing online feedback. Despite its simplicity, we demonstrate via human evaluation in several tasks that OAIF outperforms both offline DAP and RLHF methods. We further show that the feedback leveraged in OAIF is easily controllable, via instruction prompts to the LLM annotator.\n\nThe current implementation uses reward models for scoring completions -- see [Reward Bench](https://huggingface.co/spaces/allenai/reward-bench) for a leaderboard of public models you can use.\n\nThis post-training method was contributed by [Michael Noukhovitch](https://huggingface.co/mnoukhov), [Shengyi Costa Huang](https://huggingface.co/vwxyzjn), [Quentin Gallouédec](https://huggingface.co/qgallouedec), and [Edward Beeching](https://huggingface.co/edbeeching).\n\n## Quick start\n\nThis example demonstrates how to train a model using the online DPO method. We use the [Qwen 0.5B model](https://huggingface.co/Qwen/Qwen2-0.5B-Instruct) as the base model and the [Qwen 0.5B reward model](https://huggingface.co/trl-lib/Qwen2-0.5B-Reward) as the reward model. We use the prompts from the [UltraFeedback dataset](https://huggingface.co/datasets/openbmb/UltraFeedback). You can view the prompts in the dataset here:\n\n<iframe\n  src=\"https://huggingface.co/datasets/trl-lib/ultrafeedback-prompt/embed/viewer/default/train?row=0\"\n  frameborder=\"0\"\n  width=\"100%\"\n  height=\"560px\"\n></iframe>\n\nBelow is the script to train the model:\n\n```python\n# train_online_dpo.py\nfrom datasets import load_dataset\nfrom trl import OnlineDPOConfig, OnlineDPOTrainer\nfrom transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer\n\nmodel = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2-0.5B-Instruct\")\nreward_model = AutoModelForSequenceClassification.from_pretrained(\"trl-lib/Qwen2-0.5B-Reward\", num_labels=1)\ntrain_dataset = load_dataset(\"trl-lib/ultrafeedback-prompt\", split=\"train\")\n\ntraining_args = OnlineDPOConfig(output_dir=\"online-dpo-qwen2\", logging_steps=10)\ntrainer = OnlineDPOTrainer(\n    model=model,\n    reward_model=reward_model,\n    args=training_args,\n    tokenizer=tokenizer,\n    train_dataset=train_dataset,\n)\ntrainer.train()\n```\n\nExecute the script using the following command:\n\n```bash\naccelerate launch train_online_dpo.py\n```\n\nDistributed across 8 GPUs, the training takes approximately 1 hour. You can verify the training progress by checking the reward graph. An increasing trend in both the reward for rejected and chosen completions indicates that the model is improving and generating better responses over time.\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/online-dpo-qwen2-reward.png)\n\nTo see how the trained model performs, use the following code to generate completions:\n\n```python\n>>> from transformers import pipeline\n>>> generator = pipeline(\"text-generation\", model=\"online-dpo-qwen2/checkpoint-1773\", device=\"cuda\")\n>>> question = \"Why is the problem always DNS?\"\n>>> output = generator([{\"role\": \"user\", \"content\": question}], max_new_tokens=200, return_full_text=False)[0]\n>>> print(output[\"generated_text\"])\nThe reason why the problem of DNS (Domain Name System) can always be encountered is that it is designed to provide reliable and accurate information about the availability, ownership, or expiration of domain names. However, there may be some circumstances where the system fails to resolve an IP address correctly, leading to the problem of DNS.\nFor example, if the server hosting the domain name does not have the correct IP address associated with it, or if the IP address is incorrectly formatted, then the DNS system will fail to resolve the domain name correctly. Additionally, if the server hosting the domain name has been compromised, then the DNS system may also fail to resolve the domain name correctly.\nIt's worth noting that the exact cause of DNS failure can vary depending on the specific situation, so it's important to carefully check all relevant factors before attempting to resolve the issue. If you suspect that your DNS problem may be caused by a bug in the system, you should report it to the DNS provider directly for further investigation.\n```\n\n## Expected dataset format\n\nOnline DPO only requires a [prompt-only dataset](dataset_format#preference) (unlike offline DPO, that expects [preference dataset](dataset_format#preference)). The [`OnlineDPOTrainer`] supports both [conversational](dataset_format#conversational-dataset-format) and [standard](dataset_format#standard-dataset-format) dataset format. When provided with a conversational dataset, the trainer will automatically apply the chat template to the dataset.\n\n## Usage tips\n\n### ⚠️ Use the same chat template\n\nMake sure that the SFT model and reward model use the _same_ chat template. Otherwise, you may find the model completions are scored incorrectly during training.\n\n### Encourage EOS token generation\n\nWe can want the model to generate completion within a given length. During the learning, the model will generate completion up to the maximum completion length specified in the `max_new_tokens` argument of [`OnlineDPOConfig`]. I you want to penalize for not generating an EOS token before the maximum completion length, you can use the `missing_eos_penalty` argument of [`OnlineDPOConfig`]:\n\n```python\ntraining_args = OnlineDPOConfig(..., max_new_tokens=128, missing_eos_penalty=1.0)\n```\n\n### Logging Completions\n\nTo better understand your model’s behavior during training, you can log sample completions periodically using the [`LogCompletionsCallback`].\n\n```python\ntrainer = OnlineDPOTrainer(..., eval_dataset=eval_dataset)\ncompletions_callback = LogCompletionsCallback(trainer, num_prompts=8)\ntrainer.add_callback(completions_callback)\n```\n\nThis callback logs the model's generated completions directly to Weights & Biases.\n\n![Logged Completions](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/wandb_completions.png)\n\n\n## Example script\n\nWe provide an example script to train a model using the online DPO method. The script is available in [`examples/scripts/dpo_online.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/dpo_online.py)\n\nTo test the online DPO script with the [Pythia 1B model](https://huggingface.co/trl-lib/pythia-1b-deduped-tldr-sft) on the TL;DR summarization task, run the following command:\n\n```bash\npython examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-tldr-online-dpo \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 32 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 53 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --push_to_hub\n```\n\n## Logged metrics\n\nThe logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35)\n\n* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current model and reference model.\n* `objective/entropy`: The mean entropy of the model, indicating the randomness of the actions chosen by the model.\n* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.\n* `objective/rlhf_reward`: The mean RLHF reward, which is `scores - non_score_reward`. The `rlhf_reward` is the ultimate objective of online DPO training. If training works as intended, this metric should keep going up.\n* `objective/scores`: The mean scores returned by the reward mode.\n* `objective/scores_margin`: The mean score margin (according to the external reward model) between the chosen and rejected completions.\n* `rewards/chosen`: The mean reward (according to online DPO's implicit reward model)of the chosen completions.\n* `rewards/rejected`: The mean reward (according to online DPO's implicit reward model) of the rejected completions.\n* `rewards/accuracies`: The accuracies of the online DPO's implicit reward model.\n* `rewards/margins`: The mean reward margin (according to online DPO's implicit reward model) between the chosen and rejected completions.\n* `logps/chosen`: The mean log probabilities of the chosen completions.\n* `logps/rejected`: The mean log probabilities of the rejected completions.\n* `val/contain_eos_token`: The fraction of completions which contain an EOS token.\n* `beta`: The parameter that controls the weight of the loss term representing the deviation from the reference model. Typically fixed, but can be made dynamic by passing a list to [`OnlineDPOConfig`].\n\n## Benchmark experiments\n\nTo validate the online DPO implementation works, we ran experiments with the Pythia 1B, 2.8B, and 6.9B models on a single node of 8 x H100s. Here are the commands we used to run the experiments. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).\n\n\n```\n# 1B Online DPO experiment\naccelerate launch --config_file examples/accelerate_configs/multi_gpu.yaml \\\n    examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-1b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-1b-deduped-tldr-online-dpo \\\n    --beta 0.1 \\\n    --per_device_train_batch_size 8 \\\n    --gradient_accumulation_steps 2 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 53 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --logging_steps 20 \\\n    --save_steps 0.1 \\\n    --push_to_hub\n\n# 2.8B Online DPO experiment\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-2.8b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-2.8b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-2.8b-deduped-tldr-online-dpo \\\n    --beta 0.1 \\\n    --per_device_train_batch_size 8 \\\n    --gradient_accumulation_steps 2 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 53 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --bf16 \\\n    --logging_steps 20 \\\n    --save_steps 0.1 \\\n    --push_to_hub\n\n# 6.9B Online DPO experiment\naccelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \\\n    examples/scripts/dpo_online.py \\\n    --model_name_or_path trl-lib/pythia-6.9b-deduped-tldr-sft  \\\n    --reward_model_path trl-lib/pythia-6.9b-deduped-tldr-rm \\\n    --dataset_name trl-lib/tldr \\\n    --learning_rate 5.0e-7 \\\n    --output_dir pythia-6.9b-deduped-tldr-online-dpo \\\n    --beta 0.1 \\\n    --per_device_train_batch_size 4 \\\n    --gradient_accumulation_steps 4 \\\n    --num_train_epochs 3 \\\n    --max_new_tokens 53 \\\n    --warmup_ratio 0.1 \\\n    --missing_eos_penalty 1.0 \\\n    --bf16 \\\n    --gradient_checkpointing \\\n    --logging_steps 20 \\\n    --save_steps 0.1 \\\n    --push_to_hub\n```\n\nCheckpoints and experiment tracking are available at:\n\n- [🤗 Model checkpoints](https://huggingface.co/collections/trl-lib/online-dpo-66acd3fa38a331a9cd457b07)\n- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/reports/Online-DPO-experiments-for-TL-DR-summarisation--Vmlldzo5MTczMDU0)\n\n\nTo evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.\nFor more information on how to use judges, see [Judges](judges).\n\n```bash\n$ python examples/scripts/evals/judge_tldr.py --model_name_or_path trl-lib/pythia-1b-deduped-tldr-sft --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 33.00%\npython examples/scripts/evals/judge_tldr.py --model_name_or_path trl-lib/pythia-6.9b-deduped-tldr-sft --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 41.50%\npython examples/scripts/evals/judge_tldr.py --model_name_or_path trl-lib/pythia-1b-deduped-tldr-online-dpo --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 62.60%\npython examples/scripts/evals/judge_tldr.py --model_name_or_path trl-lib/pythia-6.9b-deduped-tldr-online-dpo --judge_model gpt-4o-mini --num_examples 1000\nModel win rate: 74.20%\n```\n\nWe can then plot the RLHF scaling chart.\n\n```python\nimport matplotlib.pyplot as plt\n\nresults = {\n    \"SFT\": {1.0e9: 0.21, 2.8e9: 0.27, 6.9e9: 0.316},\n    \"online-dpo\": {1.0e9: 0.542, 2.8e9: 0.746, 6.9e9: 0.796},\n    \"offline-dpo\": {1.0e9: 0.422, 2.8e9: 0.517, 6.9e9: 0.701},\n}\n\n\nplt.plot(results[\"SFT\"].keys(), results[\"SFT\"].values(), label=\"SFT\", marker=\"o\")\nplt.plot(results[\"online-dpo\"].keys(), results[\"online-dpo\"].values(), label=\"Online-dpo with RM judge\", marker=\"o\")\nplt.plot(results[\"offline-dpo\"].keys(), results[\"offline-dpo\"].values(), label=\"Offline-dpo\", marker=\"o\")\nplt.axhline(y=0.5, color=\"black\", linestyle=\"-.\", label=\"Human reference summary\")\nplt.xscale(\"log\")\nplt.xlabel(\"Model size\")\nplt.ylabel(\"Win rate against reference summaries\\n(according to GPT-4-0613)\")\nplt.title(\"DPO scaling by model size\")\nplt.legend()\nplt.xlim(5e8, 1.2e10)\nplt.xticks([1e9, 3e9, 1e10], [\"1B\", \"3B\", \"10B\"])\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\")\nplt.tight_layout()\nplt.show()\n```\n\n![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/online_dpo_scaling.png)\n\nThe online DPO checkpoint gets increasingly more win rate as we scale up the model sizes. This is a good sign that the online DPO implementation is working as intended.\n\n## OnlineDPOTrainer\n\n[[autodoc]] OnlineDPOTrainer\n\n## OnlineDPOConfig\n\n[[autodoc]] OnlineDPOConfig\n\n# Text Environments\n\nText environments provide a learning ground for language agents. It allows a language model to use tools to accomplish a task such as using a Python interpreter to answer math questions or using a search index for trivia questions. Having access to tools allows language models to solve tasks that would be very hard for the models itself but can be trivial for the appropriate tools. A good example is arithmetics of large numbers that become a simple copy-paste task once you have access to a calculator.\n\n<div style=\"text-align: center\">\n<img src=\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/textenv.png\">\n</div>\n\nLet's dive into how text environments work and start with tools!\n\n## Tools\n\nOne of the core building blocks of text environments are tools that the model can use to solve tasks. In general tools can be any Python function that takes a string as input and returns string. The `TextEnvironment` offers two options for tools: either go with predefined tools from `transformers.Tool` or define your own function or class with `__call__` method. Let's have a look at both!\n\n### `transformers.Tool`\n\nText environments fully support tools of the class `transformers.Tool`. The advantage of building tools in that framework is that they can easily be shared \n\n```Python\nfrom transformers import load_tool\n\n# simple calculator tool that runs +-/* operations\ncalc_tool = load_tool(\"ybelkada/simple-calculator\")\n\n# python interpreter that executes program and returns outputs\npy_tool = load_tool(\"lvwerra/python-interpreter\")\n\n# wikipedia search index that returns best search match\nwiki_tool = load_tool(\"vwxyzjn/pyserini-wikipedia-kilt-doc\")\n```\n\nThese tools are either loaded from the hub or from a local folder. Using the tool is as simple as calling them with a text query:\n\n```Python\ncalc_tool(\"1/2\")\n>>> \"0.5\"\n```\n\nNote that both input and return values are strings to enable easy usage with a language model.\n\n### Custom Tools\n\nThe following is an example of a tool that adds two integers:\n\n```Python\ndef add(text):\n    int_1, int_2 = text.split(\"+\")\n    result = int(int_1) + int(int_2)\n    return str(result)\n\nprint(add(\"1+1\"))\n>>> \"2\"\n```\n\nWe looked at basic examples such as a calculator but the principle holds for more complex tools as well such as a web search tool where you input the query and get the search results in return. Now let's look at how the model can use the tools with the call syntax.\n\n### Call syntax\n\nIn order to have a unified way for the model to call a tool we created a simple syntax that looks as follows:\n\n```python\n\"<request><TOOL_NAME>QUERY<call>TOOL_RESPONSE<response>\"\n```\n\nThere are a few special tokens involved so let's decompose it: First the model can signal that it wants to use a tool by emitting the `<request>` token. After that we want to know the name of the tool to call which is done by enclosing the tool name with `<>` brackets. Once we know which tool to call the tool query follows which is in free text form. The `<call>` tokens signifies the end of the query and stops the model generation. At this point the model output is parsed and the query sent to the tool. The environment appends the tool response to the string followed by the `<response>` token to show the end the tool output.\n\nLet's look at the concrete example of the calculator and assume its name is `Calculator` (more on how the name of a tool is inferred later):\n\n```python\n\"<request><Calculator>1/2<call>0.5<response>\"\n```\n\nFinally, the episode is ended and generation stops when the model generates `<submit>` which marks the interaction as completed.\n\nNow let's have a look how we can create a new text environment!\n\n## Create a `TextEnvironment`\n\n\n```python\nprompt = \"\"\"\\\nWhat is 13-3?\n<request><SimpleCalculatorTool>13-3<call>10.0<response>\nResult=10<submit>\n\"\"\"\n\ndef reward_fn(result, answer):\n    \"\"\"Simplified reward function returning 1 if result matches answer and 0 otherwise.\"\"\"\n    result_parsed = result.split(\"=\")[1].split(\"<\")[0]\n    return int(result_parsed==answer)\n\ntext_env = TextEnvironemnt(\n    model=model, \n    tokenizer=tokenizer,\n    tools= {\"SimpleCalculatorTool\": load_tool(\"ybelkada/simple-calculator\")},\n    reward_fn=exact_match_reward,\n    prompt=prompt, \n    max_turns=1\n    max_tool_response=100\n    generation_kwargs={\"do_sample\": \"true\"}\n)\n```\n\nLet's decompose the settings:\n\n| Argument           | Description     |\n|:-------------------|:----------------|\n| `model`            | Language model to interact with the environment and generate requests. |\n| `tokenizer`        | Tokenizer of language model handling tokenization of strings. |\n| `tools`            | `list` of `dict` of tools. If former the name of the tool is inferred from class name and otherwise it's the keys of the dictionary.|\n| `reward_fn`        | A function that takes a string as input and returns. Can have extra arguments that are passed to `.run()` such as ground truth.|\n| `prompt`           | Prompt to prepend to every task. Usually a few examples to demonstrate to the model how to use the tools in a few-shot fashion. |\n| `max_turns`        | Maximum number of interactions between model and tools before episode ends.|\n| `max_tool_response`| The tool response is truncated to this number to avoid running out of model context.|\n| `max_length`       |  The maximum number of tokens to allow in an episode. |\n| `generation_kwargs`| Generation settings used by the language model. |\n\nYou can customize the environment to your needs and add custom tools and settings. Let's see how you can use the environment to have the model interact with the available tools!\n\n\n## Run an Episode\n\nTo run a set of queries through the text environment one can simply use the `run` method.\n\n```python\nqueries = [\"What is 1/2?\"]\nanswers = [\"0.5\"]\n\nqueries, responses, masks, rewards, histories = text_env.run(queries, answers=answers)\n```\n\nThis will execute the model/tool feedback loop for each query until either no tool is called anymore, the maximum number of turns is reached or to maximum number of tokens in an episode is exceeded. The extra `kwargs` (e.g. `answers=answers` above) passed to `run` will be passed on to the reward function.\n\nThere are five objects that are returned by `run`: \n\n- `queries`: a list of the tokenized queries\n- `responses`: all tokens that have been generated withing the environment including model and tool tokens\n- `masks`: mask that indicates which tokens have been generated by the model and which tokens are generated by the tool\n- `rewards`: a list of reward for each query/response\n- `histories`: list of `TextHistory` objects, which are useful objects containing all the above and also the text equivalents\n\nThe masks are crucial for training as we don't want to optimize tokens that the model has not generated which are tokens produced by the tools.\n\nNext, we'll train a PPO step with the generated responses!\n\n\n### Train\nTraining on episodes from the `TextEnvironment` is straight forward and simply requires forwarding all the returned variables except the `TextHistory` objects to the `step` method:\n\n```python\ntrain_stats = ppo_trainer.step(queries, responses, rewards, masks)\n```\n\n## `TextHistory`\n\nThe `TextHistory` object stores the interactions between the model and the text environment. It stores tokens and text generated in each turn and their source in each turn (model or system) as well as rewards. Let's go through the class attributes and methods.\n\n### Attributes\n\nThe following table summarises the available attributes of the `TextEnvironment` class:\n\n| Attribute           | Description     |\n|:-------------------|:----------------|\n| `text`             | The full string of the text generated in the text environment with both model and system generated text. |\n| `text_spans`       | A list of tuples with the spans for each model or system generated text segment. |\n| `system_spans`     | A list of boolean values indicating if the segment is model or system generated. |\n| `tokens`           | All tokens generated in text environment with both model and system generated tokens. |\n| `token_spans`      | Similar to `text_spans` the `token_spans` indicate the boundaries of model andsystem generated tokens. |\n| `token_masks`      | The token masks can be used to ignore system generated tokens by masking them. |\n| `completed`        | Indicates if the interaction with the environment has completed. |\n| `truncated`        | Indicates if the interaction with the environment has completed because max length was reached. |\n\nWith these attributes you can reconstruct every interaction of the model with the `TextEnvironment`. The `TextHistory` also lets you visualize the text history. Let's have a look!\n\n### Visualization\n\nWhen the model interacts inside the `TextEnvironment` it can be useful to visualize and separate which parts of the text outputs were generated by the model and which parts come from the system and tools. For that purpose there are the two methods [`TextHistory.show_text`] and [`TextHistory.show_tokens`]. They print the text and tokens respectively and highlight the various segments using the [`rich` libray](https://github.com/Textualize/rich) (make sure to install it before using these methods).\n\nYou can see that the prompt is highlighted in gray, whereas system segments such as query and tool responses are highlighted in green. All segments generated by the model are highlighted in blue and in addition to the pure text output the reward is displayed as additional text in plum. Here an example of `show_text`:\n\n<div style=\"text-align: center\">\n<img src=\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/textenv_show_text.png\" width=600>\n</div>\n\nSometimes there can be tricky tokenization related issues that are hidden when showing the decoded text. Thus `TextHistory` also offers an option to display the same highlighting on the tokens directly with `show_tokens`:\n\n<div style=\"text-align: center\">\n<img src=\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/textenv_show_tokens.png\" width=800>\n</div>\n\nNote that you can turn on the colour legend by passing `show_legend=True`.\n\n## API Documentation\n\n[[autodoc]] TextEnvironment\n\n[[autodoc]] TextHistory\n\n\n# ORPO Trainer\n\n[Odds Ratio Preference Optimization](https://huggingface.co/papers/2403.07691) (ORPO) by Jiwoo Hong, Noah Lee, and James Thorne studies the crucial role of SFT within the context of preference alignment. Using preference data the method posits that a minor penalty for the disfavored generation together with a strong adaption signal to the chosen response via a simple log odds ratio term appended to the NLL loss is sufficient for preference-aligned SFT.\n\nThus ORPO is a reference model-free preference optimization algorithm eliminating the necessity for an additional preference alignment phase thus saving compute and memory.\n\nThe official code can be found [xfactlab/orpo](https://github.com/xfactlab/orpo).\n\n## Expected dataset format\n\nThe ORPO trainer expects a format identical to the DPO trainer, which should include three entries. These entries should be named as follows:\n\n- `prompt`\n- `chosen`\n- `rejected`\n\nfor example:\n\n```py\norpo_dataset_dict = {\n    \"prompt\": [\n        \"hello\",\n        \"how are you\",\n        \"What is your name?\",\n        \"What is your name?\",\n        \"Which is the best programming language?\",\n        \"Which is the best programming language?\",\n        \"Which is the best programming language?\",\n    ],\n    \"chosen\": [\n        \"hi nice to meet you\",\n        \"I am fine\",\n        \"My name is Mary\",\n        \"My name is Mary\",\n        \"Python\",\n        \"Python\",\n        \"Java\",\n    ],\n    \"rejected\": [\n        \"leave me alone\",\n        \"I am not fine\",\n        \"Whats it to you?\",\n        \"I dont have a name\",\n        \"Javascript\",\n        \"C++\",\n        \"C++\",\n    ],\n}\n```\nwhere the `prompt` contains the context inputs, `chosen` contains the corresponding chosen responses and `rejected` contains the corresponding negative (rejected) responses. Note that a prompt can have multiple responses and this is reflected in the entries being repeated in the dictionary's value arrays.\n\n## Expected model format\nThe ORPO trainer expects a model of `AutoModelForCausalLM`, compared to PPO that expects `AutoModelForCausalLMWithValueHead` for the value function.\n\n## Using the `ORPOTrainer`\nFor a detailed example have a look at the `examples/scripts/orpo.py` script. At a high level we need to initialize the `ORPOTrainer` with a `model` we wish to train. **Note that ORPOTrainer eliminates the need to use the reference model, simplifying the optimization process.** The `beta` refers to the hyperparameter `lambda` in eq. (6) of the paper and refers to the weighting of the relative odd ratio loss in the standard cross-entropy loss used for SFT.\n\n```py\ntraining_args = ORPOConfig(\n    beta=0.1, # the lambda/alpha hyperparameter in the paper/code\n)\n\norpo_trainer = ORPOTrainer(\n    model,\n    args=training_args,\n    train_dataset=train_dataset,\n    tokenizer=tokenizer,\n)\n```\nAfter this one can then call:\n\n```py\norpo_trainer.train()\n```\n\n### For Mixture of Experts Models: Enabling the auxiliary loss\n\nMOEs are the most efficient if the load is about equally distributed between experts.  \nTo ensure that we train MOEs similarly during preference-tuning, it is beneficial to add the auxiliary loss from the load balancer to the final loss.  \n\nThis option is enabled by setting `output_router_logits=True` in the model config (e.g. MixtralConfig).  \nTo scale how much the auxiliary loss contributes to the total loss, use the hyperparameter `router_aux_loss_coef=...` (default: 0.001).\n\n## Logging\n\nWhile training and evaluating we record the following reward metrics:\n\n* `rewards/chosen`: the mean log probabilities of the policy model for the chosen responses scaled by beta\n* `rewards/rejected`: the mean log probabilities of the policy model for the rejected responses scaled by beta\n* `rewards/accuracies`: mean of how often the chosen rewards are > than the corresponding rejected rewards\n* `rewards/margins`: the mean difference between the chosen and corresponding rejected rewards\n\n* `log_odds_chosen`: the mean log odds ratio of the chosen responses over the rejected responses\n\n* `log_odds_ratio`: the mean of the `log(sigmoid(log_odds_chosen))`\n\n* `nll_loss`: the mean negative log likelihood loss from the SFT part of the loss over chosen responses\n \n## ORPOTrainer\n\n[[autodoc]] ORPOTrainer\n\n\n## ORPOConfig\n\n[[autodoc]] ORPOConfig\n\n\n# Copyright 2023 The HuggingFace Team, the AllenNLP library authors. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nScript to close stale issue. Taken in part from the AllenNLP repository.\nhttps://github.com/allenai/allennlp.\n\"\"\"\n\nimport os\nfrom datetime import datetime as dt\nfrom datetime import timezone\n\nfrom github import Github\n\n\nLABELS_TO_EXEMPT = [\n    \"good first issue\",\n    \"good second issue\",\n    \"feature request\",\n    \"help wanted\",\n]\n\n\ndef main():\n    g = Github(os.environ[\"GITHUB_TOKEN\"])\n    repo = g.get_repo(\"huggingface/trl\")\n    open_issues = repo.get_issues(state=\"open\")\n\n    for issue in open_issues:\n        comments = sorted(issue.get_comments(), key=lambda i: i.created_at, reverse=True)\n        involved_users = [comment.user.login for comment in comments]\n        inactive_days = (dt.now(timezone.utc) - issue.updated_at).days\n        is_old = (dt.now(timezone.utc) - issue.created_at).days >= 30\n        has_comments = len([user for user in involved_users if user != \"github-actions[bot]\"]) > 0\n        to_exempt = any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels())\n\n        if is_old and not to_exempt:\n            if has_comments and inactive_days > 23:\n                issue.create_comment(\n                    \"This issue has been automatically marked as stale because it has not had \"\n                    \"recent activity. If you think this still needs to be addressed \"\n                    \"please comment on this thread.\\n\\n\"\n                )\n            elif involved_users and involved_users[0] == \"github-actions[bot]\" and inactive_days > 7:\n                issue.edit(state=\"closed\")\n\n\nif __name__ == \"__main__\":\n    main()\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport subprocess\nimport sys\nfrom datetime import datetime\n\n\nCOPYRIGHT_HEADER = f\"\"\"# Copyright {datetime.now().year} The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\n\nCOPYRIGHT_KEYWORD = \"# Copyright 20\"\n\n\ndef get_tracked_python_files():\n    \"\"\"Get a list of all tracked Python files using git.\"\"\"\n    try:\n        # Get the list of all tracked files from Git\n        result = subprocess.run([\"git\", \"ls-files\"], stdout=subprocess.PIPE, text=True, check=True)\n        # Split the result by lines to get individual file paths\n        files = result.stdout.splitlines()\n        # Filter only Python files\n        py_files = [f for f in files if f.endswith(\".py\")]\n        return py_files\n    except subprocess.CalledProcessError as e:\n        print(f\"Error fetching tracked files: {e}\")\n        return []\n\n\ndef check_and_add_copyright(file_path):\n    \"\"\"Check if the file contains a copyright notice, and add it if missing.\"\"\"\n    if not os.path.isfile(file_path):\n        print(f\"[SKIP] {file_path} does not exist.\")\n        return\n\n    with open(file_path, encoding=\"utf-8\") as f:\n        content = f.readlines()\n\n    # Check if the copyright header exists in the first 10 lines\n    for line in content[:10]:\n        if COPYRIGHT_KEYWORD in line:\n            return True\n\n    # If no copyright notice was found, prepend the header\n    print(f\"[MODIFY] Adding copyright to {file_path}.\")\n    with open(file_path, \"w\", encoding=\"utf-8\") as f:\n        # Write the copyright header followed by the original content\n        f.write(COPYRIGHT_HEADER + \"\\n\" + \"\".join(content))\n    return False\n\n\ndef main():\n    \"\"\"Main function to check and add copyright for all tracked Python files.\"\"\"\n    py_files = get_tracked_python_files()\n    if not py_files:\n        print(\"No Python files are tracked in the repository.\")\n        return\n\n    print(f\"Checking {len(py_files)} Python files for copyright notice...\")\n\n    have_copyright = [check_and_add_copyright(file_path) for file_path in py_files]\n    if not all(have_copyright):\n        print(\"❌ Some files were missing the required copyright and have been updated.\")\n        sys.exit(1)\n    else:\n        print(\"✅ All files have the required copyright.\")\n        sys.exit(0)\n\n\nif __name__ == \"__main__\":\n    main()\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\nfrom datetime import date\n\nfrom tabulate import tabulate\n\n\nMAX_LEN_MESSAGE = 2900  # slack endpoint has a limit of 3001 characters\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--slack_channel_name\", default=\"trl-push-examples-ci\")\nparser.add_argument(\"--text_file_name\", required=True)\n\n\ndef main(text_file_name, slack_channel_name=None):\n    message = \"\"\n\n    if os.path.isfile(text_file_name):\n        final_results = {}\n\n        file = open(text_file_name)\n        lines = file.readlines()\n        for line in lines:\n            result, config_name = line.split(\",\")\n            config_name = config_name.split(\"/\")[-1].split(\".yaml\")[0]\n            final_results[config_name] = int(result)\n\n        no_error_payload = {\n            \"type\": \"section\",\n            \"text\": {\n                \"type\": \"plain_text\",\n                \"text\": \"🌞 There were no failures on the example tests!\"\n                if not len(final_results) == 0\n                else \"Something went wrong there is at least one empty file - please check GH action results.\",\n                \"emoji\": True,\n            },\n        }\n\n        total_num_failed = sum(final_results.values())\n    else:\n        no_error_payload = {\n            \"type\": \"section\",\n            \"text\": {\n                \"type\": \"plain_text\",\n                \"text\": \"🔴 Something is wrong with the workflow please check ASAP!\"\n                \"Something went wrong there is no text file being produced. Please check ASAP.\",\n                \"emoji\": True,\n            },\n        }\n\n        total_num_failed = 0\n\n    test_type_name = text_file_name.replace(\".txt\", \"\").replace(\"temp_results_\", \"\").replace(\"_\", \" \").title()\n\n    payload = [\n        {\n            \"type\": \"header\",\n            \"text\": {\n                \"type\": \"plain_text\",\n                \"text\": \"🤗 Results of the {} TRL {} example tests.\".format(\n                    os.environ.get(\"TEST_TYPE\", \"\"), test_type_name\n                ),\n            },\n        },\n    ]\n\n    if total_num_failed > 0:\n        message += f\"{total_num_failed} failed tests for example tests!\"\n\n        for test_name, failed in final_results.items():\n            failed_table = tabulate(\n                [[test_name, \"🟢\" if not failed else \"🔴\"]],\n                headers=[\"Test Name\", \"Status\"],\n                showindex=\"always\",\n                tablefmt=\"grid\",\n                maxcolwidths=[12],\n            )\n            message += \"\\n```\\n\" + failed_table + \"\\n```\"\n\n        print(f\"### {message}\")\n    else:\n        payload.append(no_error_payload)\n\n    if os.environ.get(\"TEST_TYPE\", \"\") != \"\":\n        from slack_sdk import WebClient\n\n        if len(message) > MAX_LEN_MESSAGE:\n            print(f\"Truncating long message from {len(message)} to {MAX_LEN_MESSAGE}\")\n            message = message[:MAX_LEN_MESSAGE] + \"...\"\n\n        if len(message) != 0:\n            md_report = {\n                \"type\": \"section\",\n                \"text\": {\"type\": \"mrkdwn\", \"text\": message},\n            }\n            payload.append(md_report)\n            action_button = {\n                \"type\": \"section\",\n                \"text\": {\"type\": \"mrkdwn\", \"text\": \"*For more details:*\"},\n                \"accessory\": {\n                    \"type\": \"button\",\n                    \"text\": {\"type\": \"plain_text\", \"text\": \"Check Action results\", \"emoji\": True},\n                    \"url\": f\"https://github.com/huggingface/trl/actions/runs/{os.environ['GITHUB_RUN_ID']}\",\n                },\n            }\n            payload.append(action_button)\n\n        date_report = {\n            \"type\": \"context\",\n            \"elements\": [\n                {\n                    \"type\": \"plain_text\",\n                    \"text\": f\"On Push - main {os.environ.get('TEST_TYPE')} test results for {date.today()}\",\n                },\n            ],\n        }\n        payload.append(date_report)\n\n        print(payload)\n\n        client = WebClient(token=os.environ.get(\"SLACK_API_TOKEN\"))\n        client.chat_postMessage(channel=f\"#{slack_channel_name}\", text=message, blocks=payload)\n\n\nif __name__ == \"__main__\":\n    args = parser.parse_args()\n    main(args.text_file_name, args.slack_channel_name)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport json\nimport os\nfrom datetime import date\nfrom pathlib import Path\n\nfrom tabulate import tabulate\n\n\nMAX_LEN_MESSAGE = 2900  # slack endpoint has a limit of 3001 characters\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--slack_channel_name\", default=\"trl-push-ci\")\n\n\ndef main(slack_channel_name=None):\n    failed = []\n    passed = []\n\n    group_info = []\n\n    total_num_failed = 0\n    empty_file = False or len(list(Path().glob(\"*.log\"))) == 0\n\n    total_empty_files = []\n\n    for log in Path().glob(\"*.log\"):\n        section_num_failed = 0\n        i = 0\n        with open(log) as f:\n            for line in f:\n                line = json.loads(line)\n                i += 1\n                if line.get(\"nodeid\", \"\") != \"\":\n                    test = line[\"nodeid\"]\n                    if line.get(\"duration\", None) is not None:\n                        duration = f'{line[\"duration\"]:.4f}'\n                        if line.get(\"outcome\", \"\") == \"failed\":\n                            section_num_failed += 1\n                            failed.append([test, duration, log.name.split(\"_\")[0]])\n                            total_num_failed += 1\n                        else:\n                            passed.append([test, duration, log.name.split(\"_\")[0]])\n            empty_file = i == 0\n        group_info.append([str(log), section_num_failed, failed])\n        total_empty_files.append(empty_file)\n        os.remove(log)\n        failed = []\n    no_error_payload = {\n        \"type\": \"section\",\n        \"text\": {\n            \"type\": \"plain_text\",\n            \"text\": \"🌞 There were no failures!\"\n            if not any(total_empty_files)\n            else \"Something went wrong there is at least one empty file - please check GH action results.\",\n            \"emoji\": True,\n        },\n    }\n\n    message = \"\"\n    payload = [\n        {\n            \"type\": \"header\",\n            \"text\": {\n                \"type\": \"plain_text\",\n                \"text\": \"🤗 Results of the {} TRL tests.\".format(os.environ.get(\"TEST_TYPE\", \"\")),\n            },\n        },\n    ]\n    if total_num_failed > 0:\n        for i, (name, num_failed, failed_tests) in enumerate(group_info):\n            if num_failed > 0:\n                if num_failed == 1:\n                    message += f\"*{name}: {num_failed} failed test*\\n\"\n                else:\n                    message += f\"*{name}: {num_failed} failed tests*\\n\"\n                failed_table = []\n                for test in failed_tests:\n                    failed_report = test[0].split(\"::\")\n                    # Truncate the last string as some test names might be long\n                    failed_report[-1] = failed_report[-1][:30] + \"..\"\n                    failed_table.append(failed_report)\n                failed_table = tabulate(\n                    failed_table,\n                    headers=[\"Test Location\", \"Test Case\", \"Test Name\"],\n                    showindex=\"always\",\n                    tablefmt=\"grid\",\n                    maxcolwidths=[12, 12, 12],\n                )\n                message += \"\\n```\\n\" + failed_table + \"\\n```\"\n\n            if total_empty_files[i]:\n                message += f\"\\n*{name}: Warning! Empty file - please check the GitHub action job *\\n\"\n        print(f\"### {message}\")\n    else:\n        payload.append(no_error_payload)\n\n    if os.environ.get(\"TEST_TYPE\", \"\") != \"\":\n        from slack_sdk import WebClient\n\n        if len(message) > MAX_LEN_MESSAGE:\n            message = f\"There are {total_num_failed} failed tests in total ! Cannot display the entire summary - please check the action results directly\"\n\n        if len(message) != 0:\n            md_report = {\n                \"type\": \"section\",\n                \"text\": {\"type\": \"mrkdwn\", \"text\": message},\n            }\n            payload.append(md_report)\n            action_button = {\n                \"type\": \"section\",\n                \"text\": {\"type\": \"mrkdwn\", \"text\": \"*For more details:*\"},\n                \"accessory\": {\n                    \"type\": \"button\",\n                    \"text\": {\"type\": \"plain_text\", \"text\": \"Check Action results\", \"emoji\": True},\n                    \"url\": f\"https://github.com/huggingface/trl/actions/runs/{os.environ['GITHUB_RUN_ID']}\",\n                },\n            }\n            payload.append(action_button)\n\n        date_report = {\n            \"type\": \"context\",\n            \"elements\": [\n                {\n                    \"type\": \"plain_text\",\n                    \"text\": f\"On Push main {os.environ.get('TEST_TYPE')} test results for {date.today()}\",\n                },\n            ],\n        }\n        payload.append(date_report)\n\n        print(payload)\n\n        client = WebClient(token=os.environ.get(\"SLACK_API_TOKEN\"))\n        client.chat_postMessage(channel=f\"#{slack_channel_name}\", text=message, blocks=payload)\n\n\nif __name__ == \"__main__\":\n    args = parser.parse_args()\n    main(args.slack_channel_name)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\n\nfrom ghapi.all import GhApi\n\n\nFOLDER_STRING = os.environ.get(\"FOLDER_STRING\", \"\")\nfolder = f\"benchmark/trl/{FOLDER_STRING}\"\nhost_url = f\"https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/{FOLDER_STRING}\"\n\n# Create a GitHub API instance\ngithub_context = json.loads(os.environ[\"GITHUB_CONTEXT\"])\ntoken = os.environ[\"PERSONAL_ACCESS_TOKEN_GITHUB\"]  # this needs to refreshed every 12 months\nstatus_message = \"**[COSTA BENCHMARK BOT]**: Here are the results\"\nbody = status_message\nrepo = github_context[\"repository\"]\nowner, repo = repo.split(\"/\")\napi = GhApi(owner=owner, repo=repo, token=token)\n\n# for each `.png` file in the folder, add it to the comment\nfor file in os.listdir(folder):\n    if file.endswith(\".png\"):\n        body += f\"\\n![{file}]({host_url}/{file})\"\n\n# Create a comment on the issue\napi.issues.create_comment(issue_number=github_context[\"event\"][\"issue\"][\"number\"], body=body)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\n\nimport tyro\nfrom huggingface_hub import HfApi\n\n\n@dataclass\nclass Args:\n    folder_path: str = \"benchmark/trl\"\n    path_in_repo: str = \"images/benchmark\"\n    repo_id: str = \"trl-internal-testing/example-images\"\n    repo_type: str = \"dataset\"\n\n\nargs = tyro.cli(Args)\napi = HfApi()\n\napi.upload_folder(\n    folder_path=args.folder_path,\n    path_in_repo=args.path_in_repo,\n    repo_id=args.repo_id,\n    repo_type=args.repo_type,\n)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport math\nimport os\nimport shlex\nimport subprocess\nimport uuid\nfrom distutils.util import strtobool\n\nimport requests\n\n\ndef parse_args():\n    # fmt: off\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--command\", type=str, default=\"\",\n        help=\"the command to run\")\n    parser.add_argument(\"--num-seeds\", type=int, default=3,\n        help=\"the number of random seeds\")\n    parser.add_argument(\"--start-seed\", type=int, default=1,\n        help=\"the number of the starting seed\")\n    parser.add_argument(\"--workers\", type=int, default=0,\n        help=\"the number of workers to run benchmark experimenets\")\n    parser.add_argument(\"--auto-tag\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n        help=\"if toggled, the runs will be tagged with git tags, commit, and pull request number if possible\")\n    parser.add_argument(\"--slurm-template-path\", type=str, default=None,\n        help=\"the path to the slurm template file (see docs for more details)\")\n    parser.add_argument(\"--slurm-gpus-per-task\", type=int, default=1,\n        help=\"the number of gpus per task to use for slurm jobs\")\n    parser.add_argument(\"--slurm-total-cpus\", type=int, default=50,\n        help=\"the number of gpus per task to use for slurm jobs\")\n    parser.add_argument(\"--slurm-ntasks\", type=int, default=1,\n        help=\"the number of tasks to use for slurm jobs\")\n    parser.add_argument(\"--slurm-nodes\", type=int, default=None,\n        help=\"the number of nodes to use for slurm jobs\")\n    args = parser.parse_args()\n    # fmt: on\n    return args\n\n\ndef run_experiment(command: str):\n    command_list = shlex.split(command)\n    print(f\"running {command}\")\n\n    # Use subprocess.PIPE to capture the output\n    fd = subprocess.Popen(command_list, stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n    output, errors = fd.communicate()\n\n    return_code = fd.returncode\n    assert return_code == 0, f\"Command failed with error: {errors.decode('utf-8')}\"\n\n    # Convert bytes to string and strip leading/trailing whitespaces\n    return output.decode(\"utf-8\").strip()\n\n\ndef autotag() -> str:\n    wandb_tag = \"\"\n    print(\"autotag feature is enabled\")\n    git_tag = \"\"\n    try:\n        git_tag = subprocess.check_output([\"git\", \"describe\", \"--tags\"]).decode(\"ascii\").strip()\n        print(f\"identified git tag: {git_tag}\")\n    except subprocess.CalledProcessError as e:\n        print(e)\n    if len(git_tag) == 0:\n        try:\n            count = int(subprocess.check_output([\"git\", \"rev-list\", \"--count\", \"HEAD\"]).decode(\"ascii\").strip())\n            hash = subprocess.check_output([\"git\", \"rev-parse\", \"--short\", \"HEAD\"]).decode(\"ascii\").strip()\n            git_tag = f\"no-tag-{count}-g{hash}\"\n            print(f\"identified git tag: {git_tag}\")\n        except subprocess.CalledProcessError as e:\n            print(e)\n    wandb_tag = f\"{git_tag}\"\n\n    git_commit = subprocess.check_output([\"git\", \"rev-parse\", \"--verify\", \"HEAD\"]).decode(\"ascii\").strip()\n    try:\n        # try finding the pull request number on github\n        prs = requests.get(f\"https://api.github.com/search/issues?q=repo:huggingface/trl+is:pr+{git_commit}\")\n        if prs.status_code == 200:\n            prs = prs.json()\n            if len(prs[\"items\"]) > 0:\n                pr = prs[\"items\"][0]\n                pr_number = pr[\"number\"]\n                wandb_tag += f\",pr-{pr_number}\"\n        print(f\"identified github pull request: {pr_number}\")\n    except Exception as e:\n        print(e)\n\n    return wandb_tag\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    if args.auto_tag:\n        existing_wandb_tag = os.environ.get(\"WANDB_TAGS\", \"\")\n        wandb_tag = autotag()\n        if len(wandb_tag) > 0:\n            if len(existing_wandb_tag) > 0:\n                os.environ[\"WANDB_TAGS\"] = \",\".join([existing_wandb_tag, wandb_tag])\n            else:\n                os.environ[\"WANDB_TAGS\"] = wandb_tag\n    print(\"WANDB_TAGS: \", os.environ.get(\"WANDB_TAGS\", \"\"))\n    commands = []\n    for seed in range(0, args.num_seeds):\n        commands += [\" \".join([args.command, \"--seed\", str(args.start_seed + seed)])]\n\n    print(\"======= commands to run:\")\n    for command in commands:\n        print(command)\n\n    if args.workers > 0 and args.slurm_template_path is None:\n        from concurrent.futures import ThreadPoolExecutor\n\n        executor = ThreadPoolExecutor(max_workers=args.workers, thread_name_prefix=\"cleanrl-benchmark-worker-\")\n        for command in commands:\n            executor.submit(run_experiment, command)\n        executor.shutdown(wait=True)\n    else:\n        print(\"not running the experiments because --workers is set to 0; just printing the commands to run\")\n\n    # SLURM logic\n    if args.slurm_template_path is not None:\n        if not os.path.exists(\"slurm\"):\n            os.makedirs(\"slurm\")\n        if not os.path.exists(\"slurm/logs\"):\n            os.makedirs(\"slurm/logs\")\n        print(\"======= slurm commands to run:\")\n        with open(args.slurm_template_path) as f:\n            slurm_template = f.read()\n        slurm_template = slurm_template.replace(\"{{array}}\", f\"0-{len(commands) - 1}%{args.workers}\")\n        slurm_template = slurm_template.replace(\n            \"{{seeds}}\", f\"({' '.join([str(args.start_seed + int(seed)) for seed in range(args.num_seeds)])})\"\n        )\n        slurm_template = slurm_template.replace(\"{{len_seeds}}\", f\"{args.num_seeds}\")\n        slurm_template = slurm_template.replace(\"{{command}}\", args.command)\n        slurm_template = slurm_template.replace(\"{{gpus_per_task}}\", f\"{args.slurm_gpus_per_task}\")\n        total_gpus = args.slurm_gpus_per_task * args.slurm_ntasks\n        slurm_cpus_per_gpu = math.ceil(args.slurm_total_cpus / total_gpus)\n        slurm_template = slurm_template.replace(\"{{cpus_per_gpu}}\", f\"{slurm_cpus_per_gpu}\")\n        slurm_template = slurm_template.replace(\"{{ntasks}}\", f\"{args.slurm_ntasks}\")\n        if args.slurm_nodes is not None:\n            slurm_template = slurm_template.replace(\"{{nodes}}\", f\"#SBATCH --nodes={args.slurm_nodes}\")\n        else:\n            slurm_template = slurm_template.replace(\"{{nodes}}\", \"\")\n        filename = str(uuid.uuid4())\n        open(os.path.join(\"slurm\", f\"{filename}.slurm\"), \"w\").write(slurm_template)\n        slurm_path = os.path.join(\"slurm\", f\"{filename}.slurm\")\n        print(f\"saving command in {slurm_path}\")\n        if args.workers > 0:\n            job_id = run_experiment(f\"sbatch --parsable {slurm_path}\")\n            print(f\"Job ID: {job_id}\")\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport importlib\nimport os\nimport sys\nfrom itertools import chain\nfrom types import ModuleType\nfrom typing import Any\n\nfrom transformers.utils.import_utils import _is_package_available\n\n\nif sys.version_info < (3, 8):\n    _is_python_greater_3_8 = False\nelse:\n    _is_python_greater_3_8 = True\n\n# Use same as transformers.utils.import_utils\n_diffusers_available = _is_package_available(\"diffusers\")\n_unsloth_available = _is_package_available(\"unsloth\")\n_rich_available = _is_package_available(\"rich\")\n_liger_kernel_available = _is_package_available(\"liger_kernel\")\n_llmblender_available = _is_package_available(\"llm_blender\")\n\n\ndef is_diffusers_available() -> bool:\n    return _diffusers_available\n\n\ndef is_unsloth_available() -> bool:\n    return _unsloth_available\n\n\ndef is_rich_available() -> bool:\n    return _rich_available\n\n\ndef is_liger_kernel_available() -> bool:  # replace by transformers.import_utils.is_liger_kernel_available() from v4.45\n    return _liger_kernel_available\n\n\ndef is_llmblender_available() -> bool:\n    return _llmblender_available\n\n\ndef is_accelerate_greater_20_0() -> bool:\n    if _is_python_greater_3_8:\n        from importlib.metadata import version\n\n        accelerate_version = version(\"accelerate\")\n    else:\n        import pkg_resources\n\n        accelerate_version = pkg_resources.get_distribution(\"accelerate\").version\n    return accelerate_version >= \"0.20.0\"\n\n\ndef is_transformers_greater_than(current_version: str) -> bool:\n    if _is_python_greater_3_8:\n        from importlib.metadata import version\n\n        _transformers_version = version(\"transformers\")\n    else:\n        import pkg_resources\n\n        _transformers_version = pkg_resources.get_distribution(\"transformers\").version\n    return _transformers_version > current_version\n\n\ndef is_torch_greater_2_0() -> bool:\n    if _is_python_greater_3_8:\n        from importlib.metadata import version\n\n        torch_version = version(\"torch\")\n    else:\n        import pkg_resources\n\n        torch_version = pkg_resources.get_distribution(\"torch\").version\n    return torch_version >= \"2.0\"\n\n\nclass _LazyModule(ModuleType):\n    \"\"\"\n    Module class that surfaces all objects but only performs associated imports when the objects are requested.\n    \"\"\"\n\n    # Very heavily inspired by optuna.integration._IntegrationModule\n    # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py\n    def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):\n        super().__init__(name)\n        self._modules = set(import_structure.keys())\n        self._class_to_module = {}\n        for key, values in import_structure.items():\n            for value in values:\n                self._class_to_module[value] = key\n        # Needed for autocompletion in an IDE\n        self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values()))\n        self.__file__ = module_file\n        self.__spec__ = module_spec\n        self.__path__ = [os.path.dirname(module_file)]\n        self._objects = {} if extra_objects is None else extra_objects\n        self._name = name\n        self._import_structure = import_structure\n\n    # Needed for autocompletion in an IDE\n    def __dir__(self):\n        result = super().__dir__()\n        # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether\n        # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir.\n        for attr in self.__all__:\n            if attr not in result:\n                result.append(attr)\n        return result\n\n    def __getattr__(self, name: str) -> Any:\n        if name in self._objects:\n            return self._objects[name]\n        if name in self._modules:\n            value = self._get_module(name)\n        elif name in self._class_to_module.keys():\n            module = self._get_module(self._class_to_module[name])\n            value = getattr(module, name)\n        else:\n            raise AttributeError(f\"module {self.__name__} has no attribute {name}\")\n\n        setattr(self, name, value)\n        return value\n\n    def _get_module(self, module_name: str):\n        try:\n            return importlib.import_module(\".\" + module_name, self.__name__)\n        except Exception as e:\n            raise RuntimeError(\n                f\"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its\"\n                f\" traceback):\\n{e}\"\n            ) from e\n\n    def __reduce__(self):\n        return (self.__class__, (self._name, self.__file__, self._import_structure))\n\n\nclass OptionalDependencyNotAvailable(BaseException):\n    \"\"\"Internally used error class for signalling an optional dependency was not found.\"\"\"\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# Function `strtobool` copied and adapted from `distutils` (as deprected\n# in Python 3.10).\n# Reference: https://github.com/python/cpython/blob/48f9d3e3faec5faaa4f7c9849fecd27eae4da213/Lib/distutils/util.py#L308-L321\n\n\ndef strtobool(val: str) -> bool:\n    \"\"\"Convert a string representation of truth to True or False booleans.\n\n    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values\n    are 'n', 'no', 'f', 'false', 'off', and '0'.\n\n    Raises:\n        ValueError: if 'val' is anything else.\n    \"\"\"\n    val = val.lower()\n    if val in (\"y\", \"yes\", \"t\", \"true\", \"on\", \"1\"):\n        return True\n    if val in (\"n\", \"no\", \"f\", \"false\", \"off\", \"0\"):\n        return False\n    raise ValueError(f\"Invalid truth value, it should be a string but {val} was provided instead.\")\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Dict, List, Optional, TypeVar\n\nfrom datasets import Dataset, DatasetDict\nfrom transformers import PreTrainedTokenizer\n\n\nDatasetType = TypeVar(\"DatasetType\", Dataset, DatasetDict)\n\n\ndef is_conversational(example: Dict[str, Any]) -> bool:\n    r\"\"\"\n    Check if the example is in a conversational format.\n\n    Args:\n        example (`Dict[str, Any]`):\n            A single data entry of a dataset. The example can have different keys depending on the\n            dataset format.\n\n    Returns:\n        `bool`: `True` if the data is in a conversational format, `False` otherwise.\n\n    Examples:\n\n    ```python\n    >>> example = {\"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}]}\n    >>> is_conversational(example)\n    True\n    >>> example = {\"prompt\": \"The sky is\"})\n    >>> is_conversational(example)\n    False\n    ```\n    \"\"\"\n    supported_keys = [\"prompt\", \"chosen\", \"rejected\", \"completion\", \"messages\"]\n    example_keys = {key for key in example.keys() if key in supported_keys}\n\n    # It must have one of the supported keys\n    if example_keys:\n        key = example_keys.pop()  # take the first supported key\n        maybe_messages = example[key]\n        # It must be a list of messages,\n        if isinstance(maybe_messages, list):\n            maybe_message = maybe_messages[0]\n            # Each message must a list of dictionaries with keys \"role\" and \"content\"\n            if isinstance(maybe_message, dict) and \"role\" in maybe_message and \"content\" in maybe_message:\n                return True\n\n    return False\n\n\ndef apply_chat_template(example: Dict[str, List[Dict[str, str]]], tokenizer: PreTrainedTokenizer) -> Dict[str, str]:\n    r\"\"\"\n    Apply a chat template to a conversational example.\n\n    For more details, see [`maybe_apply_chat_template`].\n    \"\"\"\n    # Check that the example has the correct keys\n    supported_keys = [\"prompt\", \"chosen\", \"rejected\", \"completion\", \"messages\", \"label\"]\n    example_keys = {key for key in example.keys() if key in supported_keys}\n    if example_keys not in [\n        {\"messages\"},  # language modeling\n        {\"prompt\"},  # prompt-only\n        {\"prompt\", \"completion\"},  # prompt-completion\n        {\"prompt\", \"chosen\", \"rejected\"},  # preference\n        {\"chosen\", \"rejected\"},  # preference with implicit prompt\n        {\"prompt\", \"completion\", \"label\"},  # unpaired preference\n    ]:\n        raise KeyError(f\"Invalid keys in the example: {example_keys}\")\n\n    # Apply the chat template to the whole conversation\n    if \"messages\" in example:\n        messages = tokenizer.apply_chat_template(example[\"messages\"], tokenize=False)\n\n    # Apply the chat template to the prompt, adding the generation prompt\n    if \"prompt\" in example:\n        prompt = tokenizer.apply_chat_template(example[\"prompt\"], tokenize=False, add_generation_prompt=True)\n\n    # Apply the chat template to the entire prompt + completion\n    if \"prompt\" in example:  # explicit prompt and prompt-completion case\n        if \"chosen\" in example:\n            prompt_chosen = tokenizer.apply_chat_template(example[\"prompt\"] + example[\"chosen\"], tokenize=False)\n            chosen = prompt_chosen[len(prompt) :]\n        if \"rejected\" in example and \"prompt\" in example:  # explicit prompt\n            prompt_rejected = tokenizer.apply_chat_template(example[\"prompt\"] + example[\"rejected\"], tokenize=False)\n            rejected = prompt_rejected[len(prompt) :]\n        if \"completion\" in example:\n            prompt_completion = tokenizer.apply_chat_template(\n                example[\"prompt\"] + example[\"completion\"], tokenize=False\n            )\n            completion = prompt_completion[len(prompt) :]\n    else:  # implicit prompt case\n        if \"chosen\" in example:\n            chosen = tokenizer.apply_chat_template(example[\"chosen\"], tokenize=False)\n        if \"rejected\" in example:\n            rejected = tokenizer.apply_chat_template(example[\"rejected\"], tokenize=False)\n\n    # Ensure that the prompt is the initial part of the prompt-completion string\n    if \"prompt\" in example:\n        error_message = (\n            \"The chat template applied to the prompt + completion does not start with the chat template applied to \"\n            \"the prompt alone. This can indicate that the chat template is not supported by TRL.\"\n            \"\\n**Prompt**:\\n{}\\n\\n**Prompt + Completion**:\\n{}\"\n        )\n        if \"chosen\" in example and not prompt_chosen.startswith(prompt):\n            raise ValueError(error_message.format(prompt, prompt_chosen))\n        if \"rejected\" in example and not prompt_rejected.startswith(prompt):\n            raise ValueError(error_message.format(prompt, prompt_rejected))\n        if \"completion\" in example and not prompt_completion.startswith(prompt):\n            raise ValueError(error_message.format(prompt, prompt_completion))\n\n    # Extract the completion by removing the prompt part from the prompt-completion string\n    output = {}\n    if \"messages\" in example:\n        output[\"text\"] = messages\n    if \"prompt\" in example:\n        output[\"prompt\"] = prompt\n    if \"chosen\" in example:\n        output[\"chosen\"] = chosen\n    if \"rejected\" in example:\n        output[\"rejected\"] = rejected\n    if \"completion\" in example:\n        output[\"completion\"] = completion\n    if \"label\" in example:\n        output[\"label\"] = example[\"label\"]\n\n    return output\n\n\ndef maybe_apply_chat_template(\n    example: Dict[str, List[Dict[str, str]]], tokenizer: PreTrainedTokenizer\n) -> Dict[str, str]:\n    r\"\"\"\n    If the example is in a conversational format, apply a chat template to it.\n\n    Args:\n        example (`Dict[str, List[Dict[str, str]]`):\n            Dictionary representing a single data entry of a conversational dataset. Each data entry can have different\n            keys depending on the dataset format. The supported dataset formats are:\n\n                - Language modeling dataset: `\"messages\"`.\n                - Prompt-only dataset: `\"prompt\"`.\n                - Prompt-completion dataset: `\"prompt\"` and `\"completion\"`.\n                - Preference dataset: `\"prompt\"`, `\"chosen\"`, and `\"rejected\"`.\n                - Preference dataset with implicit prompt: `\"chosen\"` and `\"rejected\"`.\n                - Unpaired preference dataset: `\"prompt\"`, `\"completion\"`, and `\"label\"`.\n\n            For keys `\"messages\"`, `\"prompt\"`, `\"chosen\"`, `\"rejected\"`, and `\"completion\"`, the values are lists of\n            messages, where each message is a dictionary with keys `\"role\"` and `\"content\"`.\n\n        tokenizer (`PreTrainedTokenizer`):\n            The tokenizer to apply the chat template with.\n\n    Returns:\n        `Dict[str, str]`: The formatted example with the chat template applied.\n\n    Note:\n        This function does not alter the keys, except for Language modeling dataset, where `\"messages\"` is replaced by\n        `\"text\"`.\n\n    Example:\n\n    ```python\n    >>> from transformers import AutoTokenizer\n    >>> tokenizer = AutoTokenizer.from_pretrained(\"microsoft/Phi-3-mini-128k-instruct\")\n    >>> example = {\n    ...     \"prompt\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}],\n    ...     \"completion\": [{\"role\": \"assistant\", \"content\": \"It is blue.\"}]\n    ... }\n    >>> apply_chat_template(example, tokenizer)\n    {'prompt': '<|user|>\\nWhat color is the sky?<|end|>\\n<|assistant|>\\n', 'completion': 'It is blue.<|end|>\\n<|endoftext|>'}\n    ```\n    \"\"\"\n    if is_conversational(example):\n        return apply_chat_template(example, tokenizer)\n    else:\n        return example\n\n\ndef _unpair_row(examples: List[Dict[str, List[Dict[str, str]]]]) -> List[Dict[str, List[Dict[str, str]]]]:\n    batch_size = len(examples[\"chosen\"])\n    new_rows = {\n        \"completion\": examples[\"chosen\"] + examples[\"rejected\"],\n        \"label\": [True] * batch_size + [False] * batch_size,\n    }\n    if \"prompt\" in examples:\n        new_rows[\"prompt\"] = examples[\"prompt\"] + examples[\"prompt\"]\n    return new_rows\n\n\ndef unpair_preference_dataset(dataset: DatasetType, num_proc: Optional[int] = None) -> DatasetType:\n    r\"\"\"\n    Unpair a preference dataset.\n\n    Args:\n        dataset (`Dataset` or `DatasetDict`):\n            Preference dataset to unpair. The dataset must have columns `\"chosen\"`, `\"rejected\"` and optionally\n            `\"prompt\"`.\n        num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n\n    Returns:\n        `Dataset`: The unpaired preference dataset.\n\n    Example:\n\n    ```python\n    >>> from datasets import Dataset\n    >>> dataset_dict = {\n    ...     \"prompt\": [\"The sky is\", \"The sun is\"]\n    ...     \"chosen\": [\" blue.\", \"in the sky.\"],\n    ...     \"rejected\": [\" green.\", \" in the sea.\"]\n    ... }\n    >>> dataset = Dataset.from_dict(dataset_dict)\n    >>> dataset = unpair_preference_dataset(dataset)\n    >>> dataset\n    Dataset({\n        features: ['prompt', 'completion', 'label'],\n        num_rows: 4\n    })\n    >>> dataset[0]\n    {'prompt': 'The sky is', 'completion': ' blue.', 'label': True}\n    ```\n    \"\"\"\n    return dataset.map(_unpair_row, batched=True, remove_columns=[\"chosen\", \"rejected\"], num_proc=num_proc)\n\n\ndef maybe_unpair_preference_dataset(dataset: DatasetType, num_proc: Optional[int] = None) -> DatasetType:\n    r\"\"\"\n    Unpair a preference dataset if it is paired.\n\n    Args:\n        dataset (`Dataset` or `DatasetDict`):\n            Preference dataset to unpair. The dataset must have columns `\"chosen\"`, `\"rejected\"` and optionally\n            `\"prompt\"`.\n        num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n\n    Returns:\n        `Dataset` or `DatasetDict`: The unpaired preference dataset if it was paired, otherwise the original dataset.\n\n    Example:\n\n    ```python\n    >>> from datasets import Dataset\n    >>> dataset_dict = {\n    ...     \"prompt\": [\"The sky is\", \"The sun is\"]\n    ...     \"chosen\": [\" blue.\", \"in the sky.\"],\n    ...     \"rejected\": [\" green.\", \" in the sea.\"]\n    ... }\n    >>> dataset = Dataset.from_dict(dataset_dict)\n    >>> dataset = unpair_preference_dataset(dataset)\n    >>> dataset\n    Dataset({\n        features: ['prompt', 'completion', 'label'],\n        num_rows: 4\n    })\n    >>> dataset[0]\n    {'prompt': 'The sky is', 'completion': ' blue.', 'label': True}\n    ```\n    \"\"\"\n    if isinstance(dataset, DatasetDict):\n        column_names = dataset[list(dataset.keys())[0]].column_names\n    else:\n        column_names = dataset.column_names\n    if \"chosen\" in column_names and \"rejected\" in column_names:\n        return unpair_preference_dataset(dataset, num_proc=num_proc)\n    else:\n        return dataset\n\n\ndef extract_prompt(example: Dict[str, List]) -> Dict[str, List]:\n    r\"\"\"\n    Extracts the shared prompt from a preference data example, where the prompt is implicit within both\n    the chosen and rejected completions.\n\n    For more details, see [`maybe_extract_prompt`].\n    \"\"\"\n    for idx in range(min(len(example[\"chosen\"]), len(example[\"rejected\"]))):\n        if example[\"chosen\"][idx][\"content\"] != example[\"rejected\"][idx][\"content\"]:\n            break\n    return {\n        \"prompt\": example[\"chosen\"][:idx],\n        \"chosen\": example[\"chosen\"][idx:],\n        \"rejected\": example[\"rejected\"][idx:],\n    }\n\n\ndef maybe_extract_prompt(example: Dict[str, List]) -> Dict[str, List]:\n    r\"\"\"\n    Extracts the shared prompt from a preference data example, where the prompt is implicit within both\n    the chosen and rejected completions.\n\n    If the example already contains a `\"prompt\"` key, the function returns the example as is. Else, the function\n\n    identifies the longest common sequence (prefix) of conversation turns between the \"chosen\" and \"rejected\"\n    completions and extracts this as the prompt. It then removes this prompt from the respective \"chosen\" and\n    \"rejected\" completions.\n\n    Args:\n        example (`Dict[str, List]`):\n            A dictionary representing a single data entry in the preference dataset. It must contain the keys\n            `\"chosen\"` and `\"rejected\"`, where each value is a list.\n\n    Returns:\n        `Dict[str, List]`: A dictionary containing:\n            - `\"prompt\"`: The longest common prefix between the \"chosen\" and \"rejected\" completions.\n            - `\"chosen\"`: The remainder of the \"chosen\" completion, with the prompt removed.\n            - `\"rejected\"`: The remainder of the \"rejected\" completion, with the prompt removed.\n\n    Examples:\n\n    ```python\n    >>> example = {\n    ...     \"chosen\": [\n    ...         {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n    ...         {\"role\": \"assistant\", \"content\": \"It is blue.\"}\n    ...     ],\n    ...     \"rejected\": [\n    ...         {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n    ...         {\"role\": \"assistant\", \"content\": \"It is green.\"}\n    ...     ]\n    ... }\n    >>> extract_prompt(example)\n    {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],\n     'chosen': [{'role': 'assistant', 'content': 'It is blue.'}],\n     'rejected': [{'role': 'assistant', 'content': 'It is green.'}]}\n    ```\n\n    Or, with the `map` method of `datasets.Dataset`:\n\n    ```python\n    >>> from trl import extract_prompt\n    >>> from datasets import Dataset\n    >>> dataset_dict = {\n    ...     \"chosen\": [\n    ...         [\n    ...             {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n    ...             {\"role\": \"assistant\", \"content\": \"It is blue.\"},\n    ...         ],\n    ...         [\n    ...             {\"role\": \"user\", \"content\": \"Where is the sun?\"},\n    ...             {\"role\": \"assistant\", \"content\": \"In the sky.\"},\n    ...         ],\n    ...     ],\n    ...     \"rejected\": [\n    ...         [\n    ...             {\"role\": \"user\", \"content\": \"What color is the sky?\"},\n    ...             {\"role\": \"assistant\", \"content\": \"It is green.\"},\n    ...         ],\n    ...         [\n    ...             {\"role\": \"user\", \"content\": \"Where is the sun?\"},\n    ...             {\"role\": \"assistant\", \"content\": \"In the sea.\"},\n    ...         ],\n    ...     ],\n    ... }\n    >>> dataset = Dataset.from_dict(dataset_dict)\n    >>> dataset = dataset.map(extract_prompt)\n    >>> dataset[0]\n    {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],\n     'chosen': [{'role': 'assistant', 'content': 'It is blue.'}],\n     'rejected': [{'role': 'assistant', 'content': 'It is green.'}]}\n    ```\n    \"\"\"\n    # Some dataset add a `\"prompt\"` column, even though the prompt is implicit and included in the \"chosen\" and\n    # \"rejected\" completions. E.g.:\n    # {\"prompt\": \"What color is the sky?\",\n    #  \"chosen\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}, {\"role\": \"assistant\", \"content\": \"It is blue.\"}],\n    #  \"rejected\": [{\"role\": \"user\", \"content\": \"What color is the sky?\"}, {\"role\": \"assistant\", \"content\": \"It is green.\"}]}\n    # That's why we check if the prompt is also conversational before deciding not to extract it.\n    if \"prompt\" in example and is_conversational({\"prompt\": example[\"prompt\"]}):\n        return example\n    else:\n        return extract_prompt({\"chosen\": example[\"chosen\"], \"rejected\": example[\"rejected\"]})\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport random\nimport warnings\nfrom contextlib import contextmanager\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.rnn import pad_sequence\nfrom transformers import TopKLogitsWarper, TopPLogitsWarper, is_torch_npu_available, is_torch_xpu_available\n\n\ntry:\n    from collections.abc import Mapping\nexcept ImportError:\n    from collections.abc import Mapping\n\n\nWANDB_PADDING = -1\n\n\ndef top_k_top_p_filtering(\n    logits: torch.FloatTensor,\n    top_k: int = 0,\n    top_p: float = 1.0,\n    filter_value: float = -float(\"Inf\"),\n    min_tokens_to_keep: int = 1,\n) -> torch.FloatTensor:\n    \"\"\"\n    Filter a distribution of logits using top-k and/or nucleus (top-p) filtering.\n\n    Args:\n        logits: logits distribution shape (batch size, vocabulary size)\n        top_k (`int`, *optional*, defaults to 0):\n            If > 0, only keep the top k tokens with highest probability (top-k filtering)\n        top_p (`float`, *optional*, defaults to 1.0):\n            If < 1.0, only keep the top tokens with cumulative probability >= top_p (nucleus filtering). Nucleus\n            filtering is described in Holtzman et al. (https://huggingface.co/papers/1904.09751)\n        min_tokens_to_keep (`int`, *optional*, defaults to 1):\n            Minimumber of tokens we keep per batch example in the output.\n\n    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317\n    \"\"\"\n\n    if top_k > 0:\n        logits = TopKLogitsWarper(top_k=top_k, filter_value=filter_value, min_tokens_to_keep=min_tokens_to_keep)(\n            None, logits\n        )\n\n    if 0 <= top_p <= 1.0:\n        logits = TopPLogitsWarper(top_p=top_p, filter_value=filter_value, min_tokens_to_keep=min_tokens_to_keep)(\n            None, logits\n        )\n\n    return logits\n\n\ndef flatten_dict(nested: Dict, sep: str = \"/\") -> Dict:\n    \"\"\"Flatten dictionary and concatenate nested keys with separator.\"\"\"\n\n    def recurse(nest: Dict, prefix: str, into: Dict) -> None:\n        for k, v in nest.items():\n            if sep in k:\n                raise ValueError(f\"separator '{sep}' not allowed to be in key '{k}'\")\n            if isinstance(v, Mapping):\n                recurse(v, prefix + k + sep, into)\n            else:\n                into[prefix + k] = v\n\n    flat = {}\n    recurse(nested, \"\", flat)\n    return flat\n\n\ndef convert_to_scalar(stats: Dict) -> Dict:\n    \"\"\"\n    Converts the stats from a flattened dict to single scalar dicts\n    \"\"\"\n    tensorboard_stats = {}\n    for k, v in stats.items():\n        # for tensorboard compatibility - arrays and tensors are ignored with tensorboard\n        # therefore we convert single element tensors to scalars\n        if (isinstance(v, torch.Tensor) or isinstance(v, np.ndarray)) and (\n            len(v.shape) == 0 or (len(v.shape) == 1 and v.shape[0] == 1)\n        ):\n            v = v.item()\n        tensorboard_stats[k] = v\n    return tensorboard_stats\n\n\ndef stack_dicts(stats_dicts: List[Dict]) -> Dict:\n    \"\"\"Stack the values of a dict.\"\"\"\n    results = dict()\n    for k in stats_dicts[0]:\n        stats_list = [torch.flatten(d[k]) for d in stats_dicts]\n        results[k] = pad_sequence(stats_list, batch_first=True, padding_value=WANDB_PADDING)\n    return results\n\n\ndef logprobs_from_logits(logits: torch.Tensor, labels: torch.Tensor, gather: bool = True) -> torch.Tensor:\n    \"\"\"\n    See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591\n    \"\"\"\n    logp = F.log_softmax(logits, dim=2)\n\n    if not gather:\n        return logp\n    logpy = torch.gather(logp, 2, labels.unsqueeze(2)).squeeze(-1)\n    return logpy\n\n\ndef whiten(values: torch.Tensor, shift_mean: bool = True) -> torch.Tensor:\n    \"\"\"Whiten values.\"\"\"\n    mean, var = torch.mean(values), torch.var(values)\n    whitened = (values - mean) * torch.rsqrt(var + 1e-8)\n    if not shift_mean:\n        whitened += mean\n    return whitened\n\n\ndef masked_mean(values: torch.Tensor, mask: torch.Tensor, axis: Optional[bool] = None) -> torch.Tensor:\n    \"\"\"Compute mean of tensor with a masked values.\"\"\"\n    if axis is not None:\n        return (values * mask).sum(axis=axis) / mask.sum(axis=axis)\n    else:\n        return (values * mask).sum() / mask.sum()\n\n\ndef masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) -> torch.Tensor:\n    \"\"\"Compute variance of tensor with masked values.\"\"\"\n    mean = masked_mean(values, mask)\n    centered_values = values - mean\n    variance = masked_mean(centered_values**2, mask)\n    if unbiased:\n        mask_sum = mask.sum()\n        if mask_sum == 0:\n            raise ValueError(\n                \"The sum of the mask is zero, which can happen when `mini_batch_size=1`;\"\n                \"try increase the `mini_batch_size` or `gradient_accumulation_steps`\"\n            )\n        # note that if mask_sum == 1, then there is a division by zero issue\n        # to avoid it you just need to use a larger minibatch_size\n        bessel_correction = mask_sum / (mask_sum - 1)\n        variance = variance * bessel_correction\n    return variance\n\n\ndef masked_whiten(values: torch.Tensor, mask: torch.Tensor, shift_mean: bool = True) -> torch.Tensor:\n    \"\"\"Whiten values with masked values.\"\"\"\n    mean, var = masked_mean(values, mask), masked_var(values, mask)\n    whitened = (values - mean) * torch.rsqrt(var + 1e-8)\n    if not shift_mean:\n        whitened += mean\n    return whitened\n\n\ndef clip_by_value(x: torch.Tensor, tensor_min: float, tensor_max: float) -> torch.Tensor:\n    \"\"\"\n    Tensor extension to torch.clamp\n    https://github.com/pytorch/pytorch/issues/2793#issuecomment-428784713\n    \"\"\"\n    clipped = torch.max(torch.min(x, tensor_max), tensor_min)\n    return clipped\n\n\ndef entropy_from_logits(logits: torch.Tensor) -> torch.Tensor:\n    \"\"\"Calculate entropy from logits.\"\"\"\n    pd = torch.nn.functional.softmax(logits, dim=-1)\n    entropy = torch.logsumexp(logits, axis=-1) - torch.sum(pd * logits, axis=-1)\n    return entropy\n\n\ndef stats_to_np(stats_dict: Dict) -> Dict:\n    \"\"\"Cast all torch.tensors in dict to numpy arrays.\"\"\"\n    new_dict = dict()\n    for k, v in stats_dict.items():\n        if isinstance(v, torch.Tensor):\n            new_dict[k] = v.detach().cpu()\n            if new_dict[k].dtype == torch.bfloat16:\n                new_dict[k] = new_dict[k].float()\n            new_dict[k] = new_dict[k].numpy()\n        else:\n            new_dict[k] = v\n        if np.isscalar(new_dict[k]):\n            new_dict[k] = float(new_dict[k])\n    return new_dict\n\n\ndef respond_to_batch(\n    model: nn.Module, queries: List[torch.LongTensor], txt_len: int = 20, top_k: int = 0, top_p: float = 1.0\n) -> torch.LongTensor:\n    \"\"\"Sample text from language model.\"\"\"\n    input_ids = queries\n    for _i in range(txt_len):\n        # Get Logits\n        outputs = model(input_ids)\n        next_token_logits = outputs[0][:, -1, :]\n        next_token_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)\n        # Sample\n        probs = F.softmax(next_token_logits, dim=-1)\n        next_token = torch.multinomial(probs, num_samples=1).squeeze(1)\n        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)\n    return input_ids[:, -txt_len:]\n\n\ndef set_seed(seed: int) -> None:\n    \"\"\"\n    Helper function for reproducible behavior to set the seed in `random`, `numpy`, and `torch`.\n\n    Args:\n        seed (`int`): The seed to set.\n    \"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if is_torch_xpu_available():\n        torch.xpu.manual_seed_all(seed)\n    elif is_torch_npu_available():\n        torch.npu.manual_seed_all(seed)\n    else:\n        torch.cuda.manual_seed_all(seed)\n\n\nclass LengthSampler:\n    \"\"\"\n    Samples a length\n    \"\"\"\n\n    def __init__(self, min_value: int, max_value: int):\n        self.values = list(range(min_value, max_value))\n\n    def __call__(self) -> int:\n        return np.random.choice(self.values)\n\n\nclass PPODecorators:\n    optimize_device_cache = False\n\n    @classmethod\n    @contextmanager\n    def empty_device_cache(cls):\n        yield\n        if cls.optimize_device_cache:\n            if is_torch_xpu_available():\n                gc.collect()\n                torch.xpu.empty_cache()\n                gc.collect()\n            elif is_torch_npu_available():\n                gc.collect()\n                torch.npu.empty_cache()\n                gc.collect()\n            elif torch.cuda.is_available():\n                gc.collect()\n                torch.cuda.empty_cache()\n                gc.collect()\n\n\ndef randn_tensor(\n    shape: Union[Tuple, List],\n    generator: Optional[Union[List[torch.Generator], torch.Generator]] = None,\n    device: Optional[torch.device] = None,\n    dtype: Optional[torch.dtype] = None,\n    layout: Optional[torch.layout] = None,\n) -> torch.Tensor:\n    \"\"\"A helper function to create random tensors on the desired `device` with the desired `dtype`. When\n    passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor\n    is always created on the CPU.\n    \"\"\"\n    # device on which tensor is created defaults to device\n    rand_device = device\n    batch_size = shape[0]\n\n    layout = layout or torch.strided\n    device = device or torch.device(\"cpu\")\n\n    if generator is not None:\n        gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type\n        if gen_device_type != device.type and gen_device_type == \"cpu\":\n            rand_device = \"cpu\"\n            if device != \"mps\":\n                warnings.warn(\n                    f\"The passed generator was created on 'cpu' even though a tensor on {device} was expected.\"\n                    f\" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably\"\n                    f\" slighly speed up this function by passing a generator that was created on the {device} device.\"\n                )\n        elif gen_device_type != device.type and gen_device_type == \"cuda\":\n            raise ValueError(f\"Cannot generate a {device} tensor from a generator of type {gen_device_type}.\")\n\n    # make sure generator list of length 1 is treated like a non-list\n    if isinstance(generator, list) and len(generator) == 1:\n        generator = generator[0]\n\n    if isinstance(generator, list):\n        shape = (1,) + shape[1:]\n        latents = [\n            torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout)\n            for i in range(batch_size)\n        ]\n        latents = torch.cat(latents, dim=0).to(device)\n    else:\n        latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)\n\n    return latents\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# flake8: noqa\n\n__version__ = \"0.12.0.dev0\"\n\nfrom typing import TYPE_CHECKING\nfrom .import_utils import _LazyModule, is_diffusers_available, OptionalDependencyNotAvailable\n\n\n_import_structure = {\n    \"core\": [\n        \"set_seed\",\n    ],\n    \"environment\": [\n        \"TextEnvironment\",\n        \"TextHistory\",\n    ],\n    \"extras\": [\n        \"BestOfNSampler\",\n    ],\n    \"import_utils\": [\n        \"is_diffusers_available\",\n        \"is_liger_kernel_available\",\n        \"is_llmblender_available\",\n    ],\n    \"models\": [\n        \"AutoModelForCausalLMWithValueHead\",\n        \"AutoModelForSeq2SeqLMWithValueHead\",\n        \"PreTrainedModelWrapper\",\n        \"create_reference_model\",\n        \"setup_chat_format\",\n        \"SUPPORTED_ARCHITECTURES\",\n    ],\n    \"trainer\": [\n        \"DataCollatorForCompletionOnlyLM\",\n        \"DPOConfig\",\n        \"DPOTrainer\",\n        \"CPOConfig\",\n        \"CPOTrainer\",\n        \"AlignPropConfig\",\n        \"AlignPropTrainer\",\n        \"IterativeSFTTrainer\",\n        \"KTOConfig\",\n        \"KTOTrainer\",\n        \"BCOConfig\",\n        \"BCOTrainer\",\n        \"ModelConfig\",\n        \"NashMDConfig\",\n        \"NashMDTrainer\",\n        \"OnlineDPOConfig\",\n        \"OnlineDPOTrainer\",\n        \"XPOConfig\",\n        \"XPOTrainer\",\n        \"ORPOConfig\",\n        \"ORPOTrainer\",\n        \"PPOConfig\",\n        \"PPOTrainer\",\n        \"PPOv2Config\",\n        \"PPOv2Trainer\",\n        \"RewardConfig\",\n        \"RewardTrainer\",\n        \"RLOOConfig\",\n        \"RLOOTrainer\",\n        \"SFTConfig\",\n        \"SFTTrainer\",\n        \"FDivergenceConstants\",\n        \"FDivergenceType\",\n        \"GKDTrainer\",\n        \"GKDConfig\",\n        \"WinRateCallback\",\n        \"BaseJudge\",\n        \"BaseRankJudge\",\n        \"BasePairwiseJudge\",\n        \"RandomRankJudge\",\n        \"RandomPairwiseJudge\",\n        \"PairRMJudge\",\n        \"HfPairwiseJudge\",\n        \"OpenAIPairwiseJudge\",\n        \"LogCompletionsCallback\",\n    ],\n    \"commands\": [],\n    \"commands.cli_utils\": [\"init_zero_verbose\", \"SFTScriptArguments\", \"DPOScriptArguments\", \"TrlParser\"],\n    \"trainer.callbacks\": [\"RichProgressCallback\", \"SyncRefModelCallback\"],\n    \"trainer.utils\": [\"get_kbit_device_map\", \"get_peft_config\", \"get_quantization_config\"],\n    \"multitask_prompt_tuning\": [\n        \"MultitaskPromptEmbedding\",\n        \"MultitaskPromptTuningConfig\",\n        \"MultitaskPromptTuningInit\",\n    ],\n    \"data_utils\": [\n        \"apply_chat_template\",\n        \"extract_prompt\",\n        \"is_conversational\",\n        \"maybe_apply_chat_template\",\n        \"maybe_extract_prompt\",\n        \"maybe_unpair_preference_dataset\",\n        \"unpair_preference_dataset\",\n    ],\n}\n\ntry:\n    if not is_diffusers_available():\n        raise OptionalDependencyNotAvailable()\nexcept OptionalDependencyNotAvailable:\n    pass\nelse:\n    _import_structure[\"models\"].extend(\n        [\n            \"DDPOPipelineOutput\",\n            \"DDPOSchedulerOutput\",\n            \"DDPOStableDiffusionPipeline\",\n            \"DefaultDDPOStableDiffusionPipeline\",\n        ]\n    )\n    _import_structure[\"trainer\"].extend([\"DDPOConfig\", \"DDPOTrainer\"])\n\nif TYPE_CHECKING:\n    from .core import set_seed\n    from .environment import TextEnvironment, TextHistory\n    from .extras import BestOfNSampler\n    from .import_utils import is_diffusers_available, is_liger_kernel_available, is_llmblender_available\n    from .models import (\n        AutoModelForCausalLMWithValueHead,\n        AutoModelForSeq2SeqLMWithValueHead,\n        PreTrainedModelWrapper,\n        create_reference_model,\n        setup_chat_format,\n        SUPPORTED_ARCHITECTURES,\n    )\n    from .trainer import (\n        DataCollatorForCompletionOnlyLM,\n        DPOConfig,\n        DPOTrainer,\n        CPOConfig,\n        CPOTrainer,\n        AlignPropConfig,\n        AlignPropTrainer,\n        IterativeSFTTrainer,\n        KTOConfig,\n        KTOTrainer,\n        BCOConfig,\n        BCOTrainer,\n        ModelConfig,\n        NashMDConfig,\n        NashMDTrainer,\n        OnlineDPOConfig,\n        OnlineDPOTrainer,\n        XPOConfig,\n        XPOTrainer,\n        ORPOConfig,\n        ORPOTrainer,\n        PPOConfig,\n        PPOTrainer,\n        PPOv2Config,\n        PPOv2Trainer,\n        RewardConfig,\n        RewardTrainer,\n        RLOOConfig,\n        RLOOTrainer,\n        SFTConfig,\n        SFTTrainer,\n        FDivergenceConstants,\n        FDivergenceType,\n        GKDTrainer,\n        GKDConfig,\n        WinRateCallback,\n        BaseJudge,\n        BaseRankJudge,\n        BasePairwiseJudge,\n        RandomRankJudge,\n        RandomPairwiseJudge,\n        PairRMJudge,\n        HfPairwiseJudge,\n        OpenAIPairwiseJudge,\n        LogCompletionsCallback,\n    )\n    from .trainer.callbacks import RichProgressCallback, SyncRefModelCallback\n    from .trainer.utils import get_kbit_device_map, get_peft_config, get_quantization_config\n    from .commands.cli_utils import init_zero_verbose, SFTScriptArguments, DPOScriptArguments, TrlParser\n    from .data_utils import (\n        apply_chat_template,\n        extract_prompt,\n        is_conversational,\n        maybe_apply_chat_template,\n        maybe_extract_prompt,\n        maybe_unpair_preference_dataset,\n        unpair_preference_dataset,\n    )\n\n    try:\n        if not is_diffusers_available():\n            raise OptionalDependencyNotAvailable()\n    except OptionalDependencyNotAvailable:\n        pass\n    else:\n        from .models import (\n            DDPOPipelineOutput,\n            DDPOSchedulerOutput,\n            DDPOStableDiffusionPipeline,\n            DefaultDDPOStableDiffusionPipeline,\n        )\n        from .trainer import DDPOConfig, DDPOTrainer\n\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(\n        __name__,\n        globals()[\"__file__\"],\n        _import_structure,\n        module_spec=__spec__,\n        extra_objects={\"__version__\": __version__},\n    )\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\nimport warnings\nfrom typing import Optional\n\nimport torch\nfrom accelerate.utils import extract_model_from_parallel\nfrom transformers import StoppingCriteria, StoppingCriteriaList\n\nfrom ..import_utils import is_rich_available\n\n\nif is_rich_available():\n    from rich import print\n    from rich.text import Text\n\n\nclass StringStoppingCriteria(StoppingCriteria):\n    \"\"\"Custom `StoppingCriteria` which checks if all generations in the batch are completed.\"\"\"\n\n    def __init__(self, stop_strings, tokenizer):\n        self.stop_strings = stop_strings\n        self.tokenizer = tokenizer\n        self.first_call = True\n\n    def __call__(self, input_ids, scores, **kwargs):\n        \"\"\"Returns true if all generated sequences contain any of the stop strings.\"\"\"\n        if self.first_call:\n            self.generated_tokens = [1 for _ in range(input_ids.shape[0])]\n            self.start_length = input_ids.shape[-1] - 1\n            self.first_call = False\n        decoded_generations = self.tokenizer.batch_decode(input_ids[:, self.start_length :])\n        done = []\n\n        for i, decoded_generation in enumerate(decoded_generations):\n            sequence_complete = any(stop_string in decoded_generation for stop_string in self.stop_strings)\n            done.append(sequence_complete)\n            if not sequence_complete:\n                self.generated_tokens[i] += 1\n\n        if all(done):\n            self.first_call = True\n\n        return all(done)\n\n\nclass TextHistory:\n    \"\"\"The TextHistory class keeps track of the history of an interaction between the language model and the environment.\"\"\"\n\n    def __init__(self, text, tokens, system=True):\n        \"\"\"\n        Initialize TextHistory.\n\n        Args:\n            text (`str`): The text of the first segment.\n            tokens (`torch.LongTensor`): The tokens of the first segment.\n            system (`bool`, *optional*): Whether the first segment is a system or user segment.\n        \"\"\"\n        self.system_spans = []\n        self.text_spans = []\n        self.token_spans = []\n        self.token_masks = torch.tensor([], dtype=torch.long).to(tokens.device)\n        self.text = \"\"\n        self.tokens = torch.tensor([], dtype=torch.long).to(tokens.device)\n        self.completed = False\n        self.truncated = False\n        self.reward = 0.0\n\n        self.prompt_color = \"black on grey85\"\n        self.system_color = \"black on cyan3\"\n        self.model_color = \"black on deep_sky_blue1\"\n        self.reward_color = \"black on plum1\"\n\n        self.append_segment(text, tokens, system=system)\n\n    def append_segment(self, text, tokens, system=True):\n        \"\"\"\n        Append a new segment to the history.\n\n        Args:\n            text (`str`): The text of the new segment.\n            tokens (`torch.LongTensor`): The tokens of the new segment.\n            system (`bool`, *optional*): Whether the new segment is a system or user segment.\n        \"\"\"\n\n        if len(text) == 0 or len(tokens) == 0:\n            raise ValueError(\"Can't append empty text or token list to history.\")\n\n        original_text_length = len(self.text)\n\n        self.text += text\n        self.text_spans.append((original_text_length, len(self.text)))\n        self.system_spans.append(system)\n\n        original_token_length = len(self.tokens)\n\n        self.tokens = torch.cat((self.tokens, tokens))\n        if system:\n            self.token_masks = torch.cat((self.token_masks, torch.zeros_like(tokens)))\n        else:\n            self.token_masks = torch.cat((self.token_masks, torch.ones_like(tokens)))\n        self.token_spans.append((original_token_length, len(self.tokens)))\n\n    def complete(self, truncated=False):\n        \"\"\"\n        Mark the history as completed.\n        \"\"\"\n        self.completed = True\n        self.truncated = truncated\n\n    @property\n    def last_text_segment(self):\n        \"\"\"\n        Get the last text segment.\n        \"\"\"\n        start, end = self.text_spans[-1]\n        return self.text[start:end]\n\n    def split_query_response_tokens(self):\n        \"\"\"\n        Split the tokens into query and response tokens.\n        \"\"\"\n        split_index = self.token_spans[0][1]\n        query = self.tokens[:split_index]\n        response = self.tokens[split_index:]\n        mask = self.token_masks[split_index:]\n\n        return query, response, mask\n\n    def show_text(self, show_legend=False):\n        \"\"\"\n        Print the text history.\n        \"\"\"\n        if not is_rich_available():\n            warnings.warn(\"install rich to display text\")\n            return\n\n        text = Text(self.text)\n        text.stylize(self.prompt_color, self.text_spans[0][0], self.text_spans[1][0])\n        for i, (start, end) in enumerate(self.text_spans[1:]):\n            if self.system_spans[i + 1]:\n                text.stylize(self.system_color, start, end)\n            else:\n                text.stylize(self.model_color, start, end)\n\n        text.append(f\"\\n\\nReward: {self.reward}\", style=self.reward_color)\n        print(text)\n\n        if show_legend:\n            self.show_colour_legend()\n\n    def show_tokens(self, tokenizer, show_legend=False):\n        \"\"\"\n        Print the history tokens.\n        \"\"\"\n        if not is_rich_available():\n            warnings.warn(\"install rich to display tokens\")\n            return\n\n        text = Text()\n        prompt_end = self.token_spans[0][1]\n        for i, (token, mask) in enumerate(zip(self.tokens, self.token_masks)):\n            if i < prompt_end:\n                text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.prompt_color)\n                text.append(\" \")\n            elif mask == 0:\n                text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.system_color)\n                text.append(\" \")\n            else:\n                text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.model_color)\n                text.append(\" \")\n        text.append(f\"\\n\\nReward: {self.reward}\", style=self.reward_color)\n        print(text)\n        if show_legend:\n            self.show_colour_legend()\n\n    def show_colour_legend(self):\n        \"\"\"\n        Print the colour legend.\n        \"\"\"\n        if not is_rich_available():\n            warnings.warn(\"install rich to display colour legend\")\n            return\n        text = Text(\"\\n\\n(Colour Legend: \")\n        text.append(\"Prompt\", style=self.prompt_color)\n        text.append(\"|\")\n        text.append(\"System\", style=self.system_color)\n        text.append(\"|\")\n        text.append(\"Model\", style=self.model_color)\n        text.append(\"|\")\n        text.append(\"Reward\", style=self.reward_color)\n        text.append(\")\")\n        print(text)\n\n\nclass TextEnvironment:\n    \"\"\"\n    The TextEnvironment enables interaction of a LLM with an environment using tools.\n    \"\"\"\n\n    def __init__(\n        self,\n        model=None,\n        tokenizer=None,\n        tools=None,\n        reward_fn=None,\n        prompt=None,\n        max_turns=4,\n        max_tool_reponse=100,\n        max_length=None,\n        generation_kwargs=None,\n    ):\n        \"\"\"\n        Initialize TextEnvironment.\n\n        Args:\n            model (`PreTrainedModelWrapper`): The model to use for generation.\n            tokenizer (`transformers.PreTrainedTokenizer`): The tokenizer to use for generation.\n            tools (list): A list of tools to use for interaction.\n            reward_fn (function): A function that takes a string and returns a reward.\n            prompt (str): The base prompt to use for generation. Is prepended to the tasks.\n            max_turns (Optional[int]): The maximum number of turns to allow.\n            max_tool_response (Optional[int]): The maximum number of characters to allow in a tool response.\n            max_length (Optional[int]): The maximum number of tokens to allow in an episode.\n            generation_kwargs (Optional[dict]): A dictionary of keyword arguments to pass to the model's generate method.\n        \"\"\"\n        self.model = model\n        self.tokenizer = tokenizer\n        self.prompt = prompt\n        if isinstance(tools, dict):\n            self.tools = tools\n        else:\n            self.tools = {tool.__class__.__name__: tool for tool in tools}\n        self.reward_fn = reward_fn\n        self.max_length = max_length\n        self.request_token = \"<request>\"\n        self.call_token = \"<call>\"\n        self.response_token = \"<response>\"\n        self.submit_token = \"<submit>\"\n        self.max_turns = max_turns\n        self.max_tool_response = max_tool_reponse\n\n        if generation_kwargs is None:\n            self.generation_kwargs = dict()\n        else:\n            self.generation_kwargs = generation_kwargs\n\n        self.is_encoder_decoder = hasattr(self.model, \"is_encoder_decoder\")\n        self.current_device = extract_model_from_parallel(self.model).pretrained_model.device\n\n    def run(self, queries, **rewards_kwargs):\n        \"\"\"\n        Run the environment on a list of queries.\n\n        Args:\n            queries (list[str]): A list of queries to run the model in the environment on.\n        \"\"\"\n        turns = 0\n\n        queries = [self.prompt + task for task in queries]\n        queries_tokens = [\n            self.tokenizer(query, return_tensors=\"pt\").input_ids[0].to(self.model.pretrained_model.device)\n            for query in queries\n        ]\n\n        histories = [TextHistory(q, qt, system=True) for q, qt in zip(queries, queries_tokens)]\n\n        while any(not history.completed for history in histories) and turns < self.max_turns:\n            histories = self.generate(histories)\n            histories = self.tasks_end_check(histories)\n            # TODO: make this parallel rather than for-loop\n            for i in range(len(histories)):\n                histories[i] = self.step(histories[i])\n            histories = self.tasks_end_check(histories, model_turn=False)\n            turns += 1\n        self.compute_reward(histories, **rewards_kwargs)\n\n        # convert a list of (q, r, m) tuples to lists of all qs, rs, and ms respectively\n        queries, responses, masks = map(list, zip(*[history.split_query_response_tokens() for history in histories]))\n\n        rewards = [history.reward for history in histories]\n        return queries, responses, masks, rewards, histories\n\n    def step(self, history):\n        \"\"\"\n        Step the environment forward one turn.\n\n        Args:\n            history (`TextHistory`): The history to step forward.\n        \"\"\"\n        truncated, ended = self.task_end_check(history)\n        if ended:\n            history.complete(truncated=truncated)\n        if history.completed:\n            return history\n\n        tool, query = self.parse_tool_call(history.last_text_segment)\n        if tool is None or query is None:\n            response = f\"Unknown tool call: {history.last_text_segment}\"\n        else:\n            if tool not in self.tools:\n                response = f\"Unknown tool {tool}.\"\n            try:\n                response = self.tools[tool](query)\n            except Exception as error:\n                response = f\"Tool error: {str(error)}\"\n\n        if len(response) > self.max_tool_response:\n            response = response[: (self.max_tool_response - 3)] + \"...\"\n\n        history.append_segment(\n            response + self.response_token,\n            self.tokenizer(response + self.response_token, return_tensors=\"pt\")\n            .input_ids[0]\n            .to(self.model.pretrained_model.device),\n            system=True,\n        )\n\n        return history\n\n    def parse_tool_call(self, text):\n        \"\"\"\n        Parse request string. Expected format: <request><tool_name>query<call>\n        \"\"\"\n        result = re.search(f\"(?<={self.request_token}).*?(?={self.call_token})\", text, re.DOTALL)\n\n        # if we can't find a <request>/<call> span we return none\n        if result is None:\n            return None, None\n        else:\n            extracted_text = result.group()\n\n        result = re.search(r\"<(.*?)>\", extracted_text)\n\n        # if we can't find a tool name we return none\n        if result is None:\n            return None, None\n        else:\n            tool = result.group(1)\n\n        # split off the tool name\n        query = \">\".join(extracted_text.split(\">\")[1:])\n\n        return tool, query\n\n    def compute_reward(self, histories, **reward_kwargs):\n        \"\"\"\n        Compute the reward for a list of histories.\n        \"\"\"\n        rewards = self.reward_fn([history.last_text_segment for history in histories], **reward_kwargs)\n        for history, reward in zip(histories, rewards):\n            history.reward = reward\n        return histories\n\n    def generate(self, histories):\n        \"\"\"\n        Generate responses for a list of histories.\n        \"\"\"\n        active_histories = [i for i, history in enumerate(histories) if not history.completed]\n\n        query_tensors = [histories[i].tokens for i in active_histories]\n        response_tensors = self._generate_batched(query_tensors)\n        response_texts = self.tokenizer.batch_decode(response_tensors)\n\n        for i, response_text, response_tensor in zip(active_histories, response_texts, response_tensors):\n            histories[i].append_segment(response_text, response_tensor, system=False)\n\n        return histories\n\n    def tasks_end_check(self, histories, model_turn=True):\n        \"\"\"\n        Check if the current generation sequences have finished.\n        \"\"\"\n        for history in histories:\n            if not history.completed:\n                truncated, ended = self.task_end_check(history, model_turn=model_turn)\n                if ended:\n                    history.complete(truncated=truncated)\n        return histories\n\n    def task_end_check(self, history, model_turn=True):\n        \"\"\"\n        Check if the current generation sequence has finished.\n        \"\"\"\n        truncated = False\n        ended = False\n        if history.completed:\n            return truncated, ended\n        if self.max_length is not None and len(self.tokenizer(history.text).input_ids[0]) > self.max_length:\n            truncated = True\n            ended = True\n        elif self.tokenizer.eos_token in history.text:\n            ended = True\n        elif model_turn and not (\n            (self.request_token in history.last_text_segment and self.call_token in history.last_text_segment)\n            or self.submit_token in history.last_text_segment\n        ):\n            ended = True\n        elif self.submit_token in history.last_text_segment:\n            ended = True\n        return truncated, ended\n\n    def _generate_batched(\n        self,\n        query_tensors,\n        batch_size: int = 16,\n        pad_to_multiple_of: Optional[int] = None,\n    ):\n        \"\"\"\n        Generate responses for a list of query tensors.\n\n        Args:\n            query_tensors (list[torch.Tensor]): A list of query tensors to generate responses for.\n            batch_size (int): The batch size to use for generation.\n            pad_to_multiple_of (int): The padding length to use for generation.\n        \"\"\"\n        outputs = []\n        padding_side_default = self.tokenizer.padding_side\n        if not self.is_encoder_decoder:\n            self.tokenizer.padding_side = \"left\"\n\n        # in case we have fewer examples than bs\n        batch_size = min(len(query_tensors), batch_size)\n\n        for i in range(0, len(query_tensors), batch_size):\n            # prevent overflow if query tensors are not even multiple of bs\n            end_index = min(len(query_tensors), i + batch_size)\n\n            batch = query_tensors[i:end_index]\n            batch_mask = [torch.ones_like(element) for element in batch]\n            inputs = {\"input_ids\": batch, \"attention_mask\": batch_mask}\n\n            padded_inputs = self.tokenizer.pad(\n                inputs,\n                padding=True,\n                max_length=None,\n                pad_to_multiple_of=pad_to_multiple_of,\n                return_tensors=\"pt\",\n            ).to(self.current_device)\n\n            stopping_criteria = StringStoppingCriteria([self.call_token, self.submit_token], self.tokenizer)\n\n            self.generation_kwargs[\"stopping_criteria\"] = StoppingCriteriaList([stopping_criteria])\n\n            generations = extract_model_from_parallel(self.model).generate(**padded_inputs, **self.generation_kwargs)\n\n            for generation, mask, generated_tokens in zip(\n                generations, padded_inputs[\"attention_mask\"], stopping_criteria.generated_tokens\n            ):\n                if not self.is_encoder_decoder:\n                    output = generation[(1 - mask).sum() :]  # remove padding\n                else:\n                    output = generation\n\n                if not self.is_encoder_decoder:\n                    output = output[(mask).sum() :]  # remove prompt\n\n                # remove chunk generated after stopping criteria in batch mode\n                outputs.append(output[:generated_tokens])\n        self.tokenizer.padding_side = padding_side_default\n        return outputs\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# flake8: noqa\nfrom typing import TYPE_CHECKING\nfrom ..import_utils import _LazyModule\n\n\n_import_structure = {\n    \"base_environment\": [\"TextEnvironment\", \"TextHistory\"],\n}\n\nif TYPE_CHECKING:\n    from .base_environment import TextEnvironment, TextHistory\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(__name__, globals()[\"__file__\"], _import_structure, module_spec=__spec__)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nfrom typing import Callable, Literal, Optional, Union\n\nfrom datasets import Dataset, Value\nfrom transformers import AutoTokenizer\n\nfrom ..trainer.utils import ConstantLengthDataset\n\n\nFORMAT_MAPPING = {\n    \"chatml\": [{\"content\": Value(dtype=\"string\", id=None), \"role\": Value(dtype=\"string\", id=None)}],\n    \"instruction\": {\"completion\": Value(dtype=\"string\", id=None), \"prompt\": Value(dtype=\"string\", id=None)},\n}\n\n\ndef conversations_formatting_function(tokenizer: AutoTokenizer, messages_field: Literal[\"messages\", \"conversations\"]):\n    r\"\"\"\n    return a callable function that takes in a \"messages\" dataset and returns a formatted dataset, based on the tokenizer\n    apply chat template to the dataset\n    \"\"\"\n\n    def format_dataset(examples):\n        if isinstance(examples[messages_field][0], list):\n            output_texts = []\n            for i in range(len(examples[messages_field])):\n                output_texts.append(tokenizer.apply_chat_template(examples[messages_field][i], tokenize=False))\n            return output_texts\n        else:\n            return tokenizer.apply_chat_template(examples[messages_field], tokenize=False)\n\n    return format_dataset\n\n\ndef instructions_formatting_function(tokenizer: AutoTokenizer):\n    r\"\"\"\n    return a callable function that takes in an \"instructions\" dataset and returns a formatted dataset, based on the tokenizer\n    apply chat template to the dataset\n    \"\"\"\n\n    def format_dataset(examples):\n        if isinstance(examples[\"prompt\"], list):\n            output_texts = []\n            for i in range(len(examples[\"prompt\"])):\n                converted_sample = [\n                    {\"role\": \"user\", \"content\": examples[\"prompt\"][i]},\n                    {\"role\": \"assistant\", \"content\": examples[\"completion\"][i]},\n                ]\n                output_texts.append(tokenizer.apply_chat_template(converted_sample, tokenize=False))\n            return output_texts\n        else:\n            converted_sample = [\n                {\"role\": \"user\", \"content\": examples[\"prompt\"]},\n                {\"role\": \"assistant\", \"content\": examples[\"completion\"]},\n            ]\n            return tokenizer.apply_chat_template(converted_sample, tokenize=False)\n\n    return format_dataset\n\n\ndef get_formatting_func_from_dataset(\n    dataset: Union[Dataset, ConstantLengthDataset], tokenizer: AutoTokenizer\n) -> Optional[Callable]:\n    r\"\"\"\n    Finds the correct formatting function based on the dataset structure. Currently supported datasets are:\n    - `ChatML` with [{\"role\": str, \"content\": str}]\n    - `instruction` with [{\"prompt\": str, \"completion\": str}]\n\n    Args:\n        dataset (Dataset): User dataset\n        tokenizer (AutoTokenizer): Tokenizer used for formatting\n\n    Returns:\n        Callable: Formatting function if the dataset format is supported else None\n    \"\"\"\n    if isinstance(dataset, Dataset):\n        if \"messages\" in dataset.features:\n            if dataset.features[\"messages\"] == FORMAT_MAPPING[\"chatml\"]:\n                logging.info(\"Formatting dataset with chatml format\")\n                return conversations_formatting_function(tokenizer, \"messages\")\n        if \"conversations\" in dataset.features:\n            if dataset.features[\"conversations\"] == FORMAT_MAPPING[\"chatml\"]:\n                logging.info(\"Formatting dataset with chatml format\")\n                return conversations_formatting_function(tokenizer, \"conversations\")\n        elif dataset.features == FORMAT_MAPPING[\"instruction\"]:\n            logging.info(\"Formatting dataset with instruction format\")\n            return instructions_formatting_function(tokenizer)\n\n    return None\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any, Callable, List, Optional, Union\n\nimport torch\nfrom transformers import GenerationConfig, PreTrainedTokenizer, PreTrainedTokenizerFast\n\nfrom ..core import set_seed\nfrom ..models import SUPPORTED_ARCHITECTURES, PreTrainedModelWrapper\n\n\nclass BestOfNSampler:\n    def __init__(\n        self,\n        model: PreTrainedModelWrapper,\n        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],\n        queries_to_scores: Callable[[List[str]], List[float]],\n        length_sampler: Any,\n        sample_size: int = 4,\n        seed: Optional[int] = None,\n        n_candidates: int = 1,\n        generation_config: Optional[GenerationConfig] = None,\n    ) -> None:\n        r\"\"\"\n        Initialize the sampler for best-of-n generation\n\n        Args:\n            model (`PreTrainedModelWrapper`):\n                The pretrained model to use for generation\n            tokenizer (`PreTrainedTokenizer` or `PreTrainedTokenizerFast`):\n                Tokenizer associated with the pretrained model\n            queries_to_scores (`Callable[[List[str]], List[float]]`):\n                Callable that takes a list of generated texts and returns the associated reward scores\n            length_sampler (`Any`):\n                Sampler used to sample the length of the generated text\n            sample_size (`int`):\n                Number of samples to generate for each query\n            seed (`int`, *optional*):\n                Random seed used to control generation\n            n_candidates (`int`):\n                Number of candidates to return for each query\n            generation_config (`GenerationConfig`, *optional*):\n                Generation config passed to the underlying model's `generate` method.\n                See `GenerationConfig` (https://huggingface.co/docs/transformers/v4.29.1/en/main_classes/text_generation#transformers.GenerationConfig) for more details\n        \"\"\"\n        if seed is not None:\n            set_seed(seed)\n\n        if not isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):\n            raise ValueError(\n                f\"tokenizer must be a PreTrainedTokenizer or PreTrainedTokenizerFast, got {type(tokenizer)}\"\n            )\n        if not isinstance(model, (SUPPORTED_ARCHITECTURES)):\n            raise ValueError(\n                f\"model must be a PreTrainedModelWrapper, got {type(model)} - supported architectures are: {SUPPORTED_ARCHITECTURES}\"\n            )\n\n        self.model = model\n        self.tokenizer = tokenizer\n\n        self.queries_to_scores = queries_to_scores\n        self.length_sampler = length_sampler\n        self.gen_config = generation_config\n        self.sample_size = sample_size\n        self.n_candidates = n_candidates\n\n    def generate(\n        self,\n        tokenized_query: Union[List[int], torch.Tensor, List[torch.Tensor], List[List[int]]],\n        skip_special_tokens: bool = True,\n        device: Optional[Union[str, torch.device]] = None,\n        **generation_kwargs,\n    ) -> List[List[str]]:\n        r\"\"\"\n        Generate the best of n samples for input queries\n\n        Args:\n            tokenized_query (`List[int]` or `torch.Tensor` or `List[torch.Tensor]` or `List[int]`):\n                represents either a single tokenized query (a single tensor or a list of integers) or a batch of tokenized queries (a list of tensors or a list of lists of integers)\n            skip_special_tokens (`bool`):\n                Whether to remove the special tokens from the output\n            device (`str` or `torch.device`, *optional*):\n                The device on which the model will be loaded\n            **generation_kwargs (`dict`, *optional*):\n                Additional keyword arguments passed along to the underlying model's `generate` method.\n                This is used to override generation config\n\n        Returns:\n            List[List[str]]: A list of lists of generated texts\n        \"\"\"\n        queries = None\n\n        if isinstance(tokenized_query, torch.Tensor) and tokenized_query.ndim == 1:\n            queries = tokenized_query.unsqueeze(0)\n        elif isinstance(tokenized_query, List):\n            element_type = type(tokenized_query[0])\n            if element_type is int:\n                queries = torch.tensor(tokenized_query).unsqueeze(0)\n            elif element_type is torch.Tensor:\n                queries = [tensor.reshape((1, -1)) for tensor in tokenized_query]\n            else:\n                queries = [torch.tensor(query).reshape((1, -1)) for query in tokenized_query]\n\n        result = []\n\n        for query in queries:\n            queries = query.repeat((self.sample_size, 1))\n            output = self.model.generate(\n                queries.to(device),\n                max_new_tokens=self.length_sampler(),\n                generation_config=self.gen_config,\n                **generation_kwargs,\n            ).squeeze()\n            output = self.tokenizer.batch_decode(output, skip_special_tokens=skip_special_tokens)\n            scores = torch.tensor(self.queries_to_scores(output))\n            output = [output[i] for i in scores.topk(self.n_candidates).indices]\n            result.append(output)\n\n        return result\n\n\n# flake8: noqa\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import TYPE_CHECKING\n\nfrom ..import_utils import _LazyModule\n\n\n_import_structure = {\n    \"best_of_n_sampler\": [\"BestOfNSampler\"],\n}\n\nif TYPE_CHECKING:\n    from .best_of_n_sampler import BestOfNSampler\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(__name__, globals()[\"__file__\"], _import_structure, module_spec=__spec__)\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nState dict utilities: utility methods for converting state dicts easily\nFile copied from diffusers to avoid import issues and make TRL compatible\nwith most of diffusers versions.\n\"\"\"\n\nimport enum\n\n\nclass StateDictType(enum.Enum):\n    \"\"\"\n    The mode to use when converting state dicts.\n    \"\"\"\n\n    DIFFUSERS_OLD = \"diffusers_old\"\n    PEFT = \"peft\"\n\n\nPEFT_TO_DIFFUSERS = {\n    \".q_proj.lora_B\": \".q_proj.lora_linear_layer.up\",\n    \".q_proj.lora_A\": \".q_proj.lora_linear_layer.down\",\n    \".k_proj.lora_B\": \".k_proj.lora_linear_layer.up\",\n    \".k_proj.lora_A\": \".k_proj.lora_linear_layer.down\",\n    \".v_proj.lora_B\": \".v_proj.lora_linear_layer.up\",\n    \".v_proj.lora_A\": \".v_proj.lora_linear_layer.down\",\n    \".out_proj.lora_B\": \".out_proj.lora_linear_layer.up\",\n    \".out_proj.lora_A\": \".out_proj.lora_linear_layer.down\",\n    \"to_k.lora_A\": \"to_k.lora.down\",\n    \"to_k.lora_B\": \"to_k.lora.up\",\n    \"to_q.lora_A\": \"to_q.lora.down\",\n    \"to_q.lora_B\": \"to_q.lora.up\",\n    \"to_v.lora_A\": \"to_v.lora.down\",\n    \"to_v.lora_B\": \"to_v.lora.up\",\n    \"to_out.0.lora_A\": \"to_out.0.lora.down\",\n    \"to_out.0.lora_B\": \"to_out.0.lora.up\",\n}\n\nDIFFUSERS_OLD_TO_DIFFUSERS = {\n    \".to_q_lora.up\": \".q_proj.lora_linear_layer.up\",\n    \".to_q_lora.down\": \".q_proj.lora_linear_layer.down\",\n    \".to_k_lora.up\": \".k_proj.lora_linear_layer.up\",\n    \".to_k_lora.down\": \".k_proj.lora_linear_layer.down\",\n    \".to_v_lora.up\": \".v_proj.lora_linear_layer.up\",\n    \".to_v_lora.down\": \".v_proj.lora_linear_layer.down\",\n    \".to_out_lora.up\": \".out_proj.lora_linear_layer.up\",\n    \".to_out_lora.down\": \".out_proj.lora_linear_layer.down\",\n}\n\nDIFFUSERS_STATE_DICT_MAPPINGS = {\n    StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_DIFFUSERS,\n    StateDictType.PEFT: PEFT_TO_DIFFUSERS,\n}\n\nKEYS_TO_ALWAYS_REPLACE = {\n    \".processor.\": \".\",\n}\n\n\ndef convert_state_dict(state_dict, mapping):\n    r\"\"\"\n    Simply iterates over the state dict and replaces the patterns in `mapping` with the corresponding values.\n\n    Args:\n        state_dict (`dict[str, torch.Tensor]`):\n            The state dict to convert.\n        mapping (`dict[str, str]`):\n            The mapping to use for conversion, the mapping should be a dictionary with the following structure:\n                - key: the pattern to replace\n                - value: the pattern to replace with\n\n    Returns:\n        converted_state_dict (`dict`)\n            The converted state dict.\n    \"\"\"\n    converted_state_dict = {}\n    for k, v in state_dict.items():\n        # First, filter out the keys that we always want to replace\n        for pattern in KEYS_TO_ALWAYS_REPLACE.keys():\n            if pattern in k:\n                new_pattern = KEYS_TO_ALWAYS_REPLACE[pattern]\n                k = k.replace(pattern, new_pattern)\n\n        for pattern in mapping.keys():\n            if pattern in k:\n                new_pattern = mapping[pattern]\n                k = k.replace(pattern, new_pattern)\n                break\n        converted_state_dict[k] = v\n    return converted_state_dict\n\n\ndef convert_state_dict_to_diffusers(state_dict, original_type=None, **kwargs):\n    r\"\"\"\n    Converts a state dict to new diffusers format. The state dict can be from previous diffusers format\n    (`OLD_DIFFUSERS`), or PEFT format (`PEFT`) or new diffusers format (`DIFFUSERS`). In the last case the method will\n    return the state dict as is.\n\n    The method only supports the conversion from diffusers old, PEFT to diffusers new for now.\n\n    Args:\n        state_dict (`dict[str, torch.Tensor]`):\n            The state dict to convert.\n        original_type (`StateDictType`, *optional*):\n            The original type of the state dict, if not provided, the method will try to infer it automatically.\n        kwargs (`dict`, *args*):\n            Additional arguments to pass to the method.\n\n            - **adapter_name**: For example, in case of PEFT, some keys will be pre-pended\n                with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in\n                `get_peft_model_state_dict` method:\n                https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92\n                but we add it here in case we don't want to rely on that method.\n    \"\"\"\n    peft_adapter_name = kwargs.pop(\"adapter_name\", None)\n    if peft_adapter_name is not None:\n        peft_adapter_name = \".\" + peft_adapter_name\n    else:\n        peft_adapter_name = \"\"\n\n    if original_type is None:\n        # Old diffusers to PEFT\n        if any(\"to_out_lora\" in k for k in state_dict.keys()):\n            original_type = StateDictType.DIFFUSERS_OLD\n        elif any(f\".lora_A{peft_adapter_name}.weight\" in k for k in state_dict.keys()):\n            original_type = StateDictType.PEFT\n        elif any(\"lora_linear_layer\" in k for k in state_dict.keys()):\n            # nothing to do\n            return state_dict\n        else:\n            raise ValueError(\"Could not automatically infer state dict type\")\n\n    if original_type not in DIFFUSERS_STATE_DICT_MAPPINGS.keys():\n        raise ValueError(f\"Original type {original_type} is not supported\")\n\n    mapping = DIFFUSERS_STATE_DICT_MAPPINGS[original_type]\n    return convert_state_dict(state_dict, mapping)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport itertools\nfrom contextlib import contextmanager\nfrom dataclasses import dataclass\nfrom typing import TYPE_CHECKING, Literal, Optional, Tuple, Union\n\nfrom accelerate.utils import is_deepspeed_available\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\nfrom .modeling_value_head import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead\n\n\nSUPPORTED_ARCHITECTURES = (\n    AutoModelForCausalLMWithValueHead,\n    AutoModelForSeq2SeqLMWithValueHead,\n)\n\nif is_deepspeed_available():\n    import deepspeed\n\nif TYPE_CHECKING:\n    from accelerate import Accelerator\n    from deepspeed.runtime.engine import DeepSpeedEngine\n    from torch.nn.parallel.distributed import DistributedDataParallel\n\n    from .modeling_base import PreTrainedModelWrapper\n\n\n# TODO: Add Abstract Base Class if more formats are added\n@dataclass\nclass ChatMlSpecialTokens:\n    \"\"\"Dataclass for special tokens used in ChatML, including system, user, assistant, bos, eos, and pad tokens.\"\"\"\n\n    bos_token: str = \"<|im_start|>\"\n    eos_token: str = \"<|im_end|>\"\n    pad_token: str = \"<|im_end|>\"\n\n    @property\n    def system(self):\n        return f\"{self.bos_token}system\"\n\n    @property\n    def user(self):\n        return f\"{self.bos_token}user\"\n\n    @property\n    def assistant(self):\n        return f\"{self.bos_token}assistant\"\n\n    @property\n    def chat_template(self):\n        return (\n            \"{% for message in messages %}\"\n            f\"{{{{'{self.bos_token}' + message['role'] + '\\n' + message['content'] + '{self.eos_token}' + '\\n'}}}}\"\n            \"{% endfor %}\"\n            \"{% if add_generation_prompt %}\"\n            f\"{{{{ '{self.assistant}\\n' }}}}\"\n            \"{% endif %}\"\n        )\n\n\nFORMAT_MAPPING = {\"chatml\": ChatMlSpecialTokens}\n\n\ndef setup_chat_format(\n    model: PreTrainedModel,\n    tokenizer: PreTrainedTokenizer,\n    format: Optional[Literal[\"chatml\"]] = \"chatml\",\n    resize_to_multiple_of: Optional[int] = None,\n) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:\n    \"\"\"\n    Setup chat format by adding special tokens to the tokenizer, setting the correct format, and extending the embedding layer of the model based on the new special tokens.\n\n    Args:\n        model (`~transformers.PreTrainedModel`): The model to be modified.\n        tokenizer (`~transformers.PreTrainedTokenizer`): The tokenizer to be modified.\n        format (`Optional[Literal[\"chatml\"]]`): The format to be set. Defaults to \"chatml\".\n        resize_to_multiple_of (`Optional[int]`): Number to resize the embedding layer to. Defaults to None.\n\n    Returns:\n        model (`~transformers.PreTrainedModel`): The modified model.\n        tokenizer (`~transformers.PreTrainedTokenizer`): The modified tokenizer.\n    \"\"\"\n    # check if format available and retrieve\n    if format not in FORMAT_MAPPING:\n        raise ValueError(f\"Format {format} not available. Please use one of {FORMAT_MAPPING.keys()}\")\n\n    chat_format = FORMAT_MAPPING[format]()\n\n    # set special tokens and them\n    tokenizer.eos_token = chat_format.eos_token\n    tokenizer.pad_token = chat_format.pad_token\n    tokenizer.bos_token = chat_format.bos_token\n    tokenizer.add_special_tokens({\"additional_special_tokens\": [chat_format.bos_token, chat_format.eos_token]})\n    # set chat format for tokenizer\n    tokenizer.chat_template = chat_format.chat_template\n\n    # resize embedding layer to a multiple of 64, https://x.com/karpathy/status/1621578354024677377\n    model.resize_token_embeddings(\n        len(tokenizer), pad_to_multiple_of=resize_to_multiple_of if resize_to_multiple_of is not None else None\n    )\n    # Update the model config to use the new eos & bos tokens\n    if getattr(model, \"config\", None) is not None:\n        model.config.pad_token_id = tokenizer.pad_token_id\n        model.config.bos_token_id = tokenizer.bos_token_id\n        model.config.eos_token_id = tokenizer.eos_token_id\n    # Update the generation config to use the new eos & bos token\n    if getattr(model, \"generation_config\", None) is not None:\n        model.generation_config.bos_token_id = tokenizer.bos_token_id\n        model.generation_config.eos_token_id = tokenizer.eos_token_id\n        model.generation_config.pad_token_id = tokenizer.pad_token_id\n\n    return model, tokenizer\n\n\ndef remove_hooks(model: \"DeepSpeedEngine\") -> None:\n    \"\"\"Removes the optimizer hooks from a DeepSpeed ZeRO-3 model.\"\"\"\n    if model.optimizer is not None and hasattr(model.optimizer, \"parameter_offload\"):\n        optimizer_offload = model.optimizer.parameter_offload\n    elif model.optimizer is not None:\n        optimizer_offload = model.optimizer\n\n    for param in iter_params(optimizer_offload.module, recurse=True):\n        param.ds_active_sub_modules.clear()\n\n    for hook in optimizer_offload.forward_hooks:\n        hook.remove()\n    for hook in optimizer_offload.backward_hooks:\n        hook.remove()\n\n    optimizer_offload.forward_hooks = []\n    optimizer_offload.backward_hooks = []\n\n\ndef get_all_parameters(sub_module, recurse=False):\n    return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters())\n\n\ndef iter_params(module, recurse=False):\n    return [param for _, param in get_all_parameters(module, recurse)]\n\n\ndef add_hooks(model: \"DeepSpeedEngine\") -> None:\n    \"\"\"Adds the optimizer hooks from a DeepSpeed ZeRO-3 model.\"\"\"\n    if model.optimizer is not None and hasattr(model.optimizer, \"parameter_offload\"):\n        optimizer_offload = model.optimizer.parameter_offload\n    elif model.optimizer is not None:\n        optimizer_offload = model.optimizer\n    optimizer_offload._register_hooks_recursively(optimizer_offload.module)\n\n\n@contextmanager\ndef unwrap_model_for_generation(\n    model: Union[\"DistributedDataParallel\", \"DeepSpeedEngine\"], accelerator: \"Accelerator\", is_peft_model: bool = False\n) -> Union[\"PreTrainedModelWrapper\", \"DeepSpeedEngine\"]:\n    \"\"\"Context manager to unwrap a model for generation.\n    For ZeRO-3 models, we gather the weights once to speed up generation.\n    \"\"\"\n    unwrapped_model = accelerator.unwrap_model(model)\n    if is_peft_model:\n        unwrapped_model.pretrained_model.disable_adapter()\n    if accelerator.state.deepspeed_plugin is not None and accelerator.state.deepspeed_plugin.zero_stage == 3:\n        with deepspeed.zero.GatheredParameters(model.parameters()):\n            remove_hooks(model)\n            yield accelerator.unwrap_model(model)\n            add_hooks(model)\n    else:\n        yield unwrapped_model\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom huggingface_hub import hf_hub_download\nfrom huggingface_hub.utils import EntryNotFoundError\nfrom transformers import CLIPModel, is_torch_npu_available, is_torch_xpu_available\n\n\nclass MLP(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.layers = nn.Sequential(\n            nn.Linear(768, 1024),\n            nn.Dropout(0.2),\n            nn.Linear(1024, 128),\n            nn.Dropout(0.2),\n            nn.Linear(128, 64),\n            nn.Dropout(0.1),\n            nn.Linear(64, 16),\n            nn.Linear(16, 1),\n        )\n\n    def forward(self, embed):\n        return self.layers(embed)\n\n\nclass AestheticScorer(torch.nn.Module):\n    \"\"\"\n    This model attempts to predict the aesthetic score of an image. The aesthetic score\n    is a numerical approximation of how much a specific image is liked by humans on average.\n    This is from https://github.com/christophschuhmann/improved-aesthetic-predictor\n    \"\"\"\n\n    def __init__(self, *, dtype, model_id, model_filename):\n        super().__init__()\n        self.clip = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n        self.normalize = torchvision.transforms.Normalize(\n            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]\n        )\n        self.target_size = 224\n        self.mlp = MLP()\n        try:\n            cached_path = hf_hub_download(model_id, model_filename)\n        except EntryNotFoundError:\n            cached_path = os.path.join(model_id, model_filename)\n        state_dict = torch.load(cached_path, map_location=torch.device(\"cpu\"), weights_only=True)\n        self.mlp.load_state_dict(state_dict)\n        self.dtype = dtype\n        self.eval()\n\n    def __call__(self, images):\n        device = next(self.parameters()).device\n        images = torchvision.transforms.Resize(self.target_size)(images)\n        images = self.normalize(images).to(self.dtype).to(device)\n        embed = self.clip.get_image_features(pixel_values=images)\n        # normalize embedding\n        embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)\n        reward = self.mlp(embed).squeeze(1)\n        return reward\n\n\ndef aesthetic_scorer(hub_model_id, model_filename):\n    scorer = AestheticScorer(\n        model_id=hub_model_id,\n        model_filename=model_filename,\n        dtype=torch.float32,\n    )\n    if is_torch_npu_available():\n        scorer = scorer.npu()\n    elif is_torch_xpu_available():\n        scorer = scorer.xpu()\n    else:\n        scorer = scorer.cuda()\n\n    def _fn(images, prompts, metadata):\n        images = (images).clamp(0, 1)\n        scores = scorer(images)\n        return scores, {}\n\n    return _fn\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport logging\nimport os\nfrom copy import deepcopy\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom accelerate import PartialState\nfrom huggingface_hub import hf_hub_download\nfrom huggingface_hub.utils import (\n    EntryNotFoundError,\n    HFValidationError,\n    LocalEntryNotFoundError,\n    RepositoryNotFoundError,\n)\nfrom safetensors.torch import load_file as safe_load_file\nfrom transformers import GenerationMixin, PreTrainedModel, is_torch_npu_available, is_torch_xpu_available\nfrom transformers.utils import is_peft_available\n\nfrom ..import_utils import is_transformers_greater_than\n\n\nif is_peft_available():\n    from peft import (\n        PeftConfig,\n        PeftModel,\n        PeftModelForCausalLM,\n        PeftModelForSeq2SeqLM,\n        PromptLearningConfig,\n        get_peft_model,\n        prepare_model_for_kbit_training,\n    )\n\nif is_transformers_greater_than(\"4.33.0\"):\n    from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled\nelse:\n    from transformers.deepspeed import is_deepspeed_zero3_enabled\n\nLAYER_PATTERNS = [\n    \"transformer.h.{layer}\",\n    \"model.decoder.layers.{layer}\",\n    \"gpt_neox.layers.{layer}\",\n    \"model.layers.{layer}\",\n]\n\n\nclass PreTrainedModelWrapper(nn.Module):\n    r\"\"\"\n    A wrapper class around a (`transformers.PreTrainedModel`) to be compatible with the\n    (`~transformers.PreTrained`) class in order to keep some attributes and methods of the\n    (`~transformers.PreTrainedModel`) class.\n\n    Attributes:\n        pretrained_model (`transformers.PreTrainedModel`):\n            The model to be wrapped.\n        parent_class (`transformers.PreTrainedModel`):\n            The parent class of the model to be wrapped.\n        supported_args (`list`):\n            The list of arguments that are supported by the wrapper class.\n    \"\"\"\n\n    transformers_parent_class = None\n    supported_args = None\n    supported_modules = (\"v_head\",)\n    supported_rm_modules = (\"score\",)\n    supported_pretrained_model_architectures = (\n        (PreTrainedModel)\n        if not is_peft_available()\n        else (PreTrainedModel, PeftModelForCausalLM, PeftModelForSeq2SeqLM)\n    )\n\n    def __init__(\n        self, pretrained_model=None, score_module=None, supports_rm_adapter=False, rm_adapter_name=None, **kwargs\n    ):\n        super().__init__()\n        self.pretrained_model = pretrained_model\n\n        self.config = pretrained_model.config\n        self.prepare_inputs_for_generation = pretrained_model.prepare_inputs_for_generation\n        self.is_loaded_in_8bit = getattr(pretrained_model, \"is_loaded_in_8bit\", False)\n        self.is_loaded_in_4bit = getattr(pretrained_model, \"is_loaded_in_4bit\", False)\n        self.is_sequential_parallel = False\n\n        if hasattr(pretrained_model, \"gradient_checkpointing_disable\"):\n            self.gradient_checkpointing_disable = pretrained_model.gradient_checkpointing_disable\n\n        if hasattr(pretrained_model, \"gradient_checkpointing_enable\"):\n            self.gradient_checkpointing_enable = pretrained_model.gradient_checkpointing_enable\n\n        if hasattr(pretrained_model, \"enable_input_require_grads\"):\n            self.enable_input_require_grads = pretrained_model.enable_input_require_grads\n\n        self.supports_rm_adapter = supports_rm_adapter\n        self.rm_adapter_name = rm_adapter_name\n        self.policy_adapter_name = \"default\"\n        if score_module is not None:\n            self.score = score_module\n\n    @classmethod\n    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):\n        r\"\"\"\n        Instantiates a new model from a pretrained model from `transformers`. The\n        pretrained model is loaded using the `from_pretrained` method of the\n        `transformers.PreTrainedModel` class. The arguments that are specific to the\n        `transformers.PreTrainedModel` class are passed along this method and filtered\n        out from the `kwargs` argument.\n\n        Args:\n            pretrained_model_name_or_path (`str` or `transformers.PreTrainedModel`):\n                The path to the pretrained model or its name.\n            *model_args (`list`, *optional*)):\n                Additional positional arguments passed along to the underlying model's\n                `from_pretrained` method.\n            **kwargs (`dict`, *optional*):\n                Additional keyword arguments passed along to the underlying model's\n                `from_pretrained` method. We also pre-process the kwargs to extract\n                the arguments that are specific to the `transformers.PreTrainedModel`\n                class and the arguments that are specific to trl models. The kwargs\n                also support `prepare_model_for_kbit_training` arguments from\n                `peft` library.\n        \"\"\"\n        if kwargs is not None:\n            peft_config = kwargs.pop(\"peft_config\", None)\n            reward_adapter = kwargs.pop(\"reward_adapter\", None)\n            reward_adapter_name = kwargs.pop(\"reward_adapter_name\", \"reward_adapter\")\n            is_trainable = kwargs.pop(\"is_trainable\", False)\n            trl_model_args, pretrained_kwargs, peft_quantization_kwargs = cls._split_kwargs(kwargs)\n            token = pretrained_kwargs.get(\"token\", None)\n        else:\n            peft_config = None\n            is_trainable = False\n            trl_model_args = {}\n            pretrained_kwargs = {}\n            peft_quantization_kwargs = {}\n            token = None\n\n        if reward_adapter is not None and not isinstance(reward_adapter, str):\n            raise ValueError(\n                \"The `reward_adapter` argument should be a string representing the name of local path or the Hub id to the Reward Modeling adapter.\"\n            )\n\n        is_peft_model = False\n\n        current_device = cls._get_current_device()\n        if isinstance(pretrained_model_name_or_path, str):\n            is_loaded_in_8bit = pretrained_kwargs[\"load_in_8bit\"] if \"load_in_8bit\" in pretrained_kwargs else False\n            is_loaded_in_4bit = pretrained_kwargs[\"load_in_4bit\"] if \"load_in_4bit\" in pretrained_kwargs else False\n        else:\n            is_loaded_in_8bit = getattr(pretrained_model_name_or_path, \"is_loaded_in_8bit\", False)\n            is_loaded_in_4bit = getattr(pretrained_model_name_or_path, \"is_loaded_in_4bit\", False)\n\n        if (is_loaded_in_8bit or is_loaded_in_4bit) and \"device_map\" not in pretrained_kwargs:\n            # warn users\n            logging.warning(\n                \"The `device_map` argument is not provided. We will override the device_map argument.\"\n                \" to set the entire\"\n                \" model on the current device. If you want to set the model on multiple devices, please provide\"\n                \" a custom `device_map` argument.\"\n            )\n            pretrained_kwargs[\"device_map\"] = {\"\": current_device}\n\n        if is_peft_available() and peft_config is not None and not isinstance(peft_config, PeftConfig):\n            raise ValueError(\"The `peft_config` argument should be an instance of `peft.PeftConfig` class.\")\n\n        # First, load the pre-trained model using the parent-class\n        # either `AutoModelForCausalLM` or `AutoModelForSeq2SeqLM`\n        if isinstance(pretrained_model_name_or_path, str):\n            if is_peft_available():\n                try:\n                    # If there is a trained peft adapter in the hub, load its config.\n                    remote_adapter_config = hf_hub_download(\n                        pretrained_model_name_or_path,\n                        \"adapter_config.json\",\n                        token=token,\n                    )\n                except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError):\n                    remote_adapter_config = None\n            else:\n                remote_adapter_config = None\n\n            local_adapter_present = os.path.exists(os.path.join(pretrained_model_name_or_path, \"adapter_config.json\"))\n\n            if (local_adapter_present or remote_adapter_config is not None) and is_peft_available():\n                if peft_config is not None:\n                    logging.warning(\n                        \"`peft_config` argument ignored since a peft config file was found in \"\n                        f\"{pretrained_model_name_or_path}\"\n                    )\n\n                # Load the trained peft adapter config\n                if local_adapter_present:\n                    trained_adapter_config = PeftConfig.from_pretrained(pretrained_model_name_or_path)\n                else:\n                    remote_adapter_dir = os.path.dirname(remote_adapter_config)\n                    trained_adapter_config = PeftConfig.from_pretrained(remote_adapter_dir)\n\n                # Load the pretrained base model\n                pretrained_model = cls.transformers_parent_class.from_pretrained(\n                    trained_adapter_config.base_model_name_or_path, *model_args, **pretrained_kwargs\n                )\n\n                # Wrap the pretrained model with the trained peft adapter\n                pretrained_model = PeftModel.from_pretrained(\n                    pretrained_model, pretrained_model_name_or_path, is_trainable=is_trainable, token=token\n                )\n                logging.info(\"Trained peft adapter loaded\")\n            else:\n                pretrained_model = cls.transformers_parent_class.from_pretrained(\n                    pretrained_model_name_or_path, *model_args, **pretrained_kwargs\n                )\n\n                if peft_config is not None:\n                    # Initialize a new peft adapter with the given config\n                    if is_loaded_in_8bit or is_loaded_in_4bit:\n                        pretrained_model = prepare_model_for_kbit_training(\n                            pretrained_model,\n                            **peft_quantization_kwargs,\n                        )\n                    pretrained_model = get_peft_model(pretrained_model, peft_config)\n                    logging.info(\"peft adapter initialised\")\n\n        elif isinstance(pretrained_model_name_or_path, cls.supported_pretrained_model_architectures):\n            pretrained_model = pretrained_model_name_or_path\n\n            if peft_config is not None and isinstance(pretrained_model, PreTrainedModel):\n                # Initialize a new peft adapter with the given config\n                if is_loaded_in_8bit or is_loaded_in_4bit:\n                    pretrained_model = prepare_model_for_kbit_training(\n                        pretrained_model,\n                        **peft_quantization_kwargs,\n                    )\n                pretrained_model = get_peft_model(pretrained_model, peft_config)\n                logging.info(\"peft adapter initialised\")\n        else:\n            raise ValueError(\n                \"pretrained_model_name_or_path should be a string or a PreTrainedModel, \"\n                f\"but is {type(pretrained_model_name_or_path)}\"\n            )\n\n        if is_peft_available():\n            if isinstance(pretrained_model, PeftModel):\n                is_peft_model = True\n                # for backward compatibility\n                if hasattr(pretrained_model, \"active_peft_config\") and isinstance(\n                    pretrained_model.active_peft_config, PromptLearningConfig\n                ):\n                    raise ValueError(\"PromptLearningConfig is not supported for PPO training.\")\n\n        # Add reward modeling adapter if specified\n        if not is_peft_model and reward_adapter is not None:\n            raise ValueError(\"reward_adapter can only be used with a PeftModel. \")\n        elif is_peft_model and reward_adapter is not None:\n            score_module = cls.add_and_load_reward_modeling_adapter(\n                pretrained_model, reward_adapter, reward_adapter_name, token=token\n            )\n            multi_adapter_args = {\n                \"score_module\": score_module,\n                \"supports_rm_adapter\": True,\n                \"rm_adapter_name\": reward_adapter_name,\n            }\n        else:\n            multi_adapter_args = {\"supports_rm_adapter\": False}\n\n        # Then, create the full model by instantiating the wrapper class\n        model = cls(pretrained_model, **multi_adapter_args, **trl_model_args)\n\n        # if resume_training, load the state_dict again - this is ok since the\n        # state_dict is removed from the model after loading it.\n        is_resuming_training = True\n        if isinstance(pretrained_model_name_or_path, str):\n            safe_filename = os.path.join(pretrained_model_name_or_path, \"model.safetensors\")\n            filename = os.path.join(pretrained_model_name_or_path, \"pytorch_model.bin\")\n\n            sharded_index_filename = os.path.join(pretrained_model_name_or_path, \"pytorch_model.bin.index.json\")\n            safe_sharded_index_filename = os.path.join(pretrained_model_name_or_path, \"model.safetensors.index.json\")\n            is_sharded = False\n            use_safe = os.path.exists(safe_filename)\n\n            if not (os.path.exists(filename) or os.path.exists(safe_filename)):\n                # Try with `pytorch_model.bin`\n                filename, files_to_download, is_sharded, is_resuming_training = cls._get_checkpoint_from_hub(\n                    pretrained_model,\n                    pretrained_model_name_or_path,\n                    sharded_index_filename,\n                    token=token,\n                )\n                # Try with safetensors\n                if filename is None and files_to_download is None:\n                    safe_filename, files_to_download, is_sharded, is_resuming_training = cls._get_checkpoint_from_hub(\n                        pretrained_model,\n                        pretrained_model_name_or_path,\n                        safe_sharded_index_filename,\n                        token=token,\n                        model_name=\"model.safetensors\",\n                        model_index_name=\"model.safetensors.index.json\",\n                    )\n                    use_safe = True\n                else:\n                    use_safe = False\n\n            loading_func = safe_load_file if use_safe else torch.load\n            load_kwargs = {} if use_safe else {\"map_location\": \"cpu\", \"weights_only\": True}\n\n            if is_resuming_training:\n                if is_sharded:\n                    # download each file and add it to the state_dict\n                    state_dict = {}\n\n                    for shard_file in files_to_download:\n                        filename = hf_hub_download(\n                            pretrained_model_name_or_path,\n                            shard_file,\n                            token=token,\n                        )\n                        state_dict.update(loading_func(filename, **load_kwargs))\n                else:\n                    state_dict = loading_func(filename if not use_safe else safe_filename, **load_kwargs)\n\n        else:\n            state_dict = pretrained_model_name_or_path.state_dict()\n\n        model.is_peft_model = is_peft_model\n        model.current_device = current_device\n\n        if is_resuming_training:\n            model.post_init(state_dict=state_dict)\n\n        return model\n\n    @classmethod\n    def _get_checkpoint_from_hub(\n        cls,\n        pretrained_model,\n        pretrained_model_name_or_path,\n        index_filename,\n        token=None,\n        model_name=\"pytorch_model.bin\",\n        model_index_name=\"pytorch_model.bin.index.json\",\n    ):\n        files_to_download = None\n        filename = None\n        is_resuming_training = True\n        is_sharded = False\n\n        try:\n            filename = hf_hub_download(\n                pretrained_model_name_or_path,\n                model_name,\n                token=token,\n            )\n        # sharded\n        except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError):\n            if os.path.exists(index_filename):\n                index_file_name = index_filename\n            else:\n                try:\n                    index_file_name = hf_hub_download(\n                        pretrained_model_name_or_path,\n                        model_index_name,\n                        token=token,\n                    )\n                except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError):\n                    # not continue training, do not have v_head weight\n                    is_resuming_training = False\n                    logging.warning(\n                        f\"A {type(pretrained_model)} model is loaded from '{pretrained_model_name_or_path}', \"\n                        f\"and no v_head weight is found. This IS expected if you are not resuming PPO training.\"\n                    )\n            # load json\n            if is_resuming_training:\n                with open(index_file_name) as f:\n                    index = json.load(f)\n                # check filename with `v_head` or any known extra module:\n                files_to_download = set()\n                for k, v in index[\"weight_map\"].items():\n                    if any(module in k for module in cls.supported_modules):\n                        files_to_download.add(v)\n                is_sharded = True\n\n        return filename, files_to_download, is_sharded, is_resuming_training\n\n    @classmethod\n    def _get_current_device(cls):\n        r\"\"\"\n        Get the current device. For GPU, we return the local process index using the `accelerate.PartialState`\n        object to handle corner cases when running scripts in distributed environments.\n\n        Returns:\n            current_device (`Union[int, str]`):\n                The current device.\n        \"\"\"\n        state = PartialState()\n        if is_torch_xpu_available():\n            return f\"xpu:{state.local_process_index}\"\n        elif is_torch_npu_available():\n            return f\"npu:{state.local_process_index}\"\n        else:\n            return state.local_process_index if torch.cuda.is_available() else \"cpu\"\n\n    @classmethod\n    def _split_kwargs(cls, kwargs):\n        \"\"\"\n        Separate the kwargs from the arguments that we support inside\n        `supported_args` and the ones that we don't.\n        \"\"\"\n        check_peft_kwargs = False\n\n        if is_peft_available():\n            from peft import prepare_model_for_kbit_training\n\n            check_peft_kwargs = True\n\n        supported_kwargs = {}\n        unsupported_kwargs = {}\n        peft_kwargs = {}\n\n        for key, value in kwargs.items():\n            if key in cls.supported_args:\n                supported_kwargs[key] = value\n            else:\n                unsupported_kwargs[key] = value\n\n            if check_peft_kwargs:\n                if key in prepare_model_for_kbit_training.__code__.co_varnames:\n                    peft_kwargs[key] = value\n                    if key in unsupported_kwargs:\n                        unsupported_kwargs.pop(key)\n\n        return supported_kwargs, unsupported_kwargs, peft_kwargs\n\n    @classmethod\n    def add_and_load_reward_modeling_adapter(\n        cls, pretrained_model, adapter_model_id, adapter_name=\"reward_model_adapter\", token=None\n    ):\n        r\"\"\"\n        Add and load a reward modeling adapter. This method can only be used if the\n        model is a `PeftModel` and if you have initialized the model with the `reward_modeling_adapter_id`\n        argument, pointing to the id of the reward modeling adapter. The latest needs also to contain the\n        score head in order to produce the reward.\n        \"\"\"\n        pretrained_model.load_adapter(adapter_model_id, adapter_name, is_trainable=False)\n        pretrained_model.train()\n\n        filename = os.path.join(adapter_model_id, \"adapter_model.bin\")\n        safe_loading = False\n        if not os.path.exists(filename):\n            try:\n                local_filename = hf_hub_download(\n                    adapter_model_id,\n                    \"adapter_model.bin\",\n                    token=token,\n                )\n            except Exception:\n                filename = os.path.join(adapter_model_id, \"adapter_model.safetensors\")\n                safe_loading = True\n                if not os.path.exists(filename):\n                    try:\n                        local_filename = hf_hub_download(\n                            adapter_model_id,\n                            \"adapter_model.safetensors\",\n                            token=token,\n                        )\n                    except Exception as exc:\n                        raise ValueError(\n                            \"Could not find adapter model in the Hub, \"\n                            \"make sure you have the correct adapter model id.\"\n                        ) from exc\n                else:\n                    local_filename = filename\n        else:\n            local_filename = filename\n\n        loading_func = safe_load_file if safe_loading else torch.load\n        load_kwargs = {} if safe_loading else {\"map_location\": \"cpu\", \"weights_only\": True}\n\n        adapter_state_dict = loading_func(local_filename, **load_kwargs)\n\n        for score_name_candidate in cls.supported_rm_modules:\n            if any(score_name_candidate in name for name in adapter_state_dict.keys()):\n                score_name = score_name_candidate\n                # we have found the correct head name and can break\n                break\n\n        score_dict = {}\n\n        for name, param in adapter_state_dict.items():\n            if score_name in name:\n                key_name = \".\".join(name.split(\".\")[-1:])\n                score_dict[key_name] = param.to(cls._get_current_device())\n\n        num_labels, hidden_dim = score_dict[\"weight\"].shape\n        has_bias = any(\"bias\" in name for name in adapter_state_dict.keys())\n\n        score = nn.Linear(hidden_dim, num_labels, bias=has_bias).to(\n            device=cls._get_current_device(),\n            dtype=pretrained_model.dtype,\n        )\n        score.load_state_dict(score_dict)\n        for param in score.parameters():\n            param.requires_grad = False\n\n        return score\n\n    def push_to_hub(self, *args, **kwargs):\n        r\"\"\"\n        Push the pretrained model to the hub. This method is a wrapper around\n        `transformers.PreTrainedModel.push_to_hub`. Please refer to the documentation\n        of `transformers.PreTrainedModel.push_to_hub` for more information.\n\n        Args:\n            *args (`list`, *optional*):\n                Positional arguments passed along to the underlying model's\n                `push_to_hub` method.\n            **kwargs (`dict`, *optional*):\n                Keyword arguments passed along to the underlying model's\n                `push_to_hub` method.\n        \"\"\"\n        raise NotImplementedError\n\n    def save_pretrained(self, *args, **kwargs):\n        r\"\"\"\n        Save the pretrained model to a directory. This method is a wrapper around\n        `transformers.PreTrainedModel.save_pretrained`. Please refer to the documentation\n        of `transformers.PreTrainedModel.save_pretrained` for more information.\n\n        Args:\n            *args (`list`, *optional*):\n                Positional arguments passed along to the underlying model's\n                `save_pretrained` method.\n            **kwargs (`dict`, *optional*):\n                Keyword arguments passed along to the underlying model's\n                `save_pretrained` method.\n        \"\"\"\n        state_dict = kwargs.get(\"state_dict\")\n        if state_dict is None:\n            state_dict = self.state_dict()\n            kwargs[\"state_dict\"] = state_dict\n\n        # if it is a peft model only save the `v_head` state_dict and\n        # pop the `state_dict` from the kwargs to avoid slient bugs with `peft`\n        if self.is_peft_model:\n            save_path = args[0]\n            save_path = os.path.join(save_path, \"pytorch_model.bin\")\n            torch.save(state_dict, save_path)\n            _ = kwargs.pop(\"state_dict\", None)\n\n        return self.pretrained_model.save_pretrained(*args, **kwargs)\n\n    def state_dict(self, *args, **kwargs):\n        r\"\"\"\n        Return the state_dict of the pretrained model.\n        \"\"\"\n        raise NotImplementedError\n\n    def post_init(self, *args, **kwargs):\n        r\"\"\"\n        Post initialization method. This method is called after the model is\n        instantiated and loaded from a checkpoint. It can be used to perform\n        additional operations such as loading the state_dict.\n        \"\"\"\n        raise NotImplementedError\n\n    def compute_reward_score(self, input_ids, attention_mask=None, **kwargs):\n        r\"\"\"\n        Computes the reward score for a given input. The method has first to enable the adapter\n        and then compute the reward score. After that the model disables the reward modeling\n        adapter and enables the default ppo adapter again.\n        \"\"\"\n        if not self.supports_rm_adapter:\n            raise ValueError(\"This model does not support reward modeling adapter.\")\n\n        # enable rm adapter\n        self.pretrained_model.set_adapter(self.rm_adapter_name)\n        self.pretrained_model.eval()\n\n        with torch.no_grad():\n            base_model_output = self.pretrained_model(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                output_hidden_states=True,\n                return_dict=True,\n                **kwargs,\n            )\n\n            last_hidden_states = base_model_output.hidden_states[-1]\n            scores = self.score(last_hidden_states)\n\n        self.pretrained_model.set_adapter(self.policy_adapter_name)\n        self.pretrained_model.eval()\n\n        return scores\n\n\ndef create_reference_model(\n    model: PreTrainedModelWrapper, num_shared_layers: Optional[int] = None, pattern: Optional[str] = None\n) -> PreTrainedModelWrapper:\n    \"\"\"\n    Creates a static reference copy of a model. Note that model will be in `.eval()` mode.\n\n    Args:\n        model (`PreTrainedModelWrapper`): The model to be copied.\n        num_shared_layers (`int`, *optional*): The number of initial layers that are shared between both models and kept frozen.\n        pattern (`str`, *optional*): The shared layers are selected with a string pattern\n            (e.g. \"transformer.h.{layer}\" for GPT2) and if a custom pattern is necessary it can be passed here.\n\n    Returns:\n        `PreTrainedModelWrapper`\n    \"\"\"\n    if is_deepspeed_zero3_enabled():\n        raise ValueError(\n            \"DeepSpeed ZeRO-3 is enabled and is not compatible with `create_reference_model()`. Please instantiate your reference model directly with `AutoCausalLM.from_pretrained()`.\"\n        )\n\n    parameter_names = [n for n, _ in model.named_parameters()]\n    ref_model = deepcopy(model)\n\n    # if no layers are shared, return copy of model\n    if num_shared_layers is None:\n        for param_name in parameter_names:\n            param = ref_model.get_parameter(param_name)\n            param.requires_grad = False\n        return ref_model.eval()\n\n    # identify layer name pattern\n    if pattern is not None:\n        pattern = pattern.format(layer=num_shared_layers)\n    else:\n        for pattern_candidate in LAYER_PATTERNS:\n            pattern_candidate = pattern_candidate.format(layer=num_shared_layers)\n            if any(pattern_candidate in name for name in parameter_names):\n                pattern = pattern_candidate\n                break\n\n    if pattern is None:\n        raise ValueError(\"Layer pattern could not be matched.\")\n\n    # divide parameters in shared and unshared parameter lists\n    shared_param_list = []\n    unshared_param_list = []\n\n    shared_parameter = True\n    for name, _param in model.named_parameters():\n        if pattern in name:\n            shared_parameter = False\n        if shared_parameter:\n            shared_param_list.append(name)\n        else:\n            unshared_param_list.append(name)\n\n    # create reference of the original parameter if they are shared\n    for param_name in shared_param_list:\n        param = model.get_parameter(param_name)\n        param.requires_grad = False\n\n        _ref_param = ref_model.get_parameter(param_name)\n\n    # for all other parameters just make sure they don't use gradients\n    for param_name in unshared_param_list:\n        param = ref_model.get_parameter(param_name)\n        param.requires_grad = False\n\n    if pattern is not None and len(unshared_param_list) == 0:\n        logging.warning(\"Pattern passed or found, but no layers matched in the model. Check for a typo.\")\n\n    return ref_model.eval()\n\n\nclass GeometricMixtureWrapper(GenerationMixin):\n    r\"\"\"\n    Geometric Mixture generation wrapper that samples from the logits of two model's geometric mixture.\n\n    Args:\n        model (`PreTrainedModel`): The model to be wrapped.\n        ref_model (`PreTrainedModel`): The reference model.\n        generation_config (`GenerationConfig`): The generation config.\n        mixture_coef (`float`, *optional* - default: 0.5): The mixture coefficient.\n    \"\"\"\n\n    main_input_name = \"input_ids\"\n    _supports_cache_class = False\n    _supports_static_cache = False\n\n    def __init__(self, model, ref_model, generation_config, mixture_coef=0.5, device=None):\n        super().__init__()\n\n        self.model = model.eval()\n        self.config = model.config\n        self.ref_model = ref_model.eval()\n        self.generation_config = generation_config\n        self.mixture_coef = mixture_coef\n        self.device = device\n\n    def __call__(self, *args, **kwargs):\n        return self.forward(*args, **kwargs)\n\n    @torch.no_grad()\n    def forward(self, *args, **kwargs):\n        model_outputs = self.model(*args, **kwargs)\n        model_logits = model_outputs.logits\n        ref_model_logits = self.ref_model(*args, **kwargs).logits\n\n        model_outputs.logits = torch.nn.functional.log_softmax(\n            self.mixture_coef * ref_model_logits + (1 - self.mixture_coef) * model_logits, dim=-1\n        )\n\n        return model_outputs\n\n    def prepare_inputs_for_generation(self, *args, **kwargs):\n        # turn off cache in the generation config\n        kwargs[\"use_cache\"] = False\n        model_inputs = self.model.prepare_inputs_for_generation(*args, **kwargs)\n        _ = self.ref_model.prepare_inputs_for_generation(*args, **kwargs)\n\n        return model_inputs\n\n    def _validate_model_class(self):\n        self.model._validate_model_class()\n\n    def _validate_model_kwargs(self, model_kwargs):\n        return self.model._validate_model_kwargs(model_kwargs)\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport torch\nimport torch.nn as nn\nfrom transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_npu_available, is_torch_xpu_available\n\nfrom .modeling_base import PreTrainedModelWrapper\n\n\nclass ValueHead(nn.Module):\n    r\"\"\"\n    The ValueHead class implements a head for GPT2 that returns a scalar for each output token.\n    \"\"\"\n\n    def __init__(self, config, **kwargs):\n        super().__init__()\n        if not hasattr(config, \"summary_dropout_prob\"):\n            summary_dropout_prob = kwargs.pop(\"summary_dropout_prob\", 0.1)\n        else:\n            summary_dropout_prob = config.summary_dropout_prob\n\n        self.dropout = nn.Dropout(summary_dropout_prob) if summary_dropout_prob else nn.Identity()\n\n        # some models such as OPT have a projection layer before the word embeddings - e.g. OPT-350m\n        if hasattr(config, \"hidden_size\"):\n            hidden_size = config.hidden_size\n        if hasattr(config, \"word_embed_proj_dim\"):\n            hidden_size = config.word_embed_proj_dim\n        elif hasattr(config, \"is_encoder_decoder\"):\n            if config.is_encoder_decoder and hasattr(config, \"decoder\"):\n                if hasattr(config.decoder, \"hidden_size\"):\n                    hidden_size = config.decoder.hidden_size\n\n        self.summary = nn.Linear(hidden_size, 1)\n\n        self.flatten = nn.Flatten()\n\n    def forward(self, hidden_states):\n        output = self.dropout(hidden_states)\n\n        # For now force upcast in fp32 if needed. Let's keep the\n        # output in fp32 for numerical stability.\n        if output.dtype != self.summary.weight.dtype:\n            output = output.to(self.summary.weight.dtype)\n\n        output = self.summary(output)\n        return output\n\n\nclass AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):\n    r\"\"\"\n    An autoregressive model with a value head in addition to the language model head.\n    This class inherits from `~trl.PreTrainedModelWrapper` and wraps a\n    `transformers.PreTrainedModel` class. The wrapper class supports classic functions\n    such as `from_pretrained`, `push_to_hub` and `generate`. To call a method of the wrapped\n    model, simply manipulate the `pretrained_model` attribute of this class.\n\n    Class attributes:\n        - **transformers_parent_class** (`transformers.PreTrainedModel`) -- The parent class of the wrapped model. This\n            should be set to `transformers.AutoModelForCausalLM` for this class.\n        - **lm_head_namings** (`tuple`) -- A tuple of strings that are used to identify the language model head of the\n            wrapped model. This is set to `(\"lm_head\", \"embed_out\")` for this class but can be changed for other models\n            in the future\n        - **supported_args** (`tuple`) -- A tuple of strings that are used to identify the arguments that are supported\n            by the `ValueHead` class. Currently, the supported args are:\n            - **summary_dropout_prob** (`float`, `optional`, defaults to `None`) -- The dropout probability for the\n                `ValueHead` class.\n            - **v_head_initializer_range** (`float`, `optional`, defaults to `0.2`) -- The initializer range for the\n                `ValueHead` if a specific initialization strategy is selected.\n            - **v_head_init_strategy** (`str`, `optional`, defaults to `None`) -- The initialization strategy for the\n                `ValueHead`. Currently, the supported strategies are:\n                - **`None`** -- Initializes the weights of the `ValueHead` with a random distribution. This is the default\n                    strategy.\n                - **\"normal\"** -- Initializes the weights of the `ValueHead` with a normal distribution.\n    \"\"\"\n\n    transformers_parent_class = AutoModelForCausalLM\n    lm_head_namings = [\"lm_head\", \"embed_out\"]\n    supported_args = (\n        \"summary_dropout_prob\",\n        \"v_head_initializer_range\",\n        \"v_head_init_strategy\",\n    )\n\n    def __init__(self, pretrained_model, **kwargs):\n        r\"\"\"\n        Initializes the model.\n\n        Args:\n            pretrained_model (`transformers.PreTrainedModel`):\n                The model to wrap. It should be a causal language model such as GPT2.\n                or any model mapped inside the `AutoModelForCausalLM` class.\n            kwargs (`dict`, `optional`):\n                Additional keyword arguments, that are passed to the `ValueHead` class.\n        \"\"\"\n        super().__init__(pretrained_model, **kwargs)\n        v_head_kwargs, _, _ = self._split_kwargs(kwargs)\n\n        if not any(hasattr(self.pretrained_model, attribute) for attribute in self.lm_head_namings):\n            raise ValueError(\"The model does not have a language model head, please use a model that has one.\")\n\n        self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs)\n\n        self._init_weights(**v_head_kwargs)\n\n    def _init_weights(self, **kwargs):\n        r\"\"\"\n        Initializes the weights of the value head. The default initialization strategy is random.\n        Users can pass a different initialization strategy by passing the `v_head_init_strategy` argument\n        when calling `.from_pretrained`. Supported strategies are:\n        - `normal`: initializes the weights with a normal distribution.\n\n        Args:\n            **kwargs (`dict`, `optional`):\n                Additional keyword arguments, that are passed to the `ValueHead` class. These arguments\n                can contain the `v_head_init_strategy` argument as well as the `v_head_initializer_range`\n                argument.\n        \"\"\"\n        initializer_range = kwargs.pop(\"v_head_initializer_range\", 0.2)\n        # random init by default\n        init_strategy = kwargs.pop(\"v_head_init_strategy\", None)\n        if init_strategy is None:\n            # do nothing\n            pass\n        elif init_strategy == \"normal\":\n            self.v_head.summary.weight.data.normal_(mean=0.0, std=initializer_range)\n            self.v_head.summary.bias.data.zero_()\n\n    def forward(\n        self,\n        input_ids=None,\n        past_key_values=None,\n        attention_mask=None,\n        return_past_key_values=False,\n        **kwargs,\n    ):\n        r\"\"\"\n        Applies a forward pass to the wrapped model and returns the logits of the value head.\n\n        Args:\n            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n                Indices of input sequence tokens in the vocabulary.\n            past_key_values (`tuple(tuple(torch.FloatTensor))`, `optional`):\n                Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model\n                (see `past_key_values` input) to speed up sequential decoding.\n            attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):\n                Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:\n                - 1 for tokens that are **not masked**,\n                - 0 for tokens that are **masked**.\n            return_past_key_values (bool): A flag indicating if the computed hidden-states should be returned.\n            kwargs (`dict`, `optional`):\n                Additional keyword arguments, that are passed to the wrapped model.\n        \"\"\"\n        kwargs[\"output_hidden_states\"] = True  # this had already been set in the LORA / PEFT examples\n        kwargs[\"past_key_values\"] = past_key_values\n\n        if self.is_peft_model and self.pretrained_model.active_peft_config.peft_type == \"PREFIX_TUNING\":\n            kwargs.pop(\"past_key_values\")\n\n        base_model_output = self.pretrained_model(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            **kwargs,\n        )\n\n        last_hidden_state = base_model_output.hidden_states[-1]\n        lm_logits = base_model_output.logits\n        loss = base_model_output.loss\n\n        if last_hidden_state.device != self.v_head.summary.weight.device:\n            last_hidden_state = last_hidden_state.to(self.v_head.summary.weight.device)\n\n        value = self.v_head(last_hidden_state).squeeze(-1)\n\n        # force upcast in fp32 if logits are in half-precision\n        if lm_logits.dtype != torch.float32:\n            lm_logits = lm_logits.float()\n\n        if return_past_key_values:\n            return (lm_logits, loss, value, base_model_output.past_key_values)\n        else:\n            return (lm_logits, loss, value)\n\n    def generate(self, *args, **kwargs):\n        r\"\"\"\n        A simple wrapper around the `generate` method of the wrapped model.\n        Please refer to the [`generate`](https://huggingface.co/docs/transformers/internal/generation_utils)\n        method of the wrapped model for more information about the supported arguments.\n\n        Args:\n            *args (`list`, *optional*):\n                Positional arguments passed to the `generate` method of the wrapped model.\n            **kwargs (`dict`, *optional*):\n                Keyword arguments passed to the `generate` method of the wrapped model.\n        \"\"\"\n        return self.pretrained_model.generate(*args, **kwargs)\n\n    def state_dict(self, *args, **kwargs):\n        r\"\"\"\n        Returns the state dictionary of the model. We add the state dictionary of the value head\n        to the state dictionary of the wrapped model by prepending the key with `v_head.`.\n        \"\"\"\n        if not self.is_peft_model:\n            pretrained_model_state_dict = self.pretrained_model.state_dict(*args, **kwargs)\n        else:\n            # if it is a peft model, only save the v_head\n            pretrained_model_state_dict = {}\n\n        v_head_state_dict = self.v_head.state_dict(*args, **kwargs)\n        for k, v in v_head_state_dict.items():\n            pretrained_model_state_dict[f\"v_head.{k}\"] = v\n        return pretrained_model_state_dict\n\n    def push_to_hub(self, *args, **kwargs):\n        self.pretrained_model.v_head = self.v_head\n\n        return self.pretrained_model.push_to_hub(*args, **kwargs)\n\n    def post_init(self, state_dict):\n        r\"\"\"\n        We add the state dictionary of the value head to the state dictionary of the wrapped model\n        by prepending the key with `v_head.`. This function removes the `v_head.` prefix from the\n        keys of the value head state dictionary.\n        \"\"\"\n        for k in list(state_dict.keys()):\n            if \"v_head.\" in k:\n                state_dict[k.replace(\"v_head.\", \"\")] = state_dict.pop(k)\n        self.v_head.load_state_dict(state_dict, strict=False)\n        del state_dict\n\n        if hasattr(self.pretrained_model, \"hf_device_map\"):\n            if (\n                \"cpu\" in self.pretrained_model.hf_device_map.values()\n                or \"disk\" in self.pretrained_model.hf_device_map.values()\n            ):\n                raise ValueError(\n                    \"The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models.\"\n                )\n\n            first_device = list(set(self.pretrained_model.hf_device_map.values()))[0]\n            if isinstance(first_device, int):\n                if is_torch_npu_available():\n                    first_device = f\"npu:{first_device}\"\n                elif is_torch_xpu_available():\n                    first_device = f\"xpu:{first_device}\"\n                else:\n                    first_device = f\"cuda:{first_device}\"\n            self.v_head = self.v_head.to(first_device)\n\n            def set_device_hook(module, input, outputs):\n                new_output = ()\n                for output in outputs:\n                    if isinstance(output, torch.Tensor):\n                        new_output += (output.to(first_device),)\n                    else:\n                        new_output += (output,)\n                return new_output\n\n            self.register_forward_hook(set_device_hook)\n\n            self.is_sequential_parallel = True\n\n\nclass AutoModelForSeq2SeqLMWithValueHead(PreTrainedModelWrapper):\n    r\"\"\"\n    A seq2seq model with a value head in addition to the language model head.\n    This class inherits from `~trl.PreTrainedModelWrapper` and wraps a\n    `transformers.PreTrainedModel` class. The wrapper class supports classic functions\n    such as `from_pretrained` and `push_to_hub` and also provides some additional\n    functionalities such as `generate`.\n\n    Args:\n        pretrained_model (`transformers.PreTrainedModel`):\n            The model to wrap. It should be a causal language model such as GPT2.\n            or any model mapped inside the `AutoModelForSeq2SeqLM` class.\n        kwargs:\n            Additional keyword arguments passed along to the `ValueHead` class.\n    \"\"\"\n\n    transformers_parent_class = AutoModelForSeq2SeqLM\n    lm_head_namings = [\"lm_head\", \"embed_out\", \"output_projection\"]\n    supported_args = (\n        \"summary_dropout_prob\",\n        \"v_head_initializer_range\",\n        \"v_head_init_strategy\",\n    )\n\n    def __init__(self, pretrained_model, **kwargs):\n        super().__init__(pretrained_model, **kwargs)\n        v_head_kwargs, _, _ = self._split_kwargs(kwargs)\n        self.is_encoder_decoder = True\n\n        if not self._has_lm_head():\n            raise ValueError(\"The model does not have a language model head, please use a model that has one.\")\n\n        self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs)\n\n        self._init_weights(**v_head_kwargs)\n\n    def _has_lm_head(self):\n        # check module names of all modules inside `pretrained_model` to find the language model head\n        for name, _module in self.pretrained_model.named_modules():\n            if any(attribute in name for attribute in self.lm_head_namings):\n                return True\n        return False\n\n    def post_init(self, state_dict):\n        r\"\"\"\n        We add the state dictionary of the value head to the state dictionary of the wrapped model\n        by prepending the key with `v_head.`. This function removes the `v_head.` prefix from the\n        keys of the value head state dictionary.\n        \"\"\"\n        for k in list(state_dict.keys()):\n            if \"v_head.\" in k:\n                state_dict[k.replace(\"v_head.\", \"\")] = state_dict.pop(k)\n        self.v_head.load_state_dict(state_dict, strict=False)\n        del state_dict\n\n        if hasattr(self.pretrained_model, \"hf_device_map\"):\n            if (\n                \"cpu\" in self.pretrained_model.hf_device_map.values()\n                or \"disk\" in self.pretrained_model.hf_device_map.values()\n            ):\n                raise ValueError(\n                    \"The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models.\"\n                )\n\n            # get the lm_head device\n            for name, module in self.pretrained_model.named_modules():\n                if any(attribute in name for attribute in self.lm_head_namings):\n                    lm_head_device = module.weight.device\n                    break\n\n            # put v_head on the same device as the lm_head to avoid issues\n            self.v_head = self.v_head.to(lm_head_device)\n\n            def set_device_hook(module, input, outputs):\n                r\"\"\"\n                A hook that sets the device of the output of the model to the device of the first\n                parameter of the model.\n\n                Args:\n                    module (`nn.Module`):\n                        The module to which the hook is attached.\n                    input (`tuple`):\n                        The input to the module.\n                    outputs (`tuple`):\n                        The output of the module.\n                \"\"\"\n                new_output = ()\n                for output in outputs:\n                    if isinstance(output, torch.Tensor):\n                        new_output += (output.to(lm_head_device),)\n                    else:\n                        new_output += (output,)\n                return new_output\n\n            self.register_forward_hook(set_device_hook)\n            self.is_sequential_parallel = True\n\n    def state_dict(self, *args, **kwargs):\n        r\"\"\"\n        Returns the state dictionary of the model. We add the state dictionary of the value head\n        to the state dictionary of the wrapped model by prepending the key with `v_head.`.\n        \"\"\"\n        if not self.is_peft_model:\n            pretrained_model_state_dict = self.pretrained_model.state_dict(*args, **kwargs)\n        else:\n            # if it is a peft model, only save the v_head\n            pretrained_model_state_dict = {}\n\n        v_head_state_dict = self.v_head.state_dict(*args, **kwargs)\n        for k, v in v_head_state_dict.items():\n            pretrained_model_state_dict[f\"v_head.{k}\"] = v\n        return pretrained_model_state_dict\n\n    def push_to_hub(self, *args, **kwargs):\n        self.pretrained_model.v_head = self.v_head\n\n        return self.pretrained_model.push_to_hub(*args, **kwargs)\n\n    def _init_weights(self, **kwargs):\n        r\"\"\"\n        We initialize the weights of the value head.\n        \"\"\"\n        initializer_range = kwargs.pop(\"v_head_initializer_range\", 0.2)\n        # random init by default\n        init_strategy = kwargs.pop(\"v_head_init_strategy\", None)\n        if init_strategy is None:\n            # do nothing\n            pass\n        elif init_strategy == \"normal\":\n            self.v_head.summary.weight.data.normal_(mean=0.0, std=initializer_range)\n            self.v_head.summary.bias.data.zero_()\n\n    def forward(\n        self,\n        input_ids=None,\n        past_key_values=None,\n        attention_mask=None,\n        return_past_key_values=False,\n        **kwargs,\n    ):\n        kwargs[\"past_key_values\"] = past_key_values\n        if self.is_peft_model and self.pretrained_model.active_peft_config.peft_type == \"PREFIX_TUNING\":\n            kwargs.pop(\"past_key_values\")\n\n        base_model_output = self.pretrained_model(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            output_hidden_states=True,  # We force the model to output hidden states\n            **kwargs,\n        )\n\n        last_hidden_state = base_model_output.decoder_hidden_states[-1]\n        lm_logits = base_model_output.logits\n        loss = base_model_output.loss\n\n        value = self.v_head(last_hidden_state).squeeze(-1)\n\n        # force upcast in fp32 if logits are in half-precision\n        if lm_logits.dtype != torch.float32:\n            lm_logits = lm_logits.float()\n\n        if return_past_key_values:\n            return (lm_logits, loss, value, base_model_output.past_key_values)\n        else:\n            return (lm_logits, loss, value)\n\n    def generate(self, *args, **kwargs):\n        r\"\"\"\n        We call `generate` on the wrapped model.\n        \"\"\"\n        return self.pretrained_model.generate(*args, **kwargs)\n\n\n# flake8: noqa\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# flake8: noqa\n\nfrom typing import TYPE_CHECKING\nfrom ..import_utils import _LazyModule, is_diffusers_available, OptionalDependencyNotAvailable\n\n\n_import_structure = {\n    \"modeling_base\": [\"PreTrainedModelWrapper\", \"create_reference_model\", \"GeometricMixtureWrapper\"],\n    \"modeling_value_head\": [\n        \"AutoModelForCausalLMWithValueHead\",\n        \"AutoModelForSeq2SeqLMWithValueHead\",\n    ],\n    \"utils\": [\"setup_chat_format\", \"SUPPORTED_ARCHITECTURES\", \"unwrap_model_for_generation\"],\n}\n\ntry:\n    if not is_diffusers_available():\n        raise OptionalDependencyNotAvailable()\nexcept OptionalDependencyNotAvailable:\n    pass\nelse:\n    _import_structure[\"modeling_sd_base\"] = [\n        \"DDPOPipelineOutput\",\n        \"DDPOSchedulerOutput\",\n        \"DDPOStableDiffusionPipeline\",\n        \"DefaultDDPOStableDiffusionPipeline\",\n    ]\n\nif TYPE_CHECKING:\n    from .modeling_base import PreTrainedModelWrapper, create_reference_model, GeometricMixtureWrapper\n    from .modeling_value_head import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead\n    from .utils import setup_chat_format, SUPPORTED_ARCHITECTURES\n\n    try:\n        if not is_diffusers_available():\n            raise OptionalDependencyNotAvailable()\n    except OptionalDependencyNotAvailable:\n        pass\n    else:\n        from .modeling_sd_base import (\n            DDPOPipelineOutput,\n            DDPOSchedulerOutput,\n            DDPOStableDiffusionPipeline,\n            DefaultDDPOStableDiffusionPipeline,\n        )\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(__name__, globals()[\"__file__\"], _import_structure, module_spec=__spec__)\n\n\n# Copyright 2023 DDPO-pytorch authors (Kevin Black), The HuggingFace Team, metric-space. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport contextlib\nimport os\nimport random\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nimport numpy as np\nimport torch\nimport torch.utils.checkpoint as checkpoint\nfrom diffusers import DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel\nfrom diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg\nfrom transformers.utils import is_peft_available\n\nfrom ..core import randn_tensor\nfrom .sd_utils import convert_state_dict_to_diffusers\n\n\nif is_peft_available():\n    from peft import LoraConfig\n    from peft.utils import get_peft_model_state_dict\n\n\n@dataclass\nclass DDPOPipelineOutput:\n    \"\"\"\n    Output class for the diffusers pipeline to be finetuned with the DDPO trainer\n\n    Args:\n        images (`torch.Tensor`):\n            The generated images.\n        latents (`List[torch.Tensor]`):\n            The latents used to generate the images.\n        log_probs (`List[torch.Tensor]`):\n            The log probabilities of the latents.\n\n    \"\"\"\n\n    images: torch.Tensor\n    latents: torch.Tensor\n    log_probs: torch.Tensor\n\n\n@dataclass\nclass DDPOSchedulerOutput:\n    \"\"\"\n    Output class for the diffusers scheduler to be finetuned with the DDPO trainer\n\n    Args:\n        latents (`torch.Tensor`):\n            Predicted sample at the previous timestep. Shape: `(batch_size, num_channels, height, width)`\n        log_probs (`torch.Tensor`):\n            Log probability of the above mentioned sample. Shape: `(batch_size)`\n    \"\"\"\n\n    latents: torch.Tensor\n    log_probs: torch.Tensor\n\n\nclass DDPOStableDiffusionPipeline:\n    \"\"\"\n    Main class for the diffusers pipeline to be finetuned with the DDPO trainer\n    \"\"\"\n\n    def __call__(self, *args, **kwargs) -> DDPOPipelineOutput:\n        raise NotImplementedError\n\n    def scheduler_step(self, *args, **kwargs) -> DDPOSchedulerOutput:\n        raise NotImplementedError\n\n    @property\n    def unet(self):\n        \"\"\"\n        Returns the 2d U-Net model used for diffusion.\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def vae(self):\n        \"\"\"\n        Returns the Variational Autoencoder model used from mapping images to and from the latent space\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def tokenizer(self):\n        \"\"\"\n        Returns the tokenizer used for tokenizing text inputs\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def scheduler(self):\n        \"\"\"\n        Returns the scheduler associated with the pipeline used for the diffusion process\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def text_encoder(self):\n        \"\"\"\n        Returns the text encoder used for encoding text inputs\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    def autocast(self):\n        \"\"\"\n        Returns the autocast context manager\n        \"\"\"\n        raise NotImplementedError\n\n    def set_progress_bar_config(self, *args, **kwargs):\n        \"\"\"\n        Sets the progress bar config for the pipeline\n        \"\"\"\n        raise NotImplementedError\n\n    def save_pretrained(self, *args, **kwargs):\n        \"\"\"\n        Saves all of the model weights\n        \"\"\"\n        raise NotImplementedError\n\n    def get_trainable_layers(self, *args, **kwargs):\n        \"\"\"\n        Returns the trainable parameters of the pipeline\n        \"\"\"\n        raise NotImplementedError\n\n    def save_checkpoint(self, *args, **kwargs):\n        \"\"\"\n        Light wrapper around accelerate's register_save_state_pre_hook which is run before saving state\n        \"\"\"\n        raise NotImplementedError\n\n    def load_checkpoint(self, *args, **kwargs):\n        \"\"\"\n        Light wrapper around accelerate's register_lad_state_pre_hook which is run before loading state\n        \"\"\"\n        raise NotImplementedError\n\n\ndef _left_broadcast(input_tensor, shape):\n    \"\"\"\n    As opposed to the default direction of broadcasting (right to left), this function broadcasts\n    from left to right\n        Args:\n            input_tensor (`torch.FloatTensor`): is the tensor to broadcast\n            shape (`Tuple[int]`): is the shape to broadcast to\n    \"\"\"\n    input_ndim = input_tensor.ndim\n    if input_ndim > len(shape):\n        raise ValueError(\n            \"The number of dimensions of the tensor to broadcast cannot be greater than the length of the shape to broadcast to\"\n        )\n    return input_tensor.reshape(input_tensor.shape + (1,) * (len(shape) - input_ndim)).broadcast_to(shape)\n\n\ndef _get_variance(self, timestep, prev_timestep):\n    alpha_prod_t = torch.gather(self.alphas_cumprod, 0, timestep.cpu()).to(timestep.device)\n    alpha_prod_t_prev = torch.where(\n        prev_timestep.cpu() >= 0,\n        self.alphas_cumprod.gather(0, prev_timestep.cpu()),\n        self.final_alpha_cumprod,\n    ).to(timestep.device)\n    beta_prod_t = 1 - alpha_prod_t\n    beta_prod_t_prev = 1 - alpha_prod_t_prev\n\n    variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)\n\n    return variance\n\n\ndef scheduler_step(\n    self,\n    model_output: torch.FloatTensor,\n    timestep: int,\n    sample: torch.FloatTensor,\n    eta: float = 0.0,\n    use_clipped_model_output: bool = False,\n    generator=None,\n    prev_sample: Optional[torch.FloatTensor] = None,\n) -> DDPOSchedulerOutput:\n    \"\"\"\n\n    Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion\n    process from the learned model outputs (most often the predicted noise).\n\n    Args:\n        model_output (`torch.FloatTensor`): direct output from learned diffusion model.\n        timestep (`int`): current discrete timestep in the diffusion chain.\n        sample (`torch.FloatTensor`):\n            current instance of sample being created by diffusion process.\n        eta (`float`): weight of noise for added noise in diffusion step.\n        use_clipped_model_output (`bool`): if `True`, compute \"corrected\" `model_output` from the clipped\n            predicted original sample. Necessary because predicted original sample is clipped to [-1, 1] when\n            `self.config.clip_sample` is `True`. If no clipping has happened, \"corrected\" `model_output` would\n            coincide with the one provided as input and `use_clipped_model_output` will have not effect.\n        generator: random number generator.\n        variance_noise (`torch.FloatTensor`): instead of generating noise for the variance using `generator`, we\n            can directly provide the noise for the variance itself. This is useful for methods such as\n            CycleDiffusion. (https://huggingface.co/papers/2210.05559)\n\n    Returns:\n        `DDPOSchedulerOutput`: the predicted sample at the previous timestep and the log probability of the sample\n    \"\"\"\n\n    if self.num_inference_steps is None:\n        raise ValueError(\n            \"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler\"\n        )\n\n    # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502\n    # Ideally, read DDIM paper in-detail understanding\n\n    # Notation (<variable name> -> <name in paper>\n    # - pred_noise_t -> e_theta(x_t, t)\n    # - pred_original_sample -> f_theta(x_t, t) or x_0\n    # - std_dev_t -> sigma_t\n    # - eta -> η\n    # - pred_sample_direction -> \"direction pointing to x_t\"\n    # - pred_prev_sample -> \"x_t-1\"\n    # 1. get previous step value (=t-1)\n    prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps\n    # to prevent OOB on gather\n    prev_timestep = torch.clamp(prev_timestep, 0, self.config.num_train_timesteps - 1)\n\n    # 2. compute alphas, betas\n    alpha_prod_t = self.alphas_cumprod.gather(0, timestep.cpu())\n    alpha_prod_t_prev = torch.where(\n        prev_timestep.cpu() >= 0,\n        self.alphas_cumprod.gather(0, prev_timestep.cpu()),\n        self.final_alpha_cumprod,\n    )\n    alpha_prod_t = _left_broadcast(alpha_prod_t, sample.shape).to(sample.device)\n    alpha_prod_t_prev = _left_broadcast(alpha_prod_t_prev, sample.shape).to(sample.device)\n\n    beta_prod_t = 1 - alpha_prod_t\n\n    # 3. compute predicted original sample from predicted noise also called\n    # \"predicted x_0\" of formula (12) from https://huggingface.co/papers/2010.02502\n    if self.config.prediction_type == \"epsilon\":\n        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)\n        pred_epsilon = model_output\n    elif self.config.prediction_type == \"sample\":\n        pred_original_sample = model_output\n        pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)\n    elif self.config.prediction_type == \"v_prediction\":\n        pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output\n        pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample\n    else:\n        raise ValueError(\n            f\"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or\"\n            \" `v_prediction`\"\n        )\n\n    # 4. Clip or threshold \"predicted x_0\"\n    if self.config.thresholding:\n        pred_original_sample = self._threshold_sample(pred_original_sample)\n    elif self.config.clip_sample:\n        pred_original_sample = pred_original_sample.clamp(\n            -self.config.clip_sample_range, self.config.clip_sample_range\n        )\n\n    # 5. compute variance: \"sigma_t(η)\" -> see formula (16)\n    # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)\n    variance = _get_variance(self, timestep, prev_timestep)\n    std_dev_t = eta * variance ** (0.5)\n    std_dev_t = _left_broadcast(std_dev_t, sample.shape).to(sample.device)\n\n    if use_clipped_model_output:\n        # the pred_epsilon is always re-derived from the clipped x_0 in Glide\n        pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)\n\n    # 6. compute \"direction pointing to x_t\" of formula (12) from https://huggingface.co/papers/2010.02502\n    pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon\n\n    # 7. compute x_t without \"random noise\" of formula (12) from https://huggingface.co/papers/2010.02502\n    prev_sample_mean = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction\n\n    if prev_sample is not None and generator is not None:\n        raise ValueError(\n            \"Cannot pass both generator and prev_sample. Please make sure that either `generator` or\"\n            \" `prev_sample` stays `None`.\"\n        )\n\n    if prev_sample is None:\n        variance_noise = randn_tensor(\n            model_output.shape,\n            generator=generator,\n            device=model_output.device,\n            dtype=model_output.dtype,\n        )\n        prev_sample = prev_sample_mean + std_dev_t * variance_noise\n\n    # log prob of prev_sample given prev_sample_mean and std_dev_t\n    log_prob = (\n        -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * (std_dev_t**2))\n        - torch.log(std_dev_t)\n        - torch.log(torch.sqrt(2 * torch.as_tensor(np.pi)))\n    )\n    # mean along all but batch dimension\n    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))\n\n    return DDPOSchedulerOutput(prev_sample.type(sample.dtype), log_prob)\n\n\n# 1. The output type for call is different as the logprobs are now returned\n# 2. An extra method called `scheduler_step` is added which is used to constraint the scheduler output\n@torch.no_grad()\ndef pipeline_step(\n    self,\n    prompt: Optional[Union[str, List[str]]] = None,\n    height: Optional[int] = None,\n    width: Optional[int] = None,\n    num_inference_steps: int = 50,\n    guidance_scale: float = 7.5,\n    negative_prompt: Optional[Union[str, List[str]]] = None,\n    num_images_per_prompt: Optional[int] = 1,\n    eta: float = 0.0,\n    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n    latents: Optional[torch.FloatTensor] = None,\n    prompt_embeds: Optional[torch.FloatTensor] = None,\n    negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n    output_type: Optional[str] = \"pil\",\n    return_dict: bool = True,\n    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n    callback_steps: int = 1,\n    cross_attention_kwargs: Optional[Dict[str, Any]] = None,\n    guidance_rescale: float = 0.0,\n):\n    r\"\"\"\n    Function invoked when calling the pipeline for generation.  Args: prompt (`str` or `List[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.  instead.  height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image.\n        width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):\n            The width in pixels of the generated image.\n        num_inference_steps (`int`, *optional*, defaults to 50):\n            The number of denoising steps. More denoising steps usually lead to a higher quality image at the\n            expense of slower inference.\n        guidance_scale (`float`, *optional*, defaults to 7.5):\n            Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).\n            `guidance_scale` is defined as `w` of equation 2. of [Imagen\n            Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >\n            1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,\n            usually at the expense of lower image quality.\n        negative_prompt (`str` or `List[str]`, *optional*):\n            The prompt or prompts not to guide the image generation. If not defined, one has to pass\n            `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n            less than `1`).\n        num_images_per_prompt (`int`, *optional*, defaults to 1):\n            The number of images to generate per prompt.\n        eta (`float`, *optional*, defaults to 0.0):\n            Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to\n            [`schedulers.DDIMScheduler`], will be ignored for others.\n        generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n            One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n            to make generation deterministic.\n        latents (`torch.FloatTensor`, *optional*):\n            Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image\n            generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n            tensor will ge generated by sampling using the supplied random `generator`.\n        prompt_embeds (`torch.FloatTensor`, *optional*):\n            Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n            provided, text embeddings will be generated from `prompt` input argument.\n        negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n            Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n            weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n            argument.\n        output_type (`str`, *optional*, defaults to `\"pil\"`):\n            The output format of the generate image. Choose between\n            [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.\n        return_dict (`bool`, *optional*, defaults to `True`):\n            Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a\n            plain tuple.\n        callback (`Callable`, *optional*):\n            A function that will be called every `callback_steps` steps during inference. The function will be\n            called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n        callback_steps (`int`, *optional*, defaults to 1):\n            The frequency at which the `callback` function will be called. If not specified, the callback will be\n            called at every step.\n        cross_attention_kwargs (`dict`, *optional*):\n            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under\n            `self.processor` in\n            [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).\n        guidance_rescale (`float`, *optional*, defaults to 0.7):\n            Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are\n            Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of\n            [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891).\n            Guidance rescale factor should fix overexposure when using zero terminal SNR.\n\n    Examples:\n\n    Returns:\n        `DDPOPipelineOutput`: The generated image, the predicted latents used to generate the image and the associated log probabilities\n    \"\"\"\n    # 0. Default height and width to unet\n    height = height or self.unet.config.sample_size * self.vae_scale_factor\n    width = width or self.unet.config.sample_size * self.vae_scale_factor\n\n    # 1. Check inputs. Raise error if not correct\n    self.check_inputs(\n        prompt,\n        height,\n        width,\n        callback_steps,\n        negative_prompt,\n        prompt_embeds,\n        negative_prompt_embeds,\n    )\n\n    # 2. Define call parameters\n    if prompt is not None and isinstance(prompt, str):\n        batch_size = 1\n    elif prompt is not None and isinstance(prompt, list):\n        batch_size = len(prompt)\n    else:\n        batch_size = prompt_embeds.shape[0]\n\n    device = self._execution_device\n    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n    # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`\n    # corresponds to doing no classifier free guidance.\n    do_classifier_free_guidance = guidance_scale > 1.0\n\n    # 3. Encode input prompt\n    text_encoder_lora_scale = cross_attention_kwargs.get(\"scale\", None) if cross_attention_kwargs is not None else None\n    prompt_embeds = self._encode_prompt(\n        prompt,\n        device,\n        num_images_per_prompt,\n        do_classifier_free_guidance,\n        negative_prompt,\n        prompt_embeds=prompt_embeds,\n        negative_prompt_embeds=negative_prompt_embeds,\n        lora_scale=text_encoder_lora_scale,\n    )\n\n    # 4. Prepare timesteps\n    self.scheduler.set_timesteps(num_inference_steps, device=device)\n    timesteps = self.scheduler.timesteps\n\n    # 5. Prepare latent variables\n    num_channels_latents = self.unet.config.in_channels\n    latents = self.prepare_latents(\n        batch_size * num_images_per_prompt,\n        num_channels_latents,\n        height,\n        width,\n        prompt_embeds.dtype,\n        device,\n        generator,\n        latents,\n    )\n\n    # 6. Denoising loop\n    num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n    all_latents = [latents]\n    all_log_probs = []\n    with self.progress_bar(total=num_inference_steps) as progress_bar:\n        for i, t in enumerate(timesteps):\n            # expand the latents if we are doing classifier free guidance\n            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n\n            # predict the noise residual\n            noise_pred = self.unet(\n                latent_model_input,\n                t,\n                encoder_hidden_states=prompt_embeds,\n                cross_attention_kwargs=cross_attention_kwargs,\n                return_dict=False,\n            )[0]\n\n            # perform guidance\n            if do_classifier_free_guidance:\n                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n\n            if do_classifier_free_guidance and guidance_rescale > 0.0:\n                # Based on 3.4. in https://huggingface.co/papers/2305.08891\n                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)\n\n            # compute the previous noisy sample x_t -> x_t-1\n            scheduler_output = scheduler_step(self.scheduler, noise_pred, t, latents, eta)\n            latents = scheduler_output.latents\n            log_prob = scheduler_output.log_probs\n\n            all_latents.append(latents)\n            all_log_probs.append(log_prob)\n\n            # call the callback, if provided\n            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n                progress_bar.update()\n                if callback is not None and i % callback_steps == 0:\n                    callback(i, t, latents)\n\n    if not output_type == \"latent\":\n        image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]\n        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)\n    else:\n        image = latents\n        has_nsfw_concept = None\n\n    if has_nsfw_concept is None:\n        do_denormalize = [True] * image.shape[0]\n    else:\n        do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]\n\n    image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)\n\n    # Offload last model to CPU\n    if hasattr(self, \"final_offload_hook\") and self.final_offload_hook is not None:\n        self.final_offload_hook.offload()\n\n    return DDPOPipelineOutput(image, all_latents, all_log_probs)\n\n\ndef pipeline_step_with_grad(\n    pipeline,\n    prompt: Optional[Union[str, List[str]]] = None,\n    height: Optional[int] = None,\n    width: Optional[int] = None,\n    num_inference_steps: int = 50,\n    guidance_scale: float = 7.5,\n    truncated_backprop: bool = True,\n    truncated_backprop_rand: bool = True,\n    gradient_checkpoint: bool = True,\n    truncated_backprop_timestep: int = 49,\n    truncated_rand_backprop_minmax: tuple = (0, 50),\n    negative_prompt: Optional[Union[str, List[str]]] = None,\n    num_images_per_prompt: Optional[int] = 1,\n    eta: float = 0.0,\n    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n    latents: Optional[torch.FloatTensor] = None,\n    prompt_embeds: Optional[torch.FloatTensor] = None,\n    negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n    output_type: Optional[str] = \"pil\",\n    return_dict: bool = True,\n    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n    callback_steps: int = 1,\n    cross_attention_kwargs: Optional[Dict[str, Any]] = None,\n    guidance_rescale: float = 0.0,\n):\n    r\"\"\"\n    Function to get RGB image with gradients attached to the model weights.\n\n    Args:\n        prompt (`str` or `List[str]`, *optional*, defaults to `None`):\n            The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` instead.\n        height (`int`, *optional*, defaults to `pipeline.unet.config.sample_size * pipeline.vae_scale_factor`):\n            The height in pixels of the generated image.\n        width (`int`, *optional*, defaults to `pipeline.unet.config.sample_size * pipeline.vae_scale_factor`):\n            The width in pixels of the generated image.\n        num_inference_steps (`int`, *optional*, defaults to `50`):\n            The number of denoising steps. More denoising steps usually lead to a higher quality image at the\n            expense of slower inference.\n        guidance_scale (`float`, *optional*, defaults to `7.5`):\n            Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).\n            `guidance_scale` is defined as `w` of equation 2. of [Imagen\n            Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >\n            1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,\n            usually at the expense of lower image quality.\n        truncated_backprop (`bool`, *optional*, defaults to True):\n            Truncated Backpropation to fixed timesteps, helps prevent collapse during diffusion reward training as shown in AlignProp (https://huggingface.co/papers/2310.03739).\n        truncated_backprop_rand (`bool`, *optional*, defaults to True):\n            Truncated Randomized Backpropation randomizes truncation to different diffusion timesteps, this helps prevent collapse during diffusion reward training as shown in AlignProp (https://huggingface.co/papers/2310.03739).\n            Enabling truncated_backprop_rand allows adapting earlier timesteps in diffusion while not resulting in a collapse.\n        gradient_checkpoint (`bool`, *optional*, defaults to True):\n            Adds gradient checkpointing to Unet forward pass. Reduces GPU memory consumption while slightly increasing the training time.\n        truncated_backprop_timestep (`int`, *optional*, defaults to 49):\n            Absolute timestep to which the gradients are being backpropagated. Higher number reduces the memory usage and reduces the chances of collapse.\n            While a lower value, allows more semantic changes in the diffusion generations, as the earlier diffusion timesteps are getting updated.\n            However it also increases the chances of collapse.\n        truncated_rand_backprop_minmax (`Tuple`, *optional*, defaults to (0,50)):\n            Range for randomized backprop. Here the value at 0 index indicates the earlier diffusion timestep to update (closer to noise), while the value\n            at index 1 indicates the later diffusion timestep to update.\n        negative_prompt (`str` or `List[str]`, *optional*):\n            The prompt or prompts not to guide the image generation. If not defined, one has to pass\n            `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n            less than `1`).\n        num_images_per_prompt (`int`, *optional*, defaults to 1):\n            The number of images to generate per prompt.\n        eta (`float`, *optional*, defaults to 0.0):\n            Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to\n            [`schedulers.DDIMScheduler`], will be ignored for others.\n        generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n            One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n            to make generation deterministic.\n        latents (`torch.FloatTensor`, *optional*):\n            Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image\n            generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n            tensor will ge generated by sampling using the supplied random `generator`.\n        prompt_embeds (`torch.FloatTensor`, *optional*):\n            Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n            provided, text embeddings will be generated from `prompt` input argument.\n        negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n            Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n            weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n            argument.\n        output_type (`str`, *optional*, defaults to `\"pil\"`):\n            The output format of the generate image. Choose between\n            [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.\n        return_dict (`bool`, *optional*, defaults to `True`):\n            Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a\n            plain tuple.\n        callback (`Callable`, *optional*):\n            A function that will be called every `callback_steps` steps during inference. The function will be\n            called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n        callback_steps (`int`, *optional*, defaults to 1):\n            The frequency at which the `callback` function will be called. If not specified, the callback will be\n            called at every step.\n        cross_attention_kwargs (`dict`, *optional*):\n            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under\n            `pipeline.processor` in\n            [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).\n        guidance_rescale (`float`, *optional*, defaults to 0.7):\n            Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are\n            Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of\n            [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891).\n            Guidance rescale factor should fix overexposure when using zero terminal SNR.\n\n    Examples:\n\n    Returns:\n        `DDPOPipelineOutput`: The generated image, the predicted latents used to generate the image and the associated log probabilities\n    \"\"\"\n    # 0. Default height and width to unet\n    height = height or pipeline.unet.config.sample_size * pipeline.vae_scale_factor\n    width = width or pipeline.unet.config.sample_size * pipeline.vae_scale_factor\n\n    with torch.no_grad():\n        # 1. Check inputs. Raise error if not correct\n        pipeline.check_inputs(\n            prompt,\n            height,\n            width,\n            callback_steps,\n            negative_prompt,\n            prompt_embeds,\n            negative_prompt_embeds,\n        )\n\n        # 2. Define call parameters\n        if prompt is not None and isinstance(prompt, str):\n            batch_size = 1\n        elif prompt is not None and isinstance(prompt, list):\n            batch_size = len(prompt)\n        else:\n            batch_size = prompt_embeds.shape[0]\n\n        device = pipeline._execution_device\n        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n        # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`\n        # corresponds to doing no classifier free guidance.\n        do_classifier_free_guidance = guidance_scale > 1.0\n\n        # 3. Encode input prompt\n        text_encoder_lora_scale = (\n            cross_attention_kwargs.get(\"scale\", None) if cross_attention_kwargs is not None else None\n        )\n        prompt_embeds = pipeline._encode_prompt(\n            prompt,\n            device,\n            num_images_per_prompt,\n            do_classifier_free_guidance,\n            negative_prompt,\n            prompt_embeds=prompt_embeds,\n            negative_prompt_embeds=negative_prompt_embeds,\n            lora_scale=text_encoder_lora_scale,\n        )\n\n        # 4. Prepare timesteps\n        pipeline.scheduler.set_timesteps(num_inference_steps, device=device)\n        timesteps = pipeline.scheduler.timesteps\n\n        # 5. Prepare latent variables\n        num_channels_latents = pipeline.unet.config.in_channels\n        latents = pipeline.prepare_latents(\n            batch_size * num_images_per_prompt,\n            num_channels_latents,\n            height,\n            width,\n            prompt_embeds.dtype,\n            device,\n            generator,\n            latents,\n        )\n    # 6. Denoising loop\n    num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order\n    all_latents = [latents]\n    all_log_probs = []\n    with pipeline.progress_bar(total=num_inference_steps) as progress_bar:\n        for i, t in enumerate(timesteps):\n            # expand the latents if we are doing classifier free guidance\n            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n            latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t)\n\n            # predict the noise residual\n            if gradient_checkpoint:\n                noise_pred = checkpoint.checkpoint(\n                    pipeline.unet,\n                    latent_model_input,\n                    t,\n                    prompt_embeds,\n                    cross_attention_kwargs=cross_attention_kwargs,\n                    use_reentrant=False,\n                )[0]\n            else:\n                noise_pred = pipeline.unet(\n                    latent_model_input,\n                    t,\n                    encoder_hidden_states=prompt_embeds,\n                    cross_attention_kwargs=cross_attention_kwargs,\n                    return_dict=False,\n                )[0]\n\n            #  truncating backpropagation is critical for preventing overoptimization (https://huggingface.co/papers/2304.05977).\n            if truncated_backprop:\n                # Randomized truncation randomizes the truncation process (https://huggingface.co/papers/2310.03739)\n                # the range of truncation is defined by truncated_rand_backprop_minmax\n                # Setting truncated_rand_backprop_minmax[0] to be low will allow the model to update earlier timesteps in the diffusion chain, while setitng it high will reduce the memory usage.\n                if truncated_backprop_rand:\n                    rand_timestep = random.randint(\n                        truncated_rand_backprop_minmax[0], truncated_rand_backprop_minmax[1]\n                    )\n                    if i < rand_timestep:\n                        noise_pred = noise_pred.detach()\n                else:\n                    # fixed truncation process\n                    if i < truncated_backprop_timestep:\n                        noise_pred = noise_pred.detach()\n\n            # perform guidance\n            if do_classifier_free_guidance:\n                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n\n            if do_classifier_free_guidance and guidance_rescale > 0.0:\n                # Based on 3.4. in https://huggingface.co/papers/2305.08891\n                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)\n\n            # compute the previous noisy sample x_t -> x_t-1\n            scheduler_output = scheduler_step(pipeline.scheduler, noise_pred, t, latents, eta)\n            latents = scheduler_output.latents\n            log_prob = scheduler_output.log_probs\n\n            all_latents.append(latents)\n            all_log_probs.append(log_prob)\n\n            # call the callback, if provided\n            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0):\n                progress_bar.update()\n                if callback is not None and i % callback_steps == 0:\n                    callback(i, t, latents)\n\n    if not output_type == \"latent\":\n        image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0]\n        image, has_nsfw_concept = pipeline.run_safety_checker(image, device, prompt_embeds.dtype)\n    else:\n        image = latents\n        has_nsfw_concept = None\n\n    if has_nsfw_concept is None:\n        do_denormalize = [True] * image.shape[0]\n    else:\n        do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]\n\n    image = pipeline.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)\n\n    # Offload last model to CPU\n    if hasattr(pipeline, \"final_offload_hook\") and pipeline.final_offload_hook is not None:\n        pipeline.final_offload_hook.offload()\n\n    return DDPOPipelineOutput(image, all_latents, all_log_probs)\n\n\nclass DefaultDDPOStableDiffusionPipeline(DDPOStableDiffusionPipeline):\n    def __init__(self, pretrained_model_name: str, *, pretrained_model_revision: str = \"main\", use_lora: bool = True):\n        self.sd_pipeline = StableDiffusionPipeline.from_pretrained(\n            pretrained_model_name, revision=pretrained_model_revision\n        )\n\n        self.use_lora = use_lora\n        self.pretrained_model = pretrained_model_name\n        self.pretrained_revision = pretrained_model_revision\n\n        try:\n            self.sd_pipeline.load_lora_weights(\n                pretrained_model_name,\n                weight_name=\"pytorch_lora_weights.safetensors\",\n                revision=pretrained_model_revision,\n            )\n            self.use_lora = True\n        except OSError:\n            if use_lora:\n                warnings.warn(\n                    \"If you are aware that the pretrained model has no lora weights to it, ignore this message. \"\n                    \"Otherwise please check the if `pytorch_lora_weights.safetensors` exists in the model folder.\"\n                )\n\n        self.sd_pipeline.scheduler = DDIMScheduler.from_config(self.sd_pipeline.scheduler.config)\n        self.sd_pipeline.safety_checker = None\n\n        # memory optimization\n        self.sd_pipeline.vae.requires_grad_(False)\n        self.sd_pipeline.text_encoder.requires_grad_(False)\n        self.sd_pipeline.unet.requires_grad_(not self.use_lora)\n\n    def __call__(self, *args, **kwargs) -> DDPOPipelineOutput:\n        return pipeline_step(self.sd_pipeline, *args, **kwargs)\n\n    def rgb_with_grad(self, *args, **kwargs) -> DDPOPipelineOutput:\n        return pipeline_step_with_grad(self.sd_pipeline, *args, **kwargs)\n\n    def scheduler_step(self, *args, **kwargs) -> DDPOSchedulerOutput:\n        return scheduler_step(self.sd_pipeline.scheduler, *args, **kwargs)\n\n    @property\n    def unet(self):\n        return self.sd_pipeline.unet\n\n    @property\n    def vae(self):\n        return self.sd_pipeline.vae\n\n    @property\n    def tokenizer(self):\n        return self.sd_pipeline.tokenizer\n\n    @property\n    def scheduler(self):\n        return self.sd_pipeline.scheduler\n\n    @property\n    def text_encoder(self):\n        return self.sd_pipeline.text_encoder\n\n    @property\n    def autocast(self):\n        return contextlib.nullcontext if self.use_lora else None\n\n    def save_pretrained(self, output_dir):\n        if self.use_lora:\n            state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(self.sd_pipeline.unet))\n            self.sd_pipeline.save_lora_weights(save_directory=output_dir, unet_lora_layers=state_dict)\n        self.sd_pipeline.save_pretrained(output_dir)\n\n    def set_progress_bar_config(self, *args, **kwargs):\n        self.sd_pipeline.set_progress_bar_config(*args, **kwargs)\n\n    def get_trainable_layers(self):\n        if self.use_lora:\n            lora_config = LoraConfig(\n                r=4,\n                lora_alpha=4,\n                init_lora_weights=\"gaussian\",\n                target_modules=[\"to_k\", \"to_q\", \"to_v\", \"to_out.0\"],\n            )\n            self.sd_pipeline.unet.add_adapter(lora_config)\n\n            # To avoid accelerate unscaling problems in FP16.\n            for param in self.sd_pipeline.unet.parameters():\n                # only upcast trainable parameters (LoRA) into fp32\n                if param.requires_grad:\n                    param.data = param.to(torch.float32)\n            return self.sd_pipeline.unet\n        else:\n            return self.sd_pipeline.unet\n\n    def save_checkpoint(self, models, weights, output_dir):\n        if len(models) != 1:\n            raise ValueError(\"Given how the trainable params were set, this should be of length 1\")\n        if self.use_lora and hasattr(models[0], \"peft_config\") and getattr(models[0], \"peft_config\", None) is not None:\n            state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(models[0]))\n            self.sd_pipeline.save_lora_weights(save_directory=output_dir, unet_lora_layers=state_dict)\n        elif not self.use_lora and isinstance(models[0], UNet2DConditionModel):\n            models[0].save_pretrained(os.path.join(output_dir, \"unet\"))\n        else:\n            raise ValueError(f\"Unknown model type {type(models[0])}\")\n\n    def load_checkpoint(self, models, input_dir):\n        if len(models) != 1:\n            raise ValueError(\"Given how the trainable params were set, this should be of length 1\")\n        if self.use_lora:\n            lora_state_dict, network_alphas = self.sd_pipeline.lora_state_dict(\n                input_dir, weight_name=\"pytorch_lora_weights.safetensors\"\n            )\n            self.sd_pipeline.load_lora_into_unet(lora_state_dict, network_alphas=network_alphas, unet=models[0])\n\n        elif not self.use_lora and isinstance(models[0], UNet2DConditionModel):\n            load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder=\"unet\")\n            models[0].register_to_config(**load_model.config)\n            models[0].load_state_dict(load_model.state_dict())\n            del load_model\n        else:\n            raise ValueError(f\"Unknown model type {type(models[0])}\")\n\n\n# This file is a copy of trl/examples/scripts/sft.py so that we could\n# use it together with rich and the TRL CLI in a more customizable manner.\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\nimport sys\nfrom argparse import Namespace\nfrom dataclasses import dataclass, field\n\nimport yaml\nfrom transformers import HfArgumentParser\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass YamlConfigParser:\n    def parse_and_set_env(self, config_path):\n        with open(config_path) as yaml_file:\n            config = yaml.safe_load(yaml_file)\n\n        if \"env\" in config:\n            env_vars = config.pop(\"env\")\n            if isinstance(env_vars, dict):\n                for key, value in env_vars.items():\n                    os.environ[key] = str(value)\n            else:\n                raise ValueError(\"`env` field should be a dict in the YAML file.\")\n\n        return config\n\n    def to_string(self, config):\n        final_string = \"\"\n        for key, value in config.items():\n            if isinstance(value, (dict, list)):\n                if len(value) != 0:\n                    value = str(value)\n                    value = value.replace(\"'\", '\"')\n                    value = f\"'{value}'\"\n                else:\n                    continue\n\n            final_string += f\"--{key} {value} \"\n        return final_string\n\n\ndef init_zero_verbose():\n    \"\"\"\n    Perform zero verbose init - use this method on top of the CLI modules to make\n    \"\"\"\n    import logging\n    import warnings\n\n    from rich.logging import RichHandler\n\n    FORMAT = \"%(message)s\"\n    logging.basicConfig(format=FORMAT, datefmt=\"[%X]\", handlers=[RichHandler()], level=logging.ERROR)\n\n    # Custom warning handler to redirect warnings to the logging system\n    def warning_handler(message, category, filename, lineno, file=None, line=None):\n        logging.warning(f\"{filename}:{lineno}: {category.__name__}: {message}\")\n\n    # Add the custom warning handler - we need to do that before importing anything to make sure the loggers work well\n    warnings.showwarning = warning_handler\n\n\n@dataclass\nclass SFTScriptArguments:\n    dataset_name: str = field(\n        default=\"timdettmers/openassistant-guanaco\",\n        metadata={\"help\": \"the dataset name\"},\n    )\n    dataset_train_split: str = field(default=\"train\", metadata={\"help\": \"The dataset split to train on\"})\n    dataset_test_split: str = field(default=\"test\", metadata={\"help\": \"The dataset split to evaluate on\"})\n    config: str = field(default=None, metadata={\"help\": \"Path to the optional config file\"})\n    gradient_checkpointing_use_reentrant: bool = field(\n        default=False,\n        metadata={\"help\": \"Whether to apply `use_reentrant` for gradient_checkpointing\"},\n    )\n\n\n@dataclass\nclass RewardScriptArguments:\n    dataset_name: str = field(\n        default=\"trl-lib/ultrafeedback_binarized\",\n        metadata={\"help\": \"the dataset name\"},\n    )\n    dataset_train_split: str = field(default=\"train\", metadata={\"help\": \"The dataset split to train on\"})\n    dataset_test_split: str = field(default=\"test\", metadata={\"help\": \"The dataset split to evaluate on\"})\n    config: str = field(default=None, metadata={\"help\": \"Path to the optional config file\"})\n    gradient_checkpointing_use_reentrant: bool = field(\n        default=False,\n        metadata={\"help\": \"Whether to apply `use_reentrant` for gradient_checkpointing\"},\n    )\n\n\n@dataclass\nclass DPOScriptArguments:\n    dataset_name: str = field(default=None, metadata={\"help\": \"the dataset name\"})\n    dataset_train_split: str = field(default=\"train\", metadata={\"help\": \"The dataset split to use for training\"})\n    dataset_test_split: str = field(default=\"test\", metadata={\"help\": \"The dataset split to use for evaluation\"})\n    ignore_bias_buffers: bool = field(\n        default=False,\n        metadata={\n            \"help\": \"debug argument for distributed training;\"\n            \"fix for DDP issues with LM bias/mask buffers - invalid scalar type,`inplace operation. See\"\n            \"https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992\"\n        },\n    )\n    config: str = field(default=None, metadata={\"help\": \"Path to the optional config file\"})\n    gradient_checkpointing_use_reentrant: bool = field(\n        default=False,\n        metadata={\"help\": \"Whether to apply `use_reentrant` for gradient_checkpointing\"},\n    )\n\n\n@dataclass\nclass ChatArguments:\n    # general settings\n    model_name_or_path: str = field(metadata={\"help\": \"Name of the pre-trained model\"})\n    user: str = field(default=None, metadata={\"help\": \"Username to display in chat interface\"})\n    system_prompt: str = field(default=None, metadata={\"help\": \"System prompt\"})\n    save_folder: str = field(default=\"./chat_history/\", metadata={\"help\": \"Folder to save chat history\"})\n    device: str = field(\n        default=\"cpu\",\n        metadata={\"help\": \"device to use for inference.\"},\n    )\n    config: str = field(\n        default=\"default\",\n        metadata={\n            \"help\": \"Config file used for setting the configs. If `default` uses examples/scripts/config/default_chat_config.yaml\"\n        },\n    )\n    examples: str = field(default=None, metadata={\"help\": \"Empty placeholder needs to be set via config.\"})\n    # generation settings\n    max_new_tokens: int = field(default=256, metadata={\"help\": \"Maximum number of tokens to generate\"})\n    do_sample: bool = field(default=True, metadata={\"help\": \"Whether to sample outputs during generation\"})\n    num_beams: int = field(default=1, metadata={\"help\": \"Number of beams for beam search\"})\n    temperature: float = field(default=1.0, metadata={\"help\": \"Temperature parameter for generation\"})\n    top_k: int = field(default=50, metadata={\"help\": \"Value of k for top-k sampling\"})\n    top_p: float = field(default=1.0, metadata={\"help\": \"Value of p for nucleus sampling\"})\n    repetition_penalty: float = field(default=1.0, metadata={\"help\": \"Repetition penalty\"})\n    eos_tokens: str = field(\n        default=None,\n        metadata={\"help\": \"EOS tokens to stop the generation. If multiple they should be comma separated\"},\n    )\n    eos_token_ids: str = field(\n        default=None,\n        metadata={\"help\": \"EOS token IDs to stop the generation. If multiple they should be comma separated\"},\n    )\n    # model loading\n    model_revision: str = field(\n        default=\"main\",\n        metadata={\"help\": \"The specific model version to use (can be a branch name, tag name or commit id).\"},\n    )\n    torch_dtype: str = field(\n        default=None,\n        metadata={\n            \"help\": (\n                \"Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the \"\n                \"dtype will be automatically derived from the model's weights.\"\n            ),\n            \"choices\": [\"auto\", \"bfloat16\", \"float16\", \"float32\"],\n        },\n    )\n    trust_remote_code: bool = field(default=False, metadata={\"help\": \"Trust remote code when loading a model.\"})\n    attn_implementation: str = field(\n        default=None,\n        metadata={\n            \"help\": (\n                \"Which attention implementation to use; you can run --attn_implementation=flash_attention_2, in which case you must install this manually by running `pip install flash-attn --no-build-isolation`\"\n            )\n        },\n    )\n    load_in_8bit: bool = field(\n        default=False,\n        metadata={\"help\": \"use 8 bit precision for the base model - works only with LoRA\"},\n    )\n    load_in_4bit: bool = field(\n        default=False,\n        metadata={\"help\": \"use 4 bit precision for the base model - works only with LoRA\"},\n    )\n\n    bnb_4bit_quant_type: str = field(default=\"nf4\", metadata={\"help\": \"precise the quantization type (fp4 or nf4)\"})\n    use_bnb_nested_quant: bool = field(default=False, metadata={\"help\": \"use nested quantization\"})\n\n\nclass TrlParser(HfArgumentParser):\n    def __init__(self, parsers, ignore_extra_args=False):\n        \"\"\"\n        The TRL parser parses a list of parsers (TrainingArguments, trl.ModelConfig, etc.), creates a config\n        parsers for users that pass a valid `config` field and merge the values that are set in the config\n        with the processed parsers.\n\n        Args:\n            parsers (`List[argparse.ArgumentParser`]):\n                List of parsers.\n            ignore_extra_args (`bool`):\n                Whether to ignore extra arguments passed by the config\n                and not raise errors.\n        \"\"\"\n        super().__init__(parsers)\n        self.yaml_parser = YamlConfigParser()\n        self.ignore_extra_args = ignore_extra_args\n\n    def post_process_dataclasses(self, dataclasses):\n        # Apply additional post-processing in case some arguments needs a special\n        # care\n        training_args = trl_args = None\n        training_args_index = None\n\n        for i, dataclass_obj in enumerate(dataclasses):\n            if dataclass_obj.__class__.__name__ == \"TrainingArguments\":\n                training_args = dataclass_obj\n                training_args_index = i\n            elif dataclass_obj.__class__.__name__ in (\"SFTScriptArguments\", \"DPOScriptArguments\"):\n                trl_args = dataclass_obj\n            else:\n                ...\n\n        if trl_args is not None and training_args is not None:\n            training_args.gradient_checkpointing_kwargs = dict(\n                use_reentrant=trl_args.gradient_checkpointing_use_reentrant\n            )\n            dataclasses[training_args_index] = training_args\n\n        return dataclasses\n\n    def parse_args_and_config(self, return_remaining_strings=False):\n        yaml_config = None\n        if \"--config\" in sys.argv:\n            config_index = sys.argv.index(\"--config\")\n\n            _ = sys.argv.pop(config_index)  # --config\n            config_path = sys.argv.pop(config_index)  # path to config\n            yaml_config = self.yaml_parser.parse_and_set_env(config_path)\n\n            self.set_defaults_with_config(**yaml_config)\n\n        outputs = self.parse_args_into_dataclasses(return_remaining_strings=return_remaining_strings)\n\n        if yaml_config is None:\n            return outputs\n\n        if return_remaining_strings:\n            # if we have extra yaml config and command line strings\n            # outputs[-1] is remaining command line strings\n            # outputs[-2] is remaining yaml config as Namespace\n            # combine them into remaining strings object\n            remaining_strings = outputs[-1] + [f\"{key}: {value}\" for key, value in vars(outputs[-2]).items()]\n            return outputs[:-2], remaining_strings\n        else:\n            # outputs[-1] is either remaining yaml config as Namespace or parsed config as Dataclass\n            if isinstance(outputs[-1], Namespace) and not self.ignore_extra_args:\n                remaining_args = vars(outputs[-1])\n                raise ValueError(f\"Some specified config arguments are not used by the TrlParser: {remaining_args}\")\n\n            return outputs\n\n    def set_defaults_with_config(self, **kwargs):\n        \"\"\"Defaults we're setting with config allow us to change to required = False\"\"\"\n        self._defaults.update(kwargs)\n\n        # if these defaults match any existing arguments, replace\n        # the previous default on the object with the new one\n        for action in self._actions:\n            if action.dest in kwargs:\n                action.default = kwargs[action.dest]\n                action.required = False\n\n\n# This file is a copy of trl/examples/scripts/sft.py so that we could\n# use it together with rich and the TRL CLI in a more customizable manner.\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport subprocess\nimport sys\nfrom subprocess import CalledProcessError\n\nfrom rich.console import Console\n\n\nSUPPORTED_COMMANDS = [\"sft\", \"dpo\", \"chat\", \"kto\"]\n\n\ndef main():\n    console = Console()\n    # Make sure to import things locally to avoid verbose from third party libs.\n    with console.status(\"[bold purple]Welcome! Initializing the TRL CLI...\"):\n        from trl.commands.cli_utils import init_zero_verbose\n\n        init_zero_verbose()\n\n        command_name = sys.argv[1]\n\n        if command_name not in SUPPORTED_COMMANDS:\n            raise ValueError(\n                f\"Please use one of the supported commands, got {command_name} - supported commands are {SUPPORTED_COMMANDS}\"\n            )\n\n        trl_examples_dir = os.path.dirname(__file__)\n\n    if command_name == \"chat\":\n        command = f\"\"\"\n        python {trl_examples_dir}/scripts/{command_name}.py {\" \".join(sys.argv[2:])}\n        \"\"\"\n    else:\n        command = f\"\"\"\n        accelerate launch {trl_examples_dir}/scripts/{command_name}.py {\" \".join(sys.argv[2:])}\n        \"\"\"\n\n    try:\n        subprocess.run(\n            command.split(),\n            text=True,\n            check=True,\n            encoding=\"utf-8\",\n            cwd=os.getcwd(),\n            env=os.environ.copy(),\n        )\n    except (CalledProcessError, ChildProcessError) as exc:\n        console.log(f\"TRL - {command_name.upper()} failed on ! See the logs above for further details.\")\n        raise ValueError(\"TRL CLI failed! Check the traceback above..\") from exc\n\n\nif __name__ == \"__main__\":\n    main()\n\n\n# flake8: noqa\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# flake8: noqa\n\nfrom typing import TYPE_CHECKING\nfrom ..import_utils import _LazyModule, OptionalDependencyNotAvailable\n\n\n_import_structure = {\n    \"cli_utils\": [\"SFTScriptArguments\", \"init_zero_verbose\", \"DPOScriptArguments\", \"TrlParser\", \"YamlConfigParser\"],\n}\n\nif TYPE_CHECKING:\n    from .cli_utils import SFTScriptArguments, init_zero_verbose, DPOScriptArguments, TrlParser, YamlConfigParser\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(__name__, globals()[\"__file__\"], _import_structure, module_spec=__spec__)\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport warnings\nfrom functools import wraps\nfrom typing import Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    DataCollator,\n    DataCollatorForLanguageModeling,\n    DataCollatorForSeq2Seq,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainingArguments,\n)\nfrom transformers.trainer_utils import EvalLoopOutput\nfrom transformers.utils import is_peft_available\n\nfrom ..core import PPODecorators\nfrom .utils import trl_sanitze_kwargs_for_tagging\n\n\nif is_peft_available():\n    from peft import PeftModel\n\n\nclass IterativeSFTTrainer(Trainer):\n    \"\"\"\n    The IterativeSFTTrainer can be used to finetune models with methods that requires some steps between optimization.\n\n    Args:\n        model (`PreTrainedModel`):\n            Model to be optimized, either an 'AutoModelForCausalLM' or an 'AutoModelForSeq2SeqLM'.\n            Check the documentation of `PreTrainedModel` for more details.\n        args (`transformers.TrainingArguments`):\n            The arguments to use for training.\n        tokenizer (`PreTrainedTokenizerBase`):\n            Tokenizer to be used for encoding the data. Check the documentation of `transformers.PreTrainedTokenizer` and\n            `transformers.PreTrainedTokenizerFast` for more details.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        data_collator (Union[DataCollatorForLanguageModeling, DataCollatorForSeq2Seq], *optional*):\n            Data collator to be used for training and passed along the dataloader.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        max_length (`int`, defaults to `None`):\n            The maximum length of the input.\n        truncation_mode (`str`, defaults to `keep_end`):\n            The truncation mode to use, either `keep_end` or `keep_start`.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to metric values.\n        optimize_device_cache (`bool`, *optional*, defaults to `False`):\n            Optimize CUDA cache for slightly more memory-efficient training.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"iterative-sft\"]\n\n    def __init__(\n        self,\n        model: Optional[PreTrainedModel] = None,\n        args: Optional[TrainingArguments] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (\n            None,\n            None,\n        ),\n        data_collator: Optional[DataCollator] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        max_length: Optional[int] = None,\n        truncation_mode: Optional[str] = \"keep_end\",\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n        optimize_device_cache: Optional[bool] = False,\n    ):\n        # Step 0: check positional arguments validity\n        if not isinstance(tokenizer, (PreTrainedTokenizerBase)):\n            raise ValueError(\n                f\"tokenizer must be a PreTrainedTokenizerBase like a PreTrainedTokenizer or a PreTrainedTokenizerFast, got {type(tokenizer)}\"\n            )\n        if not isinstance(model, PreTrainedModel):\n            raise ValueError(f\"model must be a PreTrainedModel, got {type(model)}\")\n        if not model.can_generate():\n            warnings.warn(\n                f\"The current model class {type(model)} is not compatible with `.generate()`\"\n                \"Please make sure that this is intended.\"\n            )\n        if optimizers[1] is None and args.max_steps == -1:\n            raise ValueError(\n                \"When no scheduler is provided, you need to set the total number of training steps to perform `max_steps`\"\n            )\n\n        self.is_encoder_decoder = getattr(model.config, \"is_encoder_decoder\", False)\n        self.is_peft_model = is_peft_available() and isinstance(model, PeftModel)\n\n        self.tokenizer = tokenizer\n\n        if data_collator is None:\n            if self.is_encoder_decoder:\n                warnings.warn(\n                    \"No data collator is provided. Using 'DataCollatorForSeq2Seq' with\"\n                    \"'labels_pad_token_id' set to '-100' and 'pad_to_multiple_of' set to 8.\"\n                )\n                self.data_collator = DataCollatorForSeq2Seq(tokenizer, label_pad_token_id=-100, pad_to_multiple_of=8)\n            else:\n                warnings.warn(\"No data collator is provided. Using 'DataCollatorForLanguageModeling'\")\n                self.data_collator = DataCollatorForLanguageModeling(self.tokenizer, mlm=False)\n        else:\n            self.data_collator = data_collator\n\n        self.max_length = max_length\n        self.truncation_mode = truncation_mode\n        self.optimize_device_cache = optimize_device_cache\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=self.data_collator,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            compute_metrics=compute_metrics,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        self.create_optimizer_and_scheduler(self.args.max_steps)\n\n        # prepare model, optimizer and lr_scheduler\n        self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(\n            self.model, self.optimizer, self.lr_scheduler\n        )\n\n        self.tokenizer.truncation_side = \"left\" if self.truncation_mode == \"keep_end\" else \"right\"\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n        PPODecorators.optimize_device_cache = self.optimize_device_cache\n\n    def prepare_model_inputs(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor):\n        if attention_mask is None:\n            attention_mask = [torch.ones_like(ids) for ids in input_ids]\n\n        if self.is_encoder_decoder:\n            input_data = self.data_collator(\n                [\n                    {\"input_ids\": ids, \"attention_mask\": att, \"labels\": lab}\n                    for ids, att, lab in zip(input_ids, attention_mask, labels)\n                ]\n            ).to(self.model.device)\n\n            input_data.pop(\"decoder_input_ids\", None)  # This is directly computed inside the model\n\n            input_data[\"labels\"][input_data[\"labels\"] == self.tokenizer.pad_token_id] = -100\n\n        else:\n            input_data = self.data_collator(\n                [{\"input_ids\": ids, \"attention_mask\": att} for ids, att in zip(input_ids, attention_mask)]\n            ).to(self.model.device)\n\n        # truncate in case the user has provided input_ids, attention_mask and labels\n        if self.max_length is not None:\n            if self.truncation_mode == \"keep_start\":\n                input_data = {k: v[: self.max_length] for k, v in input_data.items()}\n            elif self.truncation_mode == \"keep_end\":\n                input_data = {k: v[-self.max_length :] for k, v in input_data.items()}\n            else:\n                raise ValueError(f\"Unknown truncation mode: {self.truncation_mode}\")\n\n        return input_data\n\n    @staticmethod\n    def _step_safety_checker(\n        input_ids: List[torch.LongTensor],\n        attention_mask: List[torch.LongTensor],\n        labels: List[torch.LongTensor],\n        texts: List[str],\n        texts_labels: List[str],\n    ):\n        \"\"\"\n        Check if the input data is valid for training.\n\n        Args:\n            input_ids (List[`torch.LongTensor`]):\n                List of tensors containing the input_ids\n            attention_mask (List[`torch.LongTensor`]):\n                List of tensors containing the attention_mask\n            labels (List[`torch.FloatTensor`]):\n                List of tensors containing the labels\n            texts (List[`str`]):\n                List of string containing the text input.\n            texts_labels (List[`str`]):\n                List of string containing the text labels.\n\n        Returns:\n            `tuple`: The input data.\n        \"\"\"\n        if texts is None:\n            if attention_mask is None:\n                for name, tensor_list in zip([\"input_ids\", \"labels\"], [input_ids, labels]):\n                    if not isinstance(tensor_list, list):\n                        raise ValueError(f\"{name} must be a list of tensors - got {type(tensor_list)}\")\n                    if not isinstance(tensor_list[0], torch.Tensor):\n                        raise ValueError(f\"Elements in {name} must be tensors - got {type(tensor_list[0])}\")\n            else:\n                for name, tensor_list in zip(\n                    [\"input_ids\", \"attention_mask\", \"labels\"], [input_ids, attention_mask, labels]\n                ):\n                    if not isinstance(tensor_list, list):\n                        raise ValueError(f\"{name} must be a list of tensors - got {type(tensor_list)}\")\n                    if not isinstance(tensor_list[0], torch.Tensor):\n                        raise ValueError(f\"Elements in {name} must be tensors - got {type(tensor_list[0])}\")\n        else:\n            if not isinstance(texts, list):\n                raise ValueError(f\"'text' must be a list of strings - got {type(texts)}\")\n            if not isinstance(texts[0], str):\n                raise ValueError(f\"Elements in 'text' must be strings - got {type(texts[0])}\")\n            if texts_labels is not None:\n                if not isinstance(texts_labels, list):\n                    raise ValueError(f\"'text_labels' must be a list of strings - got {type(texts_labels)}\")\n                if not isinstance(texts_labels[0], str):\n                    raise ValueError(f\"Elements in 'text_labels' must be strings - got {type(texts_labels[0])}\")\n\n        return input_ids, attention_mask, labels, texts, texts_labels\n\n    @PPODecorators.empty_device_cache()\n    def step(\n        self,\n        input_ids: Optional[List[torch.LongTensor]] = None,\n        attention_mask: Optional[List[torch.LongTensor]] = None,\n        labels: Optional[List[torch.LongTensor]] = None,\n        texts: Optional[List[str]] = None,\n        texts_labels: Optional[List[str]] = None,\n    ):\n        \"\"\"\n        Run an optimisation step given a list of input_ids, attention_mask, and labels or a list of text and text_labels.\n        Args:\n            input_ids (List[`torch.LongTensor`]):\n                List of tensors containing the input_ids (if not provided, text will be used)\n            attention_mask (List[`torch.LongTensor`], , *optional*):\n                List of tensors containing the attention_mask\n            labels (List[`torch.FloatTensor`], *optional*):\n                List of tensors containing the labels (if set to None, will default to input_ids)\n            texts (List[`str`], *optional*):\n                List of strings containing the text input (if not provided, input_ids will directly be used)\n            texts_labels (List[`str`], *optional*):\n                List of strings containing the text labels (if set to None, will default to text)\n\n        Returns:\n            `dict[str, Any]`: A summary of the training statistics\n        \"\"\"\n        self.model.train()\n\n        if self.state.global_step == 0:\n            self.tr_loss = torch.tensor(0.0).to(self.args.device)\n            self._globalstep_last_logged = self.state.global_step\n\n        if input_ids is None and texts is None:\n            raise ValueError(\"Step should include `input_ids` or `texts` as keyword arguments.\")\n        elif input_ids is not None and texts is not None:\n            warnings.warn(\n                \"Both 'input_ids' and 'texts' are provided. 'input_ids' will be overwritten using inputs provided by the 'texts' keyword argument.\"\n            )\n\n        if labels is None and texts_labels is None and self.is_encoder_decoder:\n            raise ValueError(\n                \"No 'labels' or 'text_labels' are provided. When using an encoder-decoder architecture, 'labels' or 'text_labels' must be passed.\"\n            )\n\n        input_ids, attention_mask, labels, texts, texts_labels = self._step_safety_checker(\n            input_ids, attention_mask, labels, texts, texts_labels\n        )\n\n        if texts is not None:\n            model_inputs = self.tokenizer(\n                texts, max_length=self.max_length, truncation=True, padding=True, return_tensors=\"pt\"\n            )\n\n            input_ids, attention_mask = model_inputs[\"input_ids\"], model_inputs[\"attention_mask\"]\n\n        if texts_labels is not None:\n            labels = self.tokenizer(\n                texts, max_length=self.max_length, truncation=True, padding=True, return_tensors=\"pt\"\n            )[\"input_ids\"]\n\n        if labels is None:\n            warnings.warn(\"No labels are provided. Setting labels to input_ids\")\n            labels = input_ids\n\n        model_inputs = self.prepare_model_inputs(input_ids, attention_mask, labels)\n\n        model_inputs_names = list(model_inputs.keys())\n\n        batch_dict = {}\n        batch_dict.update(model_inputs)\n\n        def collator(data):\n            return_dict = dict()\n            for key in data[0]:\n                if key in [\"input_ids\", \"attention_mask\", \"labels\"]:\n                    return_dict[key] = torch.stack([d[key] for d in data]).to(self.model.device)\n            return return_dict\n\n        batch_data = Dataset.from_dict(batch_dict)\n        batch_data.set_format(\"torch\")\n\n        step_dataloader = DataLoader(\n            batch_data,\n            batch_size=self.args.per_device_train_batch_size,\n            shuffle=True,\n            collate_fn=collator,\n        )\n\n        for _, batch in enumerate(step_dataloader):\n            with self.accelerator.accumulate(self.model):\n                model_inputs = {k: batch[k] for k in model_inputs_names}\n                loss = self.compute_loss(self.model, model_inputs)\n\n                if self.args.n_gpu > 1:\n                    loss = loss.mean()\n\n                tr_loss_step = loss.detach()\n\n                self.accelerator.backward(loss)\n\n                if self.accelerator.sync_gradients and self.args.max_grad_norm is not None:\n                    self.accelerator.clip_grad_norm_(\n                        self.model.parameters(),\n                        self.args.max_grad_norm,\n                    )\n\n                self.optimizer.step()\n                self.optimizer.zero_grad()\n                if self.lr_scheduler is not None:\n                    self.lr_scheduler.step()\n\n                self.state.global_step += 1\n\n                # update stats etc\n                self.tr_loss += tr_loss_step\n\n                self._maybe_log_save_evaluate()\n\n    def _maybe_log_save_evaluate(self):\n        # check if eval is required\n        if self.args.eval_steps is not None:\n            if self.state.global_step % self.args.eval_steps == 0 and self.state.global_step != 0:\n                self.evaluate(self.eval_dataset)\n\n        # check if logging is required\n        if self.args.logging_steps is not None:\n            if self.state.global_step % self.args.logging_steps == 0 and self.state.global_step != 0:\n                logs: Dict[str, float] = {}\n\n                tr_loss_scalar = self._nested_gather(self.tr_loss).mean().item()\n\n                # reset tr_loss to zero\n                self.tr_loss -= self.tr_loss\n\n                logs[\"loss\"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)\n                logs[\"learning_rate\"] = self._get_learning_rate()\n\n                self._globalstep_last_logged = self.state.global_step\n\n                self.log(logs)\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"iterative-sft\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import List, Union\n\nfrom trl.trainer.online_dpo_config import OnlineDPOConfig\n\n\n@dataclass\nclass XPOConfig(OnlineDPOConfig):\n    r\"\"\"\n    Configuration class for the [`XPOTrainer`].\n\n    Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following:\n\n    Parameters:\n        alpha (`float` or `List[float]`, *optional*, defaults to `1e-5`):\n            Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each new epoch and the last alpha is used for the rest of the epochs.\n    \"\"\"\n\n    alpha: Union[float, List[float]] = 1e-5\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass ORPOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`ORPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        max_prompt_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the prompt. This argument is required if you want to use the default data collator.\n        max_completion_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the completion. This argument is required if you want to use the default data collator\n            and your model is an encoder-decoder.\n        beta (`float`, *optional*, defaults to `0.1`):\n            Parameter controlling the relative ratio loss weight in the ORPO loss. In the [paper](https://huggingface.co/papers/2403.07691),\n            it is denoted by λ. In the [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`.\n        disable_dropout (`bool`, *optional*, defaults to `True`):\n            Whether to disable dropout in the model.\n        label_pad_token_id (`int`, *optional*, defaults to `-100`):\n            Label pad token id. This argument is required if you want to use the default data collator.\n        padding_value (`Optional[int]`, *optional*, defaults to `None`):\n            Padding value to use. If `None`, the padding value of the tokenizer is used.\n        truncation_mode (`str`, *optional*, defaults to `\"keep_end\"`):\n            Truncation mode to use when the prompt is too long. Possible values are `\"keep_end\"` or `\"keep_start\"`.\n            This argument is required if you want to use the default data collator.\n        generate_during_eval (`bool`, *optional*, defaults to `False`):\n            If `True`, generates and logs completions from the model to W&B during evaluation.\n        is_encoder_decoder (`Optional[bool]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n    \"\"\"\n\n    max_length: Optional[int] = None\n    max_prompt_length: Optional[int] = None\n    max_completion_length: Optional[int] = None\n    beta: float = 0.1\n    disable_dropout: bool = True\n    label_pad_token_id: int = -100\n    padding_value: Optional[int] = None\n    truncation_mode: str = \"keep_end\"\n    generate_during_eval: bool = False\n    is_encoder_decoder: Optional[bool] = None\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    dataset_num_proc: Optional[int] = None\n\n\n# ORPO Authors: Jiwoo Hong, Noah Lee, and James Thorne\n# Official code: https://github.com/xfactlab/orpo\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import nullcontext\nfrom copy import deepcopy\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.amp as amp\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import PartialState\nfrom accelerate.utils import is_deepspeed_available\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    AutoModelForCausalLM,\n    DataCollator,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    is_torch_xla_available,\n    is_wandb_available,\n)\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalLoopOutput\nfrom transformers.utils import is_peft_available, is_torch_fx_proxy\n\nfrom ..models import PreTrainedModelWrapper\nfrom .orpo_config import ORPOConfig\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    add_bos_token_if_needed,\n    add_eos_token_if_needed,\n    disable_dropout_in_model,\n    pad_to_length,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\n\nif is_wandb_available():\n    import wandb\n\nif is_deepspeed_available():\n    import deepspeed\n\nif is_torch_xla_available():\n    import torch_xla.core.xla_model as xm\n\n\nclass ORPOTrainer(Trainer):\n    r\"\"\"\n    Initialize ORPOTrainer.\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForSequenceClassification`.\n        args (`ORPOConfig`):\n            The ORPO config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Dict`, defaults to `None`):\n            The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"orpo\"]\n\n    def __init__(\n        self,\n        model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        args: Optional[ORPOConfig] = None,\n        data_collator: Optional[DataCollator] = None,\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n    ):\n        if args.model_init_kwargs is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\"You passed model_kwargs to the ORPOTrainer. But your model is already instantiated.\")\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the ORPOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`\n        # has been called in order to properly call autocast if needed.\n        self._peft_has_been_casted_to_bf16 = False\n\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            # if model is a peft model and we have a peft_config, we merge and unload it first\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_loaded_in_4bit\", False):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n\n                prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                if _support_gc_kwargs:\n                    prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n            elif getattr(args, \"gradient_checkpointing\", False):\n                # For backward compatibility with older versions of transformers\n                if hasattr(model, \"enable_input_require_grads\"):\n                    model.enable_input_require_grads()\n                else:\n\n                    def make_inputs_require_grad(module, input, output):\n                        output.requires_grad_(True)\n\n                    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            # get peft model with the given config\n            model = get_peft_model(model, peft_config)\n            if args.bf16 and getattr(model, \"is_loaded_in_4bit\", False):\n                peft_module_casting_to_bf16(model)\n                # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager\n                self._peft_has_been_casted_to_bf16 = True\n\n        # For models that use gradient_checkpointing, we need to attach a hook that enables input\n        # to explicitly have `requires_grad=True`, otherwise training will either silently\n        # fail or completely fail.\n        elif getattr(args, \"gradient_checkpointing\", False):\n            # For backward compatibility with older versions of transformers\n            if hasattr(model, \"enable_input_require_grads\"):\n                model.enable_input_require_grads()\n            else:\n\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        if args.generate_during_eval and not is_wandb_available():\n            raise ValueError(\n                \"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install `wandb` to resolve.\"\n            )\n\n        if model is not None:\n            self.is_encoder_decoder = model.config.is_encoder_decoder\n        elif args.is_encoder_decoder is None:\n            raise ValueError(\"When no model is provided, you need to pass the parameter is_encoder_decoder.\")\n        else:\n            self.is_encoder_decoder = args.is_encoder_decoder\n\n        if self.is_encoder_decoder:\n            self.decoder_start_token_id = model.config.decoder_start_token_id\n            self.pad_token_id = model.config.pad_token_id\n\n        if tokenizer is None:\n            raise ValueError(\"tokenizer must be specified to tokenize a ORPO dataset.\")\n        if args.max_length is None:\n            warnings.warn(\n                \"`max_length` is not set in the ORPOConfig's init\"\n                \" it will default to `512` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_length = 512\n        else:\n            max_length = args.max_length\n        if args.max_prompt_length is None:\n            warnings.warn(\n                \"`max_prompt_length` is not set in the ORPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_prompt_length = 128\n        else:\n            max_prompt_length = args.max_prompt_length\n\n        if args.max_completion_length is None and self.is_encoder_decoder:\n            warnings.warn(\n                \"When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            self.max_completion_length = 128\n        else:\n            self.max_completion_length = args.max_completion_length\n\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(\n                pad_token_id=tokenizer.pad_token_id,\n                label_pad_token_id=args.label_pad_token_id,\n                is_encoder_decoder=self.is_encoder_decoder,\n            )\n\n            if args.remove_unused_columns:\n                args.remove_unused_columns = False\n                # warn users\n                warnings.warn(\n                    \"When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_dpo_data_collator = True\n        else:\n            self.use_dpo_data_collator = False\n\n        if args.disable_dropout:\n            disable_dropout_in_model(model)\n\n        self.max_length = max_length\n        self.generate_during_eval = args.generate_during_eval\n        self.label_pad_token_id = args.label_pad_token_id\n        self.padding_value = args.padding_value if args.padding_value is not None else tokenizer.pad_token_id\n        self.max_prompt_length = max_prompt_length\n        self.truncation_mode = args.truncation_mode\n        self.tokenizer = tokenizer\n\n        self.beta = args.beta\n        self.aux_loss_enabled = getattr(model.config, \"output_router_logits\", False)\n\n        self._stored_metrics = defaultdict(lambda: defaultdict(list))\n\n        # Compute that only on the main process for faster data processing.\n        # see: https://github.com/huggingface/trl/pull/1255\n        with PartialState().local_main_process_first():\n            # tokenize the dataset\n            train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)\n\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n    def build_tokenized_answer(self, prompt, answer):\n        \"\"\"\n        Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`.\n        It does ensure `enc(a + b) = enc(a) + enc(a + b)[len(enc(a)):]`.\n        Reference:\n            https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n        \"\"\"\n\n        full_tokenized = self.tokenizer(prompt + answer, add_special_tokens=False)\n        prompt_input_ids = self.tokenizer(prompt, add_special_tokens=False)[\"input_ids\"]\n\n        answer_input_ids = full_tokenized[\"input_ids\"][len(prompt_input_ids) :]\n        answer_attention_mask = full_tokenized[\"attention_mask\"][len(prompt_input_ids) :]\n\n        # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`\n        full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids])\n\n        # Prepare input tokens for token by token comparison\n        full_input_ids = np.array(full_tokenized[\"input_ids\"])\n\n        if len(full_input_ids) != len(full_concat_input_ids):\n            raise ValueError(\"Prompt input ids and answer input ids should have the same length.\")\n\n        # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens\n        # can be merged together when tokenizing prompt+answer. This could result\n        # on the last token from the prompt being different when tokenized on its own\n        # vs when done as prompt+answer.\n        response_token_ids_start_idx = len(prompt_input_ids)\n\n        # If tokenized prompt is different than both prompt+answer, then it means the\n        # last token has changed due to merging.\n        if prompt_input_ids != full_tokenized[\"input_ids\"][:response_token_ids_start_idx]:\n            response_token_ids_start_idx -= 1\n\n        prompt_input_ids = full_tokenized[\"input_ids\"][:response_token_ids_start_idx]\n        prompt_attention_mask = full_tokenized[\"attention_mask\"][:response_token_ids_start_idx]\n\n        if len(prompt_input_ids) != len(prompt_attention_mask):\n            raise ValueError(\"Prompt input ids and attention mask should have the same length.\")\n\n        answer_input_ids = full_tokenized[\"input_ids\"][response_token_ids_start_idx:]\n        answer_attention_mask = full_tokenized[\"attention_mask\"][response_token_ids_start_idx:]\n\n        return dict(\n            prompt_input_ids=prompt_input_ids,\n            prompt_attention_mask=prompt_attention_mask,\n            input_ids=answer_input_ids,\n            attention_mask=answer_attention_mask,\n        )\n\n    def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> Dict:\n        \"\"\"Tokenize a single row from a ORPO specific dataset.\n\n        At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation\n        in case the prompt + chosen or prompt + rejected responses is/are too long. First\n        we truncate the prompt; if we're still too long, we truncate the chosen/rejected.\n\n        We also create the labels for the chosen/rejected responses, which are of length equal to\n        the sum of the length of the prompt and the chosen/rejected response, with\n        label_pad_token_id  for the prompt tokens.\n        \"\"\"\n        batch = {}\n        prompt = feature[\"prompt\"]\n        chosen = feature[\"chosen\"]\n        rejected = feature[\"rejected\"]\n\n        if not self.is_encoder_decoder:\n            # Check issues below for more details\n            #  1. https://github.com/huggingface/trl/issues/907\n            #  2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n            #  3. https://github.com/LianjiaTech/BELLE/issues/337\n\n            if not isinstance(prompt, str):\n                raise ValueError(f\"prompt should be an str but got {type(prompt)}\")\n            prompt_tokens = self.tokenizer(prompt, add_special_tokens=False)\n            prompt_tokens = {f\"prompt_{k}\": v for k, v in prompt_tokens.items()}\n\n            if not isinstance(chosen, str):\n                raise ValueError(f\"chosen should be an str but got {type(chosen)}\")\n            chosen_tokens = self.build_tokenized_answer(prompt, chosen)\n\n            if not isinstance(rejected, str):\n                raise ValueError(f\"rejected should be an str but got {type(rejected)}\")\n            rejected_tokens = self.build_tokenized_answer(prompt, rejected)\n\n            # Last prompt token might get merged by tokenizer and\n            # it should not be included for generation if that happens\n            prompt_len_input_ids = len(prompt_tokens[\"prompt_input_ids\"])\n\n            chosen_prompt_len_input_ids = len(chosen_tokens[\"prompt_input_ids\"])\n            rejected_prompt_len_input_ids = len(rejected_tokens[\"prompt_input_ids\"])\n            prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids)\n\n            for k, v in prompt_tokens.items():\n                prompt_tokens[k] = v[:prompt_len_input_ids]\n\n            # Make sure prompts only have one different token at most an\n            # and length only differs by 1 at most\n            num_diff_tokens = sum(\n                [a != b for a, b in zip(chosen_tokens[\"prompt_input_ids\"], rejected_tokens[\"prompt_input_ids\"])]\n            )\n            num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids)\n            if num_diff_tokens > 1 or num_diff_len > 1:\n                raise ValueError(\n                    \"Chosen and rejected prompt_input_ids might only differ on the \"\n                    \"last token due to tokenizer merge ops.\"\n                )\n\n            # add BOS token to head of prompt. Avoid adding if it's already there\n            prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed(\n                self.tokenizer.bos_token_id,\n                prompt_len_input_ids,\n                prompt_tokens,\n                chosen_prompt_len_input_ids,\n                chosen_tokens,\n                rejected_prompt_len_input_ids,\n                rejected_tokens,\n            )\n\n            # add EOS token to end of answer. Avoid adding if it's already there\n            chosen_tokens, rejected_tokens = add_eos_token_if_needed(\n                self.tokenizer.eos_token_id, chosen_tokens, rejected_tokens\n            )\n\n            longer_response_length = max(len(chosen_tokens[\"input_ids\"]), len(rejected_tokens[\"input_ids\"]))\n\n            # if combined sequence is too long, truncate the prompt\n            for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]:\n                if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > self.max_length:\n                    if self.truncation_mode == \"keep_start\":\n                        for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                            answer_tokens[k] = answer_tokens[k][: self.max_prompt_length]\n                    elif self.truncation_mode == \"keep_end\":\n                        for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                            answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :]\n                    else:\n                        raise ValueError(f\"Unknown truncation mode: {self.truncation_mode}\")\n\n            # if that's still too long, truncate the response\n            for answer_tokens in [chosen_tokens, rejected_tokens]:\n                if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > self.max_length:\n                    for k in [\"input_ids\", \"attention_mask\"]:\n                        answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length]\n\n            # Create labels\n            chosen_sequence_tokens = {\n                k: chosen_tokens[f\"prompt_{k}\"] + chosen_tokens[k] for k in [\"input_ids\", \"attention_mask\"]\n            }\n            rejected_sequence_tokens = {\n                k: rejected_tokens[f\"prompt_{k}\"] + rejected_tokens[k] for k in [\"input_ids\", \"attention_mask\"]\n            }\n            chosen_sequence_tokens[\"labels\"] = chosen_sequence_tokens[\"input_ids\"][:]\n            chosen_sequence_tokens[\"labels\"][: len(chosen_tokens[\"prompt_input_ids\"])] = [\n                self.label_pad_token_id\n            ] * len(chosen_tokens[\"prompt_input_ids\"])\n            rejected_sequence_tokens[\"labels\"] = rejected_sequence_tokens[\"input_ids\"][:]\n            rejected_sequence_tokens[\"labels\"][: len(rejected_tokens[\"prompt_input_ids\"])] = [\n                self.label_pad_token_id\n            ] * len(rejected_tokens[\"prompt_input_ids\"])\n\n            for k, toks in {\n                \"chosen_\": chosen_sequence_tokens,\n                \"rejected_\": rejected_sequence_tokens,\n                \"\": prompt_tokens,\n            }.items():\n                for type_key, tokens in toks.items():\n                    if type_key == \"token_type_ids\":\n                        continue\n                    batch[f\"{k}{type_key}\"] = tokens\n\n        else:\n            chosen_tokens = self.tokenizer(\n                chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True\n            )\n            rejected_tokens = self.tokenizer(\n                rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True\n            )\n            prompt_tokens = self.tokenizer(\n                prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True\n            )\n\n            batch[\"chosen_labels\"] = chosen_tokens[\"input_ids\"]\n            batch[\"rejected_labels\"] = rejected_tokens[\"input_ids\"]\n            batch[\"prompt_input_ids\"] = prompt_tokens[\"input_ids\"]\n            batch[\"prompt_attention_mask\"] = prompt_tokens[\"attention_mask\"]\n\n            if model is not None and hasattr(model, \"prepare_decoder_input_ids_from_labels\"):\n                batch[\"rejected_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                    labels=torch.tensor(batch[\"rejected_labels\"])\n                )\n                batch[\"chosen_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                    labels=torch.tensor(batch[\"chosen_labels\"])\n                )\n\n        if is_torch_xla_available():\n            # Pad the sequences to global max_length to avoid TorchXLA recompilation\n            for k in batch:\n                if \"labels\" in k or self.is_encoder_decoder:\n                    pad_value = self.label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = self.padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k]))\n        return batch\n\n    @staticmethod\n    def concatenated_inputs(\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        is_encoder_decoder: bool = False,\n        label_pad_token_id: int = -100,\n        padding_value: int = 0,\n        device: Optional[torch.device] = None,\n    ) -> Dict[str, torch.LongTensor]:\n        \"\"\"Concatenate the chosen and rejected inputs into a single tensor.\n\n        Args:\n            batch: A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors of shape (batch_size, sequence_length).\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n            label_pad_token_id: The label pad token id.\n            padding_value: The padding value to use for the concatenated inputs_ids.\n            device: The device for the concatenated inputs.\n\n        Returns:\n            A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'.\n        \"\"\"\n        concatenated_batch = {}\n\n        if is_encoder_decoder:\n            max_length = max(batch[\"chosen_labels\"].shape[1], batch[\"rejected_labels\"].shape[1])\n        else:\n            max_length = max(batch[\"chosen_input_ids\"].shape[1], batch[\"rejected_input_ids\"].shape[1])\n\n        for k in batch:\n            if k.startswith(\"chosen\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"chosen\", \"concatenated\")\n                concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value)\n        for k in batch:\n            if k.startswith(\"rejected\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"rejected\", \"concatenated\")\n                concatenated_batch[concatenated_key] = torch.cat(\n                    (\n                        concatenated_batch[concatenated_key],\n                        pad_to_length(batch[k], max_length, pad_value=pad_value),\n                    ),\n                    dim=0,\n                ).to(device=device)\n\n        if is_encoder_decoder:\n            concatenated_batch[\"concatenated_input_ids\"] = batch[\"prompt_input_ids\"].repeat(2, 1).to(device=device)\n            concatenated_batch[\"concatenated_attention_mask\"] = (\n                batch[\"prompt_attention_mask\"].repeat(2, 1).to(device=device)\n            )\n\n        return concatenated_batch\n\n    def odds_ratio_loss(\n        self,\n        policy_chosen_logps: torch.FloatTensor,\n        policy_rejected_logps: torch.FloatTensor,\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities.\n\n        Args:\n            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)\n            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)\n\n        Returns:\n            A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).\n            The losses tensor contains the ORPO loss for each example in the batch.\n            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.\n            The log odds ratio of the chosen responses over the rejected responses ratio for logging purposes.\n            The `log(sigmoid(log_odds_chosen))` for logging purposes.\n        \"\"\"\n\n        # Derived from Eqs. (4) and (7) from https://huggingface.co/papers/2403.07691 by using log identities and exp(log(P(y|x)) = P(y|x)\n        log_odds = (policy_chosen_logps - policy_rejected_logps) - (\n            torch.log1p(-torch.exp(policy_chosen_logps)) - torch.log1p(-torch.exp(policy_rejected_logps))\n        )\n        sig_ratio = F.sigmoid(log_odds)\n        ratio = torch.log(sig_ratio)\n        losses = self.beta * ratio\n\n        chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach()\n        rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach()\n\n        return losses, chosen_rewards, rejected_rewards, torch.mean(ratio), torch.mean(log_odds)\n\n    @staticmethod\n    def get_batch_logps(\n        logits: torch.FloatTensor,\n        labels: torch.LongTensor,\n        average_log_prob: bool = False,\n        label_pad_token_id: int = -100,\n        is_encoder_decoder: bool = False,\n    ) -> torch.FloatTensor:\n        \"\"\"Compute the log probabilities of the given labels under the given logits.\n\n        Args:\n            logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n            labels: Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are ignored. Shape: (batch_size, sequence_length)\n            average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the (non-masked) tokens.\n            label_pad_token_id: The label pad token id.\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n\n        Returns:\n            A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the given logits.\n        \"\"\"\n        if logits.shape[:-1] != labels.shape:\n            raise ValueError(\"Logits (batch and sequence length dim) and labels must have the same shape.\")\n\n        if not is_encoder_decoder:\n            labels = labels[:, 1:].clone()\n            logits = logits[:, :-1, :]\n        loss_mask = labels != label_pad_token_id\n\n        # dummy token; we'll ignore the losses on these tokens later\n        labels = torch.where(labels == label_pad_token_id, 0, labels)\n\n        per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n\n        if average_log_prob:\n            return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n        else:\n            return (per_token_logps * loss_mask).sum(-1)\n\n    def concatenated_forward(\n        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.\n\n        We do this to avoid doing two forward passes, because it's faster for FSDP.\n        \"\"\"\n        concatenated_batch = self.concatenated_inputs(\n            batch,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n            padding_value=self.padding_value,\n            device=self.accelerator.device,\n        )\n        len_chosen = batch[\"chosen_labels\"].shape[0]\n\n        model_kwargs = (\n            {\n                \"decoder_input_ids\": self._shift_right(concatenated_batch[\"concatenated_labels\"]),\n            }\n            if self.is_encoder_decoder\n            else {}\n        )\n\n        if self.aux_loss_enabled:\n            model_kwargs[\"output_router_logits\"] = True\n\n        outputs = model(\n            concatenated_batch[\"concatenated_input_ids\"],\n            attention_mask=concatenated_batch[\"concatenated_attention_mask\"],\n            use_cache=False,\n            **model_kwargs,\n        )\n        all_logits = outputs.logits\n\n        def cross_entropy_loss(logits, labels):\n            if not self.is_encoder_decoder:\n                # Shift so that tokens < n predict n\n                logits = logits[..., :-1, :].contiguous()\n                labels = labels[..., 1:].contiguous()\n            # Flatten the tokens\n            loss_fct = nn.CrossEntropyLoss()\n            logits = logits.view(-1, logits.shape[-1])\n            labels = labels.view(-1)\n            # Enable model parallelism\n            labels = labels.to(logits.device)\n            loss = loss_fct(logits, labels)\n            return loss\n\n        if self.is_encoder_decoder:\n            labels = concatenated_batch[\"concatenated_labels\"].clone()\n        else:\n            labels = concatenated_batch[\"concatenated_input_ids\"].clone()\n            attention_mask = concatenated_batch[\"concatenated_attention_mask\"]\n            labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id)\n\n        chosen_nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen])\n\n        all_logps = self.get_batch_logps(\n            all_logits,\n            concatenated_batch[\"concatenated_labels\"],\n            average_log_prob=True,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        chosen_logps = all_logps[:len_chosen]\n        rejected_logps = all_logps[len_chosen:]\n\n        chosen_logits = all_logits[:len_chosen]\n        rejected_logits = all_logits[len_chosen:]\n\n        if self.aux_loss_enabled:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss, outputs.aux_loss)\n\n        return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss)\n\n    def get_batch_loss_metrics(\n        self,\n        model,\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        train_eval: Literal[\"train\", \"eval\"] = \"train\",\n    ):\n        \"\"\"Compute the ORPO loss and other metrics for the given batch of inputs for train or test.\"\"\"\n        metrics = {}\n\n        forward_output = self.concatenated_forward(model, batch)\n        (\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_chosen_logits,\n            policy_rejected_logits,\n            policy_nll_loss,\n        ) = forward_output[:5]\n        if self.aux_loss_enabled:\n            aux_loss = forward_output[5]\n\n        losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = self.odds_ratio_loss(\n            policy_chosen_logps, policy_rejected_logps\n        )\n        # full ORPO loss\n        loss = policy_nll_loss - losses.mean()\n\n        reward_accuracies = (chosen_rewards > rejected_rewards).float()\n\n        prefix = \"eval_\" if train_eval == \"eval\" else \"\"\n        metrics[f\"{prefix}rewards/chosen\"] = chosen_rewards.mean()\n        metrics[f\"{prefix}rewards/rejected\"] = rejected_rewards.mean()\n        metrics[f\"{prefix}rewards/accuracies\"] = reward_accuracies.mean()\n        metrics[f\"{prefix}rewards/margins\"] = (chosen_rewards - rejected_rewards).mean()\n        metrics[f\"{prefix}logps/rejected\"] = policy_rejected_logps.detach().mean()\n        metrics[f\"{prefix}logps/chosen\"] = policy_chosen_logps.detach().mean()\n        metrics[f\"{prefix}logits/rejected\"] = policy_rejected_logits.detach().mean()\n        metrics[f\"{prefix}logits/chosen\"] = policy_chosen_logits.detach().mean()\n        metrics[f\"{prefix}nll_loss\"] = policy_nll_loss.detach().mean()\n        metrics[f\"{prefix}log_odds_ratio\"] = log_odds_ratio\n        metrics[f\"{prefix}log_odds_chosen\"] = log_odds_chosen\n        if is_torch_xla_available():\n            xm.mark_step()  # needed because .item() calls\n        for k, v in metrics.items():\n            metrics[k] = v.item()\n        if self.aux_loss_enabled:\n            loss += getattr(model.config, \"router_aux_loss_coef\", 0.0) * aux_loss\n\n        return loss, metrics\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n\n        compute_loss_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with compute_loss_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"train\")\n\n        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:\n        loss = loss.to(self.args.device)\n\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"train\")\n\n        if return_outputs:\n            return (loss, metrics)\n        return loss\n\n    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:\n        \"\"\"Generate samples from the model and reference model for the given batch of inputs.\"\"\"\n\n        # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with\n        # the torch cuda amp context manager as some hidden states are silently casted to full precision.\n        generate_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with generate_context_manager:\n            policy_output = model.generate(\n                input_ids=batch[\"prompt_input_ids\"],\n                attention_mask=batch[\"prompt_attention_mask\"],\n                max_length=self.max_length,\n                do_sample=True,\n                pad_token_id=self.tokenizer.pad_token_id,\n            )\n\n        policy_output = pad_to_length(policy_output, self.max_length, self.tokenizer.pad_token_id)\n        policy_output_decoded = self.tokenizer.batch_decode(policy_output, skip_special_tokens=True)\n\n        return policy_output_decoded\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ):\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        if ignore_keys is None:\n            if hasattr(model, \"config\"):\n                ignore_keys = getattr(model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        prediction_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with torch.no_grad(), prediction_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"eval\")\n\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"eval\")\n\n        if prediction_loss_only:\n            return (loss.detach(), None, None)\n\n        # logits for the chosen and rejected samples from model\n        logits_dict = {\n            \"eval_logits/chosen\": metrics[\"eval_logits/chosen\"],\n            \"eval_logits/rejected\": metrics[\"eval_logits/rejected\"],\n        }\n        logits = tuple(v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys)\n        logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)\n        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)\n\n        return (loss.detach(), logits, labels)\n\n    def store_metrics(self, metrics: Dict[str, float], train_eval: Literal[\"train\", \"eval\"] = \"train\") -> None:\n        for key, value in metrics.items():\n            self._stored_metrics[train_eval][key].append(value)\n\n    def evaluation_loop(\n        self,\n        dataloader: DataLoader,\n        description: str,\n        prediction_loss_only: Optional[bool] = None,\n        ignore_keys: Optional[List[str]] = None,\n        metric_key_prefix: str = \"eval\",\n    ) -> EvalLoopOutput:\n        \"\"\"\n        Overriding built-in evaluation loop to store metrics for each batch.\n        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.\n\n        Works both with or without labels.\n        \"\"\"\n\n        # Sample and save to game log if requested (for one batch to save time)\n        if self.generate_during_eval:\n            # Generate random indices within the range of the total number of samples\n            num_samples = len(dataloader.dataset)\n            random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)\n\n            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader\n            random_batch_dataset = dataloader.dataset.select(random_indices)\n            random_batch = self.data_collator(random_batch_dataset)\n            random_batch = self._prepare_inputs(random_batch)\n\n            policy_output_decoded = self.get_batch_samples(self.model, random_batch)\n\n            self.log(\n                {\n                    \"game_log\": wandb.Table(\n                        columns=[\"Prompt\", \"Policy\"],\n                        rows=[\n                            [prompt, pol[len(prompt) :]]\n                            for prompt, pol in zip(random_batch[\"prompt\"], policy_output_decoded)\n                        ],\n                    )\n                }\n            )\n            self.state.log_history.pop()\n\n        # Base evaluation\n        initial_output = super().evaluation_loop(\n            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix\n        )\n\n        return initial_output\n\n    def log(self, logs: Dict[str, float]) -> None:\n        \"\"\"\n        Log `logs` on the various objects watching training, including stored metrics.\n\n        Args:\n            logs (`Dict[str, float]`):\n                The values to log.\n        \"\"\"\n        # logs either has 'loss' or 'eval_loss'\n        train_eval = \"train\" if \"loss\" in logs else \"eval\"\n        # Add averaged stored metrics to logs\n        for key, metrics in self._stored_metrics[train_eval].items():\n            logs[key] = torch.tensor(metrics).mean().item()\n        del self._stored_metrics[train_eval]\n        return super().log(logs)\n\n    def _shift_right(self, input_ids):\n        if self.decoder_start_token_id is None:\n            raise ValueError(\n                \"model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id.\"\n            )\n\n        # shift inputs to the right\n        if is_torch_fx_proxy(input_ids):\n            # Item assignment is not supported natively for proxies.\n            shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id)\n            shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)\n        else:\n            shifted_input_ids = input_ids.new_zeros(input_ids.shape)\n            shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()\n            shifted_input_ids[..., 0] = self.decoder_start_token_id\n\n        if self.pad_token_id is None:\n            raise ValueError(\"model.config.pad_token_id has to be defined.\")\n        # replace possible -100 values in labels by `pad_token_id`\n        shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id)\n\n        return shifted_input_ids\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"orpo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport warnings\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Any, Dict, Literal, Optional\n\nfrom transformers import TrainingArguments\n\n\nclass FDivergenceType(Enum):\n    REVERSE_KL = \"reverse_kl\"\n    JS_DIVERGENCE = \"js_divergence\"\n    ALPHA_DIVERGENCE = \"alpha_divergence\"\n\n\nclass FDivergenceConstants:\n    ALPHA_DIVERGENCE_COEF_KEY = \"alpha_divergence_coef\"\n    ALPHA_DIVERGENCE_COEF_DEFAULT = 1.0\n\n\n@dataclass\nclass DPOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`DPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        beta (`float`, *optional*, defaults to `0.1`):\n            Parameter controlling the deviation from the reference model. Higher β means less deviation from the\n            reference model. For the IPO loss (`loss_type=\"ipo\"`), β is the regularization parameter denoted by τ in\n            the [paper](https://huggingface.co/papers/2310.12036).\n        label_smoothing (`float`, *optional*, defaults to `0.0`):\n            Robust DPO label smoothing parameter from the [cDPO](https://ericmitchell.ai/cdpo.pdf) report and\n            [Robust DPO](https://huggingface.co/papers/2403.00409) paper that should be between `0.0` and `0.5`.\n        loss_type (`str`, *optional*, defaults to `\"sigmoid\"`):\n            Type of loss to use. Possible values are:\n\n                - `\"sigmoid\"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper.\n                - `\"hinge\"`: hinge loss on the normalized likelihood from the [SLiC](https://huggingface.co/papers/2305.10425) paper.\n                - `\"ipo\"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper.\n                - `\"exo_pair\"`: pairwise EXO loss from the [EXO](https://huggingface.co/papers/2402.00856) paper.\n                - `\"nca_pair\"`: pairwise NCA loss from the [NCA](https://huggingface.co/papers/2402.05369) paper.\n                - `\"robust\"`: unbiased estimate of the DPO loss that is robust to preference noise from the [Robust DPO](https://huggingface.co/papers/2403.00409) paper.\n                - `\"bco_pair\"`: pairwise BCO loss from the [BCO](https://huggingface.co/papers/2404.04656) paper.\n                - `\"sppo_hard\"`: SPPO loss with hard label from the [SPPO](https://huggingface.co/papers/2405.00675) paper.\n                - `\"aot\"`: AOT loss for paired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper.\n                - `\"aot_pair\"`: AOT loss for unpaired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper.\n                - `\"apo_zero\"`: APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper.\n                - `\"apo_down\"`: APO-down loss from the [APO](https://huggingface.co/papers/2408.06266) paper.\n\n        label_pad_token_id (`int`, *optional*, defaults to `-100`):\n            Label pad token id. This argument is required if you want to use the default data collator.\n        padding_value (`Optional[int]`, *optional*, defaults to `None`):\n            Padding value to use. If `None`, the padding value of the tokenizer is used.\n        truncation_mode (`str`, *optional*, defaults to `\"keep_end\"`):\n            Truncation mode to use, either `keep_end` or `keep_start`. This argument is required if you want to use the\n            default data collator.\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        max_prompt_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the prompt. This argument is required if you want to use the default data collator.\n        max_completion_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the target. This argument is required if you want to use the default data collator and\n            your model is an encoder-decoder.\n        is_encoder_decoder(`Optional[int]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        disable_dropout (`bool`, *optional*, defaults to `True`):\n            Whether to disable dropout in the model and reference model.\n        generate_during_eval (`bool`, *optional*, defaults to `False`):\n            Truncation mode to use when the prompt is too long. Possible values are `\"keep_end\"` or `\"keep_start\"`.\n            This argument is required if you want to use the default data collator.\n        precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):\n            Whether to precompute reference model log probabilities for training and evaluation datasets. This is\n            useful when training without the reference model to reduce the total GPU memory needed.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        ref_model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model\n            from a string.\n        model_adapter_name (`Optional[str]`, *optional*, defaults to `None`):\n            Name of the train target PEFT adapter, when using LoRA with multiple adapters.\n        ref_adapter_name (`Optional[str]`, *optional*, defaults to `None`):\n            Name of the reference PEFT adapter, when using LoRA with multiple adapters.\n        reference_free (`bool`, *optional*, defaults to `False`):\n            If `True`, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal\n            probability to all responses.\n        force_use_ref_model (`bool`, *optional*, defaults to `False`):\n            In case one passes a PEFT model for the active model and you want to use a different model for the\n            ref_model, set this flag to `True`.\n        f_divergence_type (`str`, *optional*, defaults to `FDivergenceType.REVERSE_KL`):\n            Type of f-divergence regularization function to compute divergence between policy and reference model.\n        f_alpha_divergence_coef (`float`, *optional*, defaults to `1.0`):\n            α coefficient in the α-divergence \\\\(u^{-\\\\alpha}\\\\) regularization function for DPO loss.\n        sync_ref_model (`bool`, *optional*, defaults to `False`):\n            When set to `True`, the reference model is synchronized with the active model every `ref_model_sync_steps`\n            steps, using the `ref_model_mixup_alpha` parameter. This synchronization originites from the\n            [TR-DPO](https://huggingface.co/papers/2404.09656) paper.\n        ref_model_mixup_alpha (`float`, *optional*, defaults to `0.9`):\n            α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix\n            between the current policy and the previous reference policy during updates. The reference policy is\n            updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`\n            To use this parameter, you must set `sync_ref_model=True`.\n        ref_model_sync_steps (`int`, *optional*, defaults to `64`):\n            τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how\n            frequently the current policy is synchronized with the reference policy. To use this parameter, you must\n            set `sync_ref_model=True`.\n        rpo_alpha (`float`, *optional*, defaults to `None`):\n            α parameter from the [RPO](https://huggingface.co/papers/2404.19733) paper (v3), which controls the\n            weighting of the NLL term in the loss. If `None`, no weighting is applied and the loss is the same as the\n            DPO loss. The paper recommends `rpo_alpha=1.0`.\n    \"\"\"\n\n    beta: float = 0.1\n    label_smoothing: float = 0.0\n    loss_type: Literal[\n        \"sigmoid\",\n        \"hinge\",\n        \"ipo\",\n        \"exo_pair\",\n        \"nca_pair\",\n        \"robust\",\n        \"bco_pair\",\n        \"sppo_hard\",\n        \"aot\",\n        \"aot_pair\",\n        \"apo_zero\",\n        \"apo_down\",\n    ] = \"sigmoid\"\n    label_pad_token_id: int = -100\n    padding_value: Optional[int] = None\n    truncation_mode: str = \"keep_end\"\n    max_length: Optional[int] = None\n    max_prompt_length: Optional[int] = None\n    max_target_length: Optional[int] = None  # deprecated in favor of max_completion_length\n    max_completion_length: Optional[int] = None\n    is_encoder_decoder: Optional[bool] = None\n    disable_dropout: bool = True\n    generate_during_eval: bool = False\n    precompute_ref_log_probs: bool = False\n    dataset_num_proc: Optional[int] = None\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    ref_model_init_kwargs: Optional[Dict[str, Any]] = None\n    model_adapter_name: Optional[str] = None\n    ref_adapter_name: Optional[str] = None\n    reference_free: bool = False\n    force_use_ref_model: bool = False\n    f_divergence_type: FDivergenceType = FDivergenceType.REVERSE_KL\n    f_alpha_divergence_coef: float = 1.0\n    sync_ref_model: bool = False\n    ref_model_mixup_alpha: float = 0.9\n    ref_model_sync_steps: int = 64\n    rpo_alpha: Optional[float] = None\n\n    def __post_init__(self):\n        if self.max_target_length is not None:\n            warnings.warn(\n                \"The `max_target_length` argument is deprecated in favor of `max_completion_length` and will be removed in a future version.\",\n                FutureWarning,\n            )\n            if self.max_completion_length is None:\n                self.max_completion_length = self.max_target_length\n\n        return super().__post_init__()\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport math\nimport os\nimport time\nimport typing\nimport warnings\nfrom contextlib import nullcontext\nfrom typing import Callable, List, Optional, Union\n\nimport datasets\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom accelerate import Accelerator\nfrom accelerate.utils import ProjectConfiguration, gather_object, is_deepspeed_available\nfrom datasets import Dataset\nfrom huggingface_hub import whoami\nfrom packaging import version\nfrom torch.optim import Adam\nfrom transformers import (\n    DataCollatorForLanguageModeling,\n    PreTrainedTokenizer,\n    PreTrainedTokenizerBase,\n    PreTrainedTokenizerFast,\n    is_torch_npu_available,\n    is_torch_xpu_available,\n)\n\nfrom ..core import (\n    WANDB_PADDING,\n    PPODecorators,\n    clip_by_value,\n    convert_to_scalar,\n    entropy_from_logits,\n    flatten_dict,\n    logprobs_from_logits,\n    masked_mean,\n    masked_var,\n    masked_whiten,\n    set_seed,\n    stack_dicts,\n    stats_to_np,\n)\nfrom ..import_utils import is_torch_greater_2_0\nfrom ..models import (\n    SUPPORTED_ARCHITECTURES,\n    PreTrainedModelWrapper,\n    create_reference_model,\n    unwrap_model_for_generation,\n)\nfrom . import AdaptiveKLController, BaseTrainer, FixedKLController, PPOConfig, RunningMoments\n\n\nif is_deepspeed_available():\n    import deepspeed\n\nMODEL_CARD_TEMPLATE = \"\"\"---\nlicense: apache-2.0\nlibrary_name: transformers\ntags:\n- trl\n- ppo\n- transformers\n- reinforcement-learning\n---\n\n# {model_name}\n\nThis is a [TRL language model](https://github.com/huggingface/trl) that has been fine-tuned with reinforcement learning to\n guide the model outputs according to a value, function, or human feedback. The model can be used for text generation.\n\n## Usage\n\nTo use this model for inference, first install the TRL library:\n\n```bash\npython -m pip install trl\n```\n\nYou can then generate text as follows:\n\n```python\nfrom transformers import pipeline\n\ngenerator = pipeline(\"text-generation\", model=\"{model_id}\")\noutputs = generator(\"Hello, my llama is cute\")\n```\n\nIf you want to use the model for training or to obtain the outputs from the value head, load the model as follows:\n\n```python\nfrom transformers import AutoTokenizer\nfrom trl import AutoModelForCausalLMWithValueHead\n\ntokenizer = AutoTokenizer.from_pretrained(\"{model_id}\")\nmodel = AutoModelForCausalLMWithValueHead.from_pretrained(\"{model_id}\")\n\ninputs = tokenizer(\"Hello, my llama is cute\", return_tensors=\"pt\")\noutputs = model(**inputs, labels=inputs[\"input_ids\"])\n```\n\"\"\"\n\n\nclass PPOTrainer(BaseTrainer):\n    \"\"\"\n    The PPOTrainer uses Proximal Policy Optimization to optimise language models.\n    Note, this trainer is heavily inspired by the original OpenAI learning to summarize work here:\n    https://github.com/openai/summarize-from-feedback\n\n    Attributes:\n        **config** (`PPOConfig`) -- Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more\n            details.\n        **model** (`PreTrainedModelWrapper`) -- Model to be optimized, Hugging Face transformer model with a value head.\n            Check the documentation of `PreTrainedModelWrapper` for more details.\n        **ref_model** (`PreTrainedModelWrapper`, *optional*) -- Reference model to be used for KL penalty, Hugging Face\n            transformer model with a casual language modelling head. Check the documentation of `PreTrainedModelWrapper`\n            for more details. If no reference model is provided, the trainer will create a reference model with the same\n             architecture as the model to be optimized with shared layers.\n        **tokenizer** (`PreTrainedTokenizerBase`) -- Tokenizer to be used for encoding the\n            data. Check the documentation of `transformers.PreTrainedTokenizer` and\n            `transformers.PreTrainedTokenizerFast` for more details.\n        **dataset** (Union[`torch.utils.data.Dataset`, `datasets.Dataset`], *optional*) -- PyTorch dataset or Hugging\n            Face dataset. This is used to create a PyTorch dataloader. If no dataset is provided, the dataloader must be\n             created outside the trainer users needs to design their own dataloader and make sure the batch\n            size that is used is the same as the one specified in the configuration object.\n        **optimizer** (`torch.optim.Optimizer`, *optional*) -- Optimizer to be used for training. If no optimizer is\n            provided, the trainer will create an Adam optimizer with the learning rate specified in the configuration\n            object.\n        **data_collator** (DataCollatorForLanguageModeling, *optional*) -- Data collator to be used for training and\n            passed along the dataloader\n        **num_shared_layers** (int, *optional*) -- Number of layers to be shared between the model and the reference\n            model, if no reference model is passed. If no number is provided, all the layers will be shared.\n        **lr_scheduler** (`torch.optim.lr_scheduler`, *optional*) -- Learning rate scheduler to be used for training.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"ppo\"]\n\n    def __init__(\n        self,\n        config: Optional[PPOConfig] = None,\n        model: Optional[PreTrainedModelWrapper] = None,\n        ref_model: Optional[PreTrainedModelWrapper] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        dataset: Optional[Union[torch.utils.data.Dataset, Dataset]] = None,\n        optimizer: Optional[torch.optim.Optimizer] = None,\n        data_collator: Optional[typing.Callable] = None,\n        num_shared_layers: Optional[int] = None,\n        lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,\n        training_data_collator: Optional[typing.Callable] = None,\n    ):\n        \"\"\"\n        Initialize PPOTrainer.\n\n        Args:\n            config (`PPOConfig`):\n                Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more details.\n            model (`PreTrainedModelWrapper`):\n                Hugging Face transformer model with a value head.\n            ref_model (`PreTrainedModelWrapper`):\n                Hugging Face transformer model with a casual language modelling head. Used for KL penalty\n            tokenizer (`transformers.PreTrainedTokenizerBase`):\n                Hugging Face tokenizer\n            dataset (Optional[Union[`torch.utils.data.Dataset`, `datasets.Dataset`]]):\n                PyTorch dataset or Hugging Face dataset. If a Hugging Face dataset is passed, the dataset\n                will be preprocessed by removing the columns that are not used by the model. If none is passed,\n                a warning will be raised in a multi-GPU setting.\n            optimizer (`Optional[torch.optim.Optimizer]`):\n                Optimizer used for training. If `None`, the `Adam` is used as default.\n            data_collator (Optional[function]):\n                Data collator function that is going to be used for `prepare_dataloader` method. Note this collator\n                is different from the one we use for training. Pass a valid `training_data_collator` instead.\n            num_shared_layers (Optional[int]):\n                Number of shared layers between the model and the reference model. If `None`, all layers are shared.\n                used only if `ref_model` is `None`.\n            lr_scheduler (`Optional[torch.optim.lr_scheduler]`):\n                Learning rate scheduler used for training.\n            training_data_collator (Optional[function]):\n                Custom data collator used for training.\n        \"\"\"\n        warnings.warn(\n            \"`PPOTrainer` is deprecated and will be removed in trl v0.12. Please use `PPOv2Trainer` instead.\",\n            FutureWarning,\n        )\n        super().__init__(config)\n\n        # initial seed for reproducible experiments\n        set_seed(config.seed)\n\n        # Step 0: check positional arguments validity\n        if not isinstance(config, PPOConfig):\n            raise ValueError(f\"config must be a PPOConfig, got {type(config)}\")\n        if not isinstance(tokenizer, (PreTrainedTokenizerBase)):\n            raise ValueError(\n                f\"tokenizer must be a PreTrainedTokenizerBase like a PreTrainedTokenizer or a PreTrainedTokenizerFast, got {type(tokenizer)}\"\n            )\n        if not isinstance(model, (SUPPORTED_ARCHITECTURES)):\n            raise ValueError(\n                f\"model must be a PreTrainedModelWrapper, got {type(model)} - supported architectures are: {SUPPORTED_ARCHITECTURES}\"\n            )\n        # Step 1: Initialize Accelerator\n        self.accelerator = Accelerator(\n            log_with=config.log_with,\n            gradient_accumulation_steps=config.gradient_accumulation_steps,\n            project_config=ProjectConfiguration(**config.project_kwargs),\n            **config.accelerator_kwargs,\n        )\n\n        # Step 1.1 Runtime variables filled by the accelerator\n        config.world_size = self.accelerator.num_processes\n        config.global_backward_batch_size = config.backward_batch_size * config.world_size\n        config.global_batch_size = config.batch_size * config.world_size\n\n        self.model = model\n        self.model_params = filter(lambda p: p.requires_grad, self.model.parameters())\n        self.is_encoder_decoder = hasattr(self.model, \"is_encoder_decoder\")\n        self.is_peft_model = getattr(self.model, \"is_peft_model\", False)\n        config.is_encoder_decoder = self.is_encoder_decoder\n        config.is_peft_model = self.is_peft_model\n\n        is_using_tensorboard = config.log_with is not None and config.log_with == \"tensorboard\"\n        self.accelerator.init_trackers(\n            config.tracker_project_name,\n            config=dict(trl_ppo_trainer_config=config.to_dict()) if not is_using_tensorboard else config.to_dict(),\n            init_kwargs=config.tracker_kwargs,\n        )\n        self.is_using_text_environment = getattr(config, \"use_text_environment\", False)\n\n        if isinstance(ref_model, SUPPORTED_ARCHITECTURES):\n            self.ref_model = ref_model\n            if num_shared_layers is not None:\n                warnings.warn(\n                    \"num_shared_layers is ignored when ref_model is provided. Two different models are used for the \"\n                    \"model and the reference model and no layers are shared.\",\n                    UserWarning,\n                )\n        elif ref_model is None and not self.is_peft_model:\n            self.ref_model = create_reference_model(self.model, num_shared_layers=num_shared_layers)\n        elif self.is_peft_model:\n            self.ref_model = None\n        else:\n            raise ValueError(\n                f\"ref_model must be a PreTrainedModelWrapper or `None`, got {type(ref_model)} - supported \"\n                f\"architectures are: {SUPPORTED_ARCHITECTURES} \"\n            )\n        self.optional_peft_ctx = (\n            self.accelerator.unwrap_model(self.model).pretrained_model.disable_adapter\n            if self.is_peft_model\n            else nullcontext\n        )\n\n        if not (isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast)):\n            raise ValueError(\n                \"tokenizer must be a transformers.PreTrainedTokenizer or transformers.PreTrainedTokenizerFast\"\n            )\n        self.tokenizer = tokenizer\n\n        if dataset is not None and not (isinstance(dataset, torch.utils.data.Dataset) or isinstance(dataset, Dataset)):\n            raise ValueError(\"dataset must be a torch.utils.data.Dataset or datasets.Dataset\")\n        elif dataset is None:\n            warnings.warn(\n                \"No dataset is provided. Make sure to set config.batch_size to the correct value before training.\",\n                UserWarning,\n            )\n        self.dataset = dataset\n        self._signature_columns = None\n        if self.dataset is not None:\n            self.dataloader = self.prepare_dataloader(self.dataset, data_collator)\n        elif self.dataset is None and self.accelerator.num_processes > 1:\n            warnings.warn(\n                \"No dataset is provided. In a multi-GPU setting, this will lead to an error. You should\"\n                \" prepare your dataloader yourself with `dataloader = ppo_trainer.accelerator.prepare(dataloader)`\"\n                \" and using `torch.utils.data.DataLoader`, or pass a dataset to the `PPOTrainer`. Please \"\n                \" refer to the documentation for more details.\",\n                UserWarning,\n            )\n            self.dataloader = None\n        else:\n            self.dataloader = None\n\n        # Step 3: Initialize optimizer and data collator\n        if training_data_collator is None:\n            self.data_collator = DataCollatorForLanguageModeling(self.tokenizer, mlm=False)\n        else:\n            self.data_collator = training_data_collator\n        if optimizer is None:\n            self.optimizer = Adam(\n                filter(lambda p: p.requires_grad, self.model.parameters()),\n                lr=self.config.learning_rate,\n            )\n        else:\n            self.optimizer = optimizer\n\n        self.lr_scheduler = lr_scheduler\n        if self.lr_scheduler is not None:\n            lr_scheduler_class = (\n                torch.optim.lr_scheduler._LRScheduler\n                if not is_torch_greater_2_0()\n                else torch.optim.lr_scheduler.LRScheduler\n            )\n\n            if not isinstance(self.lr_scheduler, lr_scheduler_class):\n                raise ValueError(\n                    \"lr_scheduler must be a torch.optim.lr_scheduler._LRScheduler or torch.optim.lr_scheduler.LRScheduler (for torch >= 2.0)\"\n                )\n\n        if self.config.adap_kl_ctrl:\n            self.kl_ctl = AdaptiveKLController(self.config.init_kl_coef, self.config.target, self.config.horizon)\n        else:\n            self.kl_ctl = FixedKLController(self.config.init_kl_coef)\n\n        # Safety checkers for DS integration\n        is_deepspeed_used = self.accelerator.distributed_type == \"DEEPSPEED\" and hasattr(\n            self.accelerator.state, \"deepspeed_plugin\"\n        )\n\n        if config.gradient_checkpointing:\n            self.model.gradient_checkpointing_enable()\n\n            if hasattr(self.model, \"enable_input_require_grads\"):\n                self.model.enable_input_require_grads()\n            else:\n                # For backward compatibility with older versions of transformers\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                self.model.pretrained_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        (\n            self.model,\n            self.optimizer,\n            self.data_collator,\n            self.dataloader,\n            self.lr_scheduler,\n        ) = self.accelerator.prepare(\n            self.model,\n            self.optimizer,\n            self.data_collator,\n            self.dataloader,\n            self.lr_scheduler,\n        )\n        if is_deepspeed_used:\n            # Quantized models are already set on the correct device\n            if not self.is_peft_model and not (\n                getattr(self.ref_model.pretrained_model, \"is_loaded_in_8bit\", False)\n                or getattr(self.ref_model.pretrained_model, \"is_loaded_in_4bit\", False)\n            ):\n                self.ref_model = self._prepare_deepspeed(self.ref_model)\n        else:\n            self.ref_model = self.accelerator.prepare(self.ref_model)\n\n        # In a distributed setup, only logging needs to be performed on the main process\n        # check: https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html\n        # or: https://discuss.pytorch.org/t/use-distributed-data-parallel-correctly/82500/11\n        self.is_distributed = self.accelerator.num_processes > 1\n\n        # init the current step\n        self.current_step = 0\n\n        # init variables for pushing model to hub\n        if config.push_to_hub_if_best_kwargs:\n            if \"repo_id\" not in config.push_to_hub_if_best_kwargs:\n                raise ValueError(\"You have to specify repo_id in order to push the model to the hub!\")\n            self.push_to_hub_kwargs = config.push_to_hub_if_best_kwargs\n            self.compare_step = 0\n            self.highest_reward = torch.tensor(-float(\"inf\"))\n\n        # post process for PP\n        if not getattr(self.model, \"is_sequential_parallel\", False):\n            self.current_device = self.accelerator.device\n        else:\n            if is_torch_xpu_available():\n                self.current_device = torch.device(\"xpu:0\")\n            elif is_torch_npu_available():\n                self.current_device = torch.device(\"npu:0\")\n            else:\n                self.current_device = torch.device(\"cuda:0\")\n\n        PPODecorators.optimize_device_cache = self.config.optimize_device_cache\n\n        self.running = RunningMoments(self.accelerator)\n\n    def _filter_kwargs(self, kwargs, target_func):\n        \"\"\"\n        filter the keyword arguments that are supported by the target function.\n\n        Args:\n            kwargs (dict):\n                Keyword arguments\n            target_func (function):\n                Target function\n        \"\"\"\n        return {k: v for k, v in kwargs.items() if k in inspect.signature(target_func).parameters.keys()}\n\n    def prepare_dataloader(self, dataset: Union[torch.utils.data.Dataset, Dataset], data_collator=None):\n        \"\"\"\n        Prepare the dataloader for training.\n\n        Args:\n            dataset (Union[`torch.utils.data.Dataset`, `datasets.Dataset`]):\n                PyTorch dataset or Hugging Face dataset. If a Hugging Face dataset is passed, the dataset\n                will be preprocessed by removing the columns that are not used by the model.\n            data_collator (Optional[function]):\n                Data collator function.\n\n        Returns:\n            `torch.utils.data.DataLoader`: PyTorch dataloader\n        \"\"\"\n        if isinstance(dataset, Dataset):\n            dataset = self._remove_unused_columns(dataset)\n        dataloader = torch.utils.data.DataLoader(\n            dataset,\n            batch_size=self.config.batch_size,\n            collate_fn=data_collator,\n            shuffle=True,\n            drop_last=True,\n        )\n        return dataloader\n\n    # Adapted from transformers.Trainer._set_signature_columns_if_needed\n    def _set_signature_columns_if_needed(self):\n        if self._signature_columns is None:\n            # Inspect model forward signature to keep only the arguments it accepts.\n            signature = inspect.signature(self.model.forward)\n            self._signature_columns = list(signature.parameters.keys())\n            # label => sentiment | we need query and response for logging purpose\n            self._signature_columns += [\"label\", \"query\", \"response\"]\n\n    # Adapted from transformers.Trainer._remove_unused_columns\n    def _remove_unused_columns(self, dataset: \"Dataset\"):\n        if not self.config.remove_unused_columns:\n            return dataset\n        self._set_signature_columns_if_needed()\n        signature_columns = self._signature_columns\n\n        ignored_columns = list(set(dataset.column_names) - set(signature_columns))\n\n        columns = [k for k in signature_columns if k in dataset.column_names]\n\n        if version.parse(datasets.__version__) < version.parse(\"1.4.0\"):\n            dataset.set_format(\n                type=dataset.format[\"type\"],\n                columns=columns,\n                format_kwargs=dataset.format[\"format_kwargs\"],\n            )\n            return dataset\n        else:\n            return dataset.remove_columns(ignored_columns)\n\n    def generate(\n        self,\n        query_tensor: Union[torch.Tensor, List[torch.Tensor]],\n        length_sampler: Optional[Callable] = None,\n        batch_size: int = 4,\n        return_prompt: bool = True,\n        generate_ref_response: bool = False,\n        **generation_kwargs,\n    ):\n        \"\"\"\n        Generate response with the model given the query tensor.\n        call the `generate` method of the model.\n\n        Args:\n            query_tensor (`torch.LongTensor`):\n                A tensor of shape (`seq_len`) containing query tokens or a list of tensors of shape (`seq_len`).\n            length_sampler (`Callable`, *optional*):\n                Callable that returns the number of newly generated tokens.\n            batch_size (`int`, *optional):\n                Batch size used for generation, defaults to `4`.\n            return_prompt (`bool`, *optional*):\n                If set to `False` the prompt is not returned but only the newly generated tokens, defaults to `True`.\n            generate_ref_response (`bool`, *optional*):\n                If set to `True` the reference response is also generated, defaults to `False`.\n            generation_kwargs (dict[str, Any]):\n                Keyword arguments for generation.\n\n        Returns:\n            `torch.LongTensor`: A tensor of shape (`batch_size`, `gen_len`) containing response tokens.\n        \"\"\"\n        if generate_ref_response:\n            ref_model = self.model if self.is_peft_model else self.ref_model\n        if isinstance(query_tensor, List):\n            response = self._generate_batched(\n                self.model,\n                query_tensor,\n                length_sampler=length_sampler,\n                batch_size=batch_size,\n                return_prompt=return_prompt,\n                **generation_kwargs,\n            )\n            if generate_ref_response:\n                ref_response = self._generate_batched(\n                    ref_model,\n                    query_tensor,\n                    length_sampler=length_sampler,\n                    batch_size=batch_size,\n                    return_prompt=return_prompt,\n                    **generation_kwargs,\n                )\n\n        else:\n            if len(query_tensor.shape) == 2:\n                raise ValueError(\n                    \"query_tensor must be a tensor of shape (`seq_len`) or a list of tensors of shape (`seq_len`)\"\n                )\n\n            if length_sampler is not None:\n                generation_kwargs[\"max_new_tokens\"] = length_sampler()\n\n            with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:\n                response = unwrapped_model.generate(input_ids=query_tensor.unsqueeze(dim=0), **generation_kwargs)\n\n            if generate_ref_response:\n                with unwrap_model_for_generation(\n                    ref_model, self.accelerator, is_peft_model=self.is_peft_model\n                ) as unwrapped_model:\n                    ref_response = unwrapped_model.generate(\n                        input_ids=query_tensor.unsqueeze(dim=0), **generation_kwargs\n                    )\n\n            if not return_prompt and not self.is_encoder_decoder:\n                response = response[:, query_tensor.shape[0] :]\n                if generate_ref_response:\n                    ref_response = ref_response[:, query_tensor.shape[0] :]\n\n        if generate_ref_response:\n            return response, ref_response\n        return response\n\n    def _generate_batched(\n        self,\n        model: PreTrainedModelWrapper,\n        query_tensors: List[torch.Tensor],\n        length_sampler: Optional[Callable] = None,\n        batch_size: int = 4,\n        return_prompt: bool = True,\n        pad_to_multiple_of: Optional[int] = None,\n        remove_padding: bool = True,\n        **generation_kwargs,\n    ):\n        outputs = []\n\n        padding_side_default = self.tokenizer.padding_side\n        if not self.is_encoder_decoder:\n            self.tokenizer.padding_side = \"left\"\n\n        # in case we have fewer examples than bs\n        batch_size = min(len(query_tensors), batch_size)\n\n        for i in range(0, len(query_tensors), batch_size):\n            if length_sampler is not None:\n                generation_kwargs[\"max_new_tokens\"] = length_sampler()\n\n            # prevent overflow if query tensors are not even multiple of bs\n            end_index = min(len(query_tensors), i + batch_size)\n\n            batch = query_tensors[i:end_index]\n            batch_mask = [torch.ones_like(element) for element in batch]\n            inputs = {\"input_ids\": batch, \"attention_mask\": batch_mask}\n\n            padded_inputs = self.tokenizer.pad(\n                inputs,\n                padding=True,\n                max_length=None,\n                pad_to_multiple_of=pad_to_multiple_of,\n                return_tensors=\"pt\",\n            ).to(self.current_device)\n\n            with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n                generations = unwrapped_model.generate(**padded_inputs, **generation_kwargs)\n\n            for generation, mask in zip(generations, padded_inputs[\"attention_mask\"]):\n                if not self.is_encoder_decoder:\n                    output = generation[(1 - mask).sum() :]  # remove padding\n                else:\n                    output = generation\n\n                if not return_prompt and not self.is_encoder_decoder:\n                    output = output[(mask).sum() :]  # remove prompt\n\n                if remove_padding and self.tokenizer.eos_token_id in output:\n                    pad_mask = output == self.tokenizer.eos_token_id\n                    pad_start = torch.nonzero(pad_mask, as_tuple=False)[0, 0].item()\n                    output = output[: pad_start + 1]  # keep the eos token at the end\n\n                outputs.append(output)\n\n        self.tokenizer.padding_side = padding_side_default\n        return outputs\n\n    def _step_safety_checker(\n        self,\n        batch_size: int,\n        queries: List[torch.LongTensor],\n        responses: List[torch.LongTensor],\n        scores: List[torch.FloatTensor],\n        masks: Optional[List[torch.LongTensor]] = None,\n    ):\n        \"\"\"\n        Check if the input data is valid for training.\n\n        Args:\n            batch_size (int):\n                Batch size from the config file.\n            queries (List[`torch.LongTensor`]):\n                List of tensors containing the encoded queries of shape (`query_length`)\n            responses (List[`torch.LongTensor`]):\n                List of tensors containing the encoded responses of shape (`response_length`)\n            scores (List[`torch.FloatTensor`]):\n                List of tensors containing the scores.\n            masks (List[`torch.LongTensor`], *optional*):\n                list of optional tensors containing the masks of shape (`response_length`)\n\n        Returns:\n            `tuple`: The input processed data.\n        \"\"\"\n        for name, tensor_list in zip([\"queries\", \"responses\", \"scores\"], [queries, responses, scores]):\n            if not isinstance(tensor_list, list):\n                raise ValueError(f\"{name} must be a list of tensors - got {type(tensor_list)}\")\n            if not isinstance(tensor_list[0], torch.Tensor):\n                raise ValueError(f\"Elements in {name} must be tensors - got {type(tensor_list[0])}\")\n            if batch_size is not None and len(tensor_list) != batch_size:\n                raise ValueError(\n                    f\"Batch size ({batch_size}) does not match number of examples - but got {len(tensor_list)} for: {name}\"\n                )\n\n        # add queries, scores and responses on the correct device\n        queries = [tensor.to(self.current_device) for tensor in queries]\n        responses = [tensor.to(self.current_device) for tensor in responses]\n        scores = [tensor.to(self.current_device) for tensor in scores]\n        masks = [tensor.to(self.current_device) for tensor in masks] if masks is not None else None\n\n        # squeeze scores if needed\n        for i, score in enumerate(scores):\n            if score.dim() > 1:\n                raise ValueError(f\"Scores must be 1-dimensional - got {score.dim()} for {score}\")\n            elif score.dim() == 1:\n                scores[i] = score.squeeze()\n\n        return queries, responses, scores, masks\n\n    @PPODecorators.empty_device_cache()\n    def step(\n        self,\n        queries: List[torch.LongTensor],\n        responses: List[torch.LongTensor],\n        scores: List[torch.FloatTensor],\n        response_masks: Optional[List[torch.LongTensor]] = None,\n    ):\n        \"\"\"\n        Run a PPO optimisation step given a list of queries, model responses, and rewards.\n\n        Args:\n            queries (List[`torch.LongTensor`]):\n                List of tensors containing the encoded queries of shape (`query_length`)\n            responses (List[`torch.LongTensor`]):\n                List of tensors containing the encoded responses of shape (`response_length`)\n            scores (List[`torch.FloatTensor`]):\n                List of tensors containing the scores.\n            response_masks (List[`torch.FloatTensor`], *optional*)):\n                List of tensors containing masks of the response tokens.\n\n        Returns:\n            `dict[str, Any]`: A summary of the training statistics\n        \"\"\"\n        bs = self.config.batch_size\n\n        queries, responses, scores, response_masks = self._step_safety_checker(\n            bs, queries, responses, scores, response_masks\n        )\n        scores = torch.tensor(scores, device=self.current_device)\n        if self.config.use_score_scaling:\n            # Score scaling\n            scores_mean, scores_std = self.running.update(scores)\n            tensor_to_kwargs = dict(dtype=scores.dtype, device=scores.device)\n            score_scaling_factor = self.running.std.to(**tensor_to_kwargs) + torch.finfo(scores.dtype).eps\n            if self.config.use_score_norm:\n                scores = (scores - self.running.mean.to(**tensor_to_kwargs)) / score_scaling_factor\n            else:\n                scores /= score_scaling_factor\n\n        if self.config.score_clip is not None:\n            # Score clipping\n            scores_dtype = scores.dtype\n            scores = torch.clip(scores.float(), -self.config.score_clip, self.config.score_clip).to(dtype=scores_dtype)\n\n        # if we want to push best model to the hub\n        if hasattr(self, \"highest_reward\"):\n            if self.compare_step % self.config.compare_steps == 0:\n                curr_mean_reward = scores.mean()\n                # if the best reward ever seen\n                if curr_mean_reward > self.highest_reward:\n                    self.highest_reward = curr_mean_reward\n                    # push model to hub\n                    self.push_to_hub(**self.push_to_hub_kwargs)\n            self.compare_step += 1\n\n        timing = dict()\n        t0 = time.time()\n\n        t = time.time()\n\n        model_inputs = self.prepare_model_inputs(queries, responses)\n\n        if self.is_distributed:\n            pad_first = self.tokenizer.padding_side == \"left\"\n\n            model_inputs[\"input_ids\"] = self.accelerator.pad_across_processes(\n                model_inputs[\"input_ids\"],\n                dim=1,\n                pad_index=self.tokenizer.pad_token_id,\n                pad_first=pad_first,\n            )\n            model_inputs[\"attention_mask\"] = self.accelerator.pad_across_processes(\n                model_inputs[\"attention_mask\"], dim=1, pad_index=0, pad_first=pad_first\n            )\n            if self.is_encoder_decoder:\n                model_inputs[\"decoder_input_ids\"] = self.accelerator.pad_across_processes(\n                    model_inputs[\"decoder_input_ids\"],\n                    dim=1,\n                    pad_index=self.tokenizer.pad_token_id,\n                    pad_first=pad_first,\n                )\n                model_inputs[\"decoder_attention_mask\"] = self.accelerator.pad_across_processes(\n                    model_inputs[\"decoder_attention_mask\"],\n                    dim=1,\n                    pad_index=0,\n                    pad_first=pad_first,\n                )\n\n        model_inputs_names = list(model_inputs.keys())\n\n        full_kl_penalty = self.config.kl_penalty == \"full\"\n\n        with torch.no_grad():\n            all_logprobs, logits_or_none, values, masks = self.batched_forward_pass(\n                self.model,\n                queries,\n                responses,\n                model_inputs,\n                response_masks=response_masks,\n                return_logits=full_kl_penalty,\n            )\n            with self.optional_peft_ctx():\n                ref_logprobs, ref_logits_or_none, _, _ = self.batched_forward_pass(\n                    self.model if self.is_peft_model else self.ref_model,\n                    queries,\n                    responses,\n                    model_inputs,\n                    return_logits=full_kl_penalty,\n                )\n\n        timing[\"time/ppo/forward_pass\"] = time.time() - t\n\n        with torch.no_grad():\n            t = time.time()\n            if full_kl_penalty:\n                active_full_logprobs = logprobs_from_logits(logits_or_none, None, gather=False)\n                ref_full_logprobs = logprobs_from_logits(ref_logits_or_none, None, gather=False)\n\n                rewards, non_score_reward, kls = self.compute_rewards(\n                    scores, active_full_logprobs, ref_full_logprobs, masks\n                )\n            else:\n                rewards, non_score_reward, kls = self.compute_rewards(scores, all_logprobs, ref_logprobs, masks)\n            timing[\"time/ppo/compute_rewards\"] = time.time() - t\n\n            t = time.time()\n            values, advantages, returns = self.compute_advantages(values, rewards, masks)\n            timing[\"time/ppo/compute_advantages\"] = time.time() - t\n\n        # upcast to float32 to avoid dataset issues\n        batch_dict = {\n            \"queries\": queries,\n            \"responses\": responses,\n            \"logprobs\": all_logprobs.to(torch.float32),\n            \"values\": values.to(torch.float32),\n            \"masks\": masks,\n            \"advantages\": advantages,\n            \"returns\": returns,\n        }\n        batch_dict.update(model_inputs)\n\n        t = time.time()\n        all_stats = []\n        early_stop = False\n        for _ in range(self.config.ppo_epochs):\n            if early_stop:\n                break\n            b_inds = np.random.permutation(bs)\n            for backward_batch_start in range(0, bs, self.config.backward_batch_size):\n                backward_batch_end = backward_batch_start + self.config.backward_batch_size\n                backward_batch_inds = b_inds[backward_batch_start:backward_batch_end]\n\n                for mini_batch_start in range(0, self.config.backward_batch_size, self.config.mini_batch_size):\n                    mini_batch_end = mini_batch_start + self.config.mini_batch_size\n                    mini_batch_inds = backward_batch_inds[mini_batch_start:mini_batch_end]\n                    mini_batch_dict = {\n                        \"logprobs\": batch_dict[\"logprobs\"][mini_batch_inds],\n                        \"values\": batch_dict[\"values\"][mini_batch_inds],\n                        \"masks\": batch_dict[\"masks\"][mini_batch_inds],\n                        # hacks: the queries and responses are ragged.\n                        \"queries\": [batch_dict[\"queries\"][i] for i in mini_batch_inds],\n                        \"responses\": [batch_dict[\"responses\"][i] for i in mini_batch_inds],\n                        \"advantages\": batch_dict[\"advantages\"][mini_batch_inds],\n                        \"returns\": batch_dict[\"returns\"][mini_batch_inds],\n                    }\n                    for k in model_inputs_names:\n                        mini_batch_dict[k] = batch_dict[k][mini_batch_inds]\n                    with self.accelerator.accumulate(self.model):\n                        model_inputs = {k: mini_batch_dict[k] for k in model_inputs_names}\n\n                        logprobs, logits, vpreds, _ = self.batched_forward_pass(\n                            self.model,\n                            mini_batch_dict[\"queries\"],\n                            mini_batch_dict[\"responses\"],\n                            model_inputs,\n                            return_logits=True,\n                        )\n                        train_stats = self.train_minibatch(\n                            mini_batch_dict[\"logprobs\"],\n                            mini_batch_dict[\"values\"],\n                            logprobs,\n                            logits,\n                            vpreds,\n                            mini_batch_dict[\"masks\"],\n                            mini_batch_dict[\"advantages\"],\n                            mini_batch_dict[\"returns\"],\n                        )\n                        all_stats.append(train_stats)\n\n            # typically, early stopping is done at the epoch level\n            if self.config.early_stopping:\n                policykl = train_stats[\"policy/policykl\"]\n                early_stop = self._early_stop(policykl)\n                if early_stop:\n                    break\n\n        timing[\"time/ppo/optimize_step\"] = time.time() - t\n\n        t = time.time()\n        train_stats = stack_dicts(all_stats)\n\n        # reshape advantages/ratios such that they are not averaged.\n        train_stats[\"policy/advantages\"] = torch.flatten(train_stats[\"policy/advantages\"]).unsqueeze(0)\n        train_stats[\"policy/advantages\"] = torch.nan_to_num(train_stats[\"policy/advantages\"], WANDB_PADDING)\n        train_stats[\"policy/ratio\"] = torch.flatten(train_stats[\"policy/ratio\"]).unsqueeze(0)\n\n        stats = self.record_step_stats(\n            scores=scores,\n            logprobs=all_logprobs,\n            ref_logprobs=ref_logprobs,\n            non_score_reward=non_score_reward,\n            train_stats=train_stats,\n            kl_coef=self.kl_ctl.value,\n            masks=masks,\n            queries=queries,\n            responses=responses,\n            kls=kls,\n        )\n        # Gather/Reduce stats from all processes\n        if self.is_distributed:\n            stats = self.gather_stats(stats)\n        stats = stats_to_np(stats)\n        timing[\"time/ppo/calc_stats\"] = time.time() - t\n        stats[\"ppo/learning_rate\"] = self.optimizer.param_groups[0][\"lr\"]\n\n        # Update the KL control - multiply the batch_size by the number of processes\n        self.kl_ctl.update(\n            stats[\"objective/kl\"],\n            self.config.batch_size * self.accelerator.num_processes,\n        )\n\n        # Log the total ppo time\n        timing[\"time/ppo/total\"] = time.time() - t0\n        stats.update(timing)\n\n        # post-process stats for tensorboard and other loggers\n        if self.config.log_with != \"wandb\":\n            stats = convert_to_scalar(stats)\n\n        if self.lr_scheduler is not None:\n            self.lr_scheduler.step()\n\n        return stats\n\n    def _early_stop(self, policykl):\n        r\"\"\"\n        Handles the early stopping logic. If the policy KL is greater than the target KL, then the gradient is zeroed and\n        the optimization step is skipped.\n        This also handles the multi-gpu case where the policy KL is averaged across all processes.\n\n        Args:\n            policy_kl (torch.Tensor):\n                the policy KL\n\n        Returns:\n            `bool`: whether to early stop or not\n        \"\"\"\n        early_stop = False\n        if not self.config.early_stopping:\n            return early_stop\n\n        if not self.is_distributed and policykl > 1.5 * self.config.target_kl:\n            self.optimizer.zero_grad()\n            early_stop = True\n        elif self.is_distributed:\n            import torch.distributed as dist\n\n            # Wait for all processes to finish\n            dist.barrier()\n\n            # all gather the policykl\n            dist.all_reduce(policykl, dist.ReduceOp.SUM)\n            policykl /= self.accelerator.num_processes\n\n            if policykl > 1.5 * self.config.target_kl:\n                self.optimizer.zero_grad()\n                early_stop = True\n        return early_stop\n\n    def gather_stats(self, stats):\n        \"\"\"\n        Gather stats from all processes. Useful in the context of distributed training.\n\n        Args:\n            stats (dict[str, Any]):\n            a dictionary of stats to be gathered. The stats should contain torch tensors.\n\n        Returns:\n            `dict[str, Any]`: A dictionary of stats with the tensors gathered.\n        \"\"\"\n        import torch.distributed as dist\n\n        # Wait for all processes to finish\n        dist.barrier()\n\n        for k, v in stats.items():\n            if isinstance(v, torch.Tensor):\n                dist.all_reduce(v.to(self.accelerator.device), dist.ReduceOp.SUM)\n                v /= self.accelerator.num_processes\n            stats[k] = v\n        return stats\n\n    def prepare_model_inputs(self, queries: torch.Tensor, responses: torch.Tensor):\n        if self.is_encoder_decoder:\n            input_data = self.data_collator(\n                [{\"input_ids\": q, \"attention_mask\": torch.ones_like(q)} for q in queries]\n            ).to(self.current_device)\n\n            decoder_inputs = self.data_collator(\n                [{\"input_ids\": r, \"attention_mask\": torch.ones_like(r)} for r in responses]\n            ).to(self.current_device)\n\n            input_data[\"decoder_input_ids\"] = decoder_inputs[\"input_ids\"]\n            input_data[\"decoder_attention_mask\"] = decoder_inputs[\"attention_mask\"]\n        else:\n            input_ids = [torch.cat([q, r]) for q, r in zip(queries, responses)]\n            input_data = self.data_collator(\n                [{\"input_ids\": ids, \"attention_mask\": torch.ones_like(ids)} for ids in input_ids]\n            ).to(self.current_device)\n\n        input_data.pop(\"labels\", None)  # we don't want to compute LM losses\n        return input_data\n\n    @PPODecorators.empty_device_cache()\n    def batched_forward_pass(\n        self,\n        model: PreTrainedModelWrapper,\n        queries: torch.Tensor,\n        responses: torch.Tensor,\n        model_inputs: dict,\n        return_logits: bool = False,\n        response_masks: Optional[torch.Tensor] = None,\n    ):\n        \"\"\"\n        Calculate model outputs in multiple batches.\n\n        Args:\n            queries (`torch.LongTensor`):\n                List of tensors containing the encoded queries, shape (`batch_size`, `query_length`)\n            responses (`torch.LongTensor`):\n                List of tensors containing the encoded responses, shape (`batch_size`, `response_length`)\n            return_logits (`bool`, *optional*, defaults to `False`):\n                Whether to return all_logits. Set to `False` if logits are not needed to reduce memory consumption.\n\n        Returns:\n            (tuple):\n                - all_logprobs (`torch.FloatTensor`): Log probabilities of the responses,\n                    shape (`batch_size`, `response_length`)\n                - all_ref_logprobs (`torch.FloatTensor`): Log probabilities of the responses,\n                    shape (`batch_size`, `response_length`)\n                - all_values (`torch.FloatTensor`): Values of the responses, shape (`batch_size`, `response_length`)\n        \"\"\"\n        bs = len(queries)\n        fbs = self.config.mini_batch_size\n        all_logprobs = []\n        all_logits = []\n        all_masks = []\n        all_values = []\n\n        model.eval()\n\n        for i in range(math.ceil(bs / fbs)):\n            input_kwargs = {key: value[i * fbs : (i + 1) * fbs] for key, value in model_inputs.items()}\n            query_batch = queries[i * fbs : (i + 1) * fbs]\n            response_batch = responses[i * fbs : (i + 1) * fbs]\n            if response_masks is not None:\n                response_masks_batch = response_masks[i * fbs : (i + 1) * fbs]\n            logits, _, values = model(**input_kwargs)\n\n            if self.is_encoder_decoder:\n                input_ids = input_kwargs[\"decoder_input_ids\"]\n                attention_mask = input_kwargs[\"decoder_attention_mask\"]\n            else:\n                input_ids = input_kwargs[\"input_ids\"]\n                attention_mask = input_kwargs[\"attention_mask\"]\n\n            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])\n            masks = torch.zeros_like(attention_mask)\n            masks[:, :-1] = attention_mask[:, 1:]\n\n            for j in range(len(query_batch)):\n                if self.is_encoder_decoder:\n                    # Decoder sentence starts always in the index 1 after padding in the Enc-Dec Models\n                    start = 1\n                    end = attention_mask[j, :].sum() - 1\n                else:\n                    start = len(query_batch[j]) - 1  # logprobs starts from the second query token\n                    if attention_mask[j, 0] == 0:  # offset left padding\n                        start += attention_mask[j, :].nonzero()[0]\n                    end = start + len(response_batch[j])\n\n                masks[j, :start] = 0\n                masks[j, end:] = 0\n                if response_masks is not None:\n                    masks[j, start:end] = masks[j, start:end] * response_masks_batch[j]\n\n            if return_logits:\n                all_logits.append(logits)\n            else:\n                del logits\n            all_values.append(values)\n            all_logprobs.append(logprobs)\n            all_masks.append(masks)\n\n        return (\n            torch.cat(all_logprobs),\n            torch.cat(all_logits)[:, :-1] if return_logits else None,\n            torch.cat(all_values)[:, :-1],\n            torch.cat(all_masks)[:, :-1],\n        )\n\n    @PPODecorators.empty_device_cache()\n    def train_minibatch(\n        self,\n        old_logprobs: torch.FloatTensor,\n        values: torch.FloatTensor,\n        logprobs: torch.FloatTensor,\n        logits: torch.FloatTensor,\n        vpreds: torch.FloatTensor,\n        mask: torch.LongTensor,\n        advantages: torch.FloatTensor,\n        returns: torch.FloatTensor,\n    ):\n        \"\"\"\n        Train one PPO minibatch\n\n        Args:\n            logprobs (`torch.FloatTensor`):\n                Log probabilities of the model, shape [mini_batch_size, response_length]\n            values (`torch.FloatTensor`):\n                Values of the value head, shape [mini_batch_size, response_length]\n            query (`torch.LongTensor`):\n                Encoded queries, shape [mini_batch_size, query_length]\n            response (`torch.LongTensor`):\n                Encoded responses, shape [mini_batch_size, response_length]\n            model_input (`torch.LongTensor`):\n                Concatenated queries and responses, shape [mini_batch_size, query_length+response_length]\n\n        Returns:\n            train_stats (dict[str, `torch.Tensor`]):\n                Dictionary of training statistics\n        \"\"\"\n        self.model.train()\n        loss_p, loss_v, train_stats = self.loss(\n            old_logprobs, values, logits, vpreds, logprobs, mask, advantages, returns\n        )\n        loss = loss_p + loss_v\n        self.accelerator.backward(loss)\n        if self.config.max_grad_norm is not None:\n            if self.accelerator.sync_gradients:\n                self.accelerator.clip_grad_norm_(self.model_params, self.config.max_grad_norm)\n        self.optimizer.step()\n        # we call optimizer.zero_grad() every time and let `accelerator` handle accumulation\n        # see https://huggingface.co/docs/accelerate/usage_guides/gradient_accumulation#the-finished-code\n        self.optimizer.zero_grad()\n        return train_stats\n\n    def compute_rewards(\n        self,\n        scores: torch.FloatTensor,\n        logprobs: torch.FloatTensor,\n        ref_logprobs: torch.FloatTensor,\n        masks: torch.LongTensor,\n    ):\n        \"\"\"\n        Compute per token rewards from scores and KL-penalty.\n\n        Args:\n            scores (`torch.FloatTensor`):\n                Scores from the reward model, shape (`batch_size`)\n            logprobs (`torch.FloatTensor`):\n                Log probabilities of the model, shape (`batch_size`, `response_length`)\n            ref_logprobs (`torch.FloatTensor`):\n                Log probabilities of the reference model, shape (`batch_size`, `response_length`)\n\n        Returns:\n            `torch.FloatTensor`: Per token rewards, shape (`batch_size`, `response_length`)\n            `torch.FloatTensor`: Non score rewards, shape (`batch_size`, `response_length`)\n            `torch.FloatTensor`: KL penalty, shape (`batch_size`, `response_length`)\n        \"\"\"\n        rewards, non_score_rewards, kls = [], [], []\n        for score, logprob, ref_logprob, mask in zip(scores, logprobs, ref_logprobs, masks):\n            # compute KL penalty (from difference in logprobs)\n            kl = self._kl_penalty(logprob, ref_logprob)\n            kls.append(kl)\n            non_score_reward = -self.kl_ctl.value * kl\n            non_score_rewards.append(non_score_reward)\n            reward = non_score_reward.clone()\n            last_non_masked_index = mask.nonzero()[-1]\n\n            # reward is preference model score + KL penalty\n            reward[last_non_masked_index] += score\n            rewards.append(reward)\n        return torch.stack(rewards), torch.stack(non_score_rewards), torch.stack(kls)\n\n    def _kl_penalty(self, logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor) -> torch.FloatTensor:\n        if self.config.kl_penalty == \"kl\":\n            return logprob - ref_logprob\n\n        if self.config.kl_penalty == \"abs\":\n            return (logprob - ref_logprob).abs()\n\n        if self.config.kl_penalty == \"mse\":\n            return 0.5 * (logprob - ref_logprob).square()\n\n        if self.config.kl_penalty == \"full\":\n            # Flip is required due to this issue? :https://github.com/pytorch/pytorch/issues/57459\n            return F.kl_div(ref_logprob, logprob, log_target=True, reduction=\"none\").sum(-1)\n\n        raise NotImplementedError\n\n    def compute_advantages(\n        self,\n        values: torch.FloatTensor,\n        rewards: torch.FloatTensor,\n        mask: torch.FloatTensor,\n    ):\n        lastgaelam = 0\n        advantages_reversed = []\n        gen_len = rewards.shape[-1]\n\n        values = values * mask\n        rewards = rewards * mask\n\n        if self.config.whiten_rewards:\n            rewards = masked_whiten(rewards, mask, shift_mean=False)\n\n        for t in reversed(range(gen_len)):\n            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0\n            delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]\n            lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam\n            advantages_reversed.append(lastgaelam)\n        advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)\n\n        returns = advantages + values\n        advantages = masked_whiten(advantages, mask)\n        advantages = advantages.detach()\n        return values, advantages, returns\n\n    def loss(\n        self,\n        old_logprobs: torch.FloatTensor,\n        values: torch.FloatTensor,\n        logits: torch.FloatTensor,\n        vpreds: torch.FloatTensor,\n        logprobs: torch.FloatTensor,\n        mask: torch.LongTensor,\n        advantages: torch.FloatTensor,\n        returns: torch.FloatTensor,\n    ):\n        \"\"\"\n        Calculate policy and value losses.\n\n        Args:\n            old_logprobs (`torch.FloatTensor`):\n                Log probabilities of the model, shape (`batch_size`, `response_length`)\n            values (`torch.FloatTensor`):\n                Values of the value head, shape (`batch_size`, `response_length`)\n            rewards (`torch.FloatTensor`):\n                Rewards from the reward model, shape (`batch_size`, `response_length`)\n            logits (`torch.FloatTensor`):\n                Logits of the model, shape (`batch_size`, `response_length`, `vocab_size`)\n            v_pred (`torch.FloatTensor`):\n                Values of the value head, shape (`batch_size`, `response_length`)\n            logprobs (`torch.FloatTensor`):\n                Log probabilities of the model, shape (`batch_size`, `response_length`)\n        \"\"\"\n\n        vpredclipped = clip_by_value(\n            vpreds,\n            values - self.config.cliprange_value,\n            values + self.config.cliprange_value,\n        )\n\n        vf_losses1 = (vpreds - returns) ** 2\n        vf_losses2 = (vpredclipped - returns) ** 2\n        vf_loss = 0.5 * masked_mean(torch.max(vf_losses1, vf_losses2), mask)\n        vf_clipfrac = masked_mean(torch.gt(vf_losses2, vf_losses1).float(), mask)\n\n        ratio = torch.exp(logprobs - old_logprobs)\n\n        pg_losses = -advantages * ratio\n        pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - self.config.cliprange, 1.0 + self.config.cliprange)\n\n        pg_loss = masked_mean(torch.max(pg_losses, pg_losses2), mask)\n        pg_clipfrac = masked_mean(torch.gt(pg_losses2, pg_losses).float(), mask)\n\n        loss = pg_loss + self.config.vf_coef * vf_loss\n\n        avg_ratio = masked_mean(ratio, mask).item()\n        if avg_ratio > self.config.ratio_threshold:\n            warnings.warn(\n                f\"The average ratio of batch ({avg_ratio:.2f}) exceeds threshold {self.config.ratio_threshold:.2f}. Skipping batch.\"\n            )\n            pg_loss = pg_loss * 0.0\n            vf_loss = vf_loss * 0.0\n            loss = loss * 0.0\n\n        entropy = masked_mean(entropy_from_logits(logits), mask)\n\n        approxkl = 0.5 * masked_mean((logprobs - old_logprobs) ** 2, mask)\n        policykl = masked_mean(old_logprobs - logprobs, mask)\n\n        return_mean, return_var = masked_mean(returns, mask), masked_var(returns, mask)\n        value_mean, value_var = masked_mean(values, mask), masked_var(values, mask)\n\n        stats = dict(\n            loss=dict(policy=pg_loss.detach(), value=vf_loss.detach(), total=loss.detach()),\n            policy=dict(\n                entropy=entropy.detach(),\n                approxkl=approxkl.detach(),\n                policykl=policykl.detach(),\n                clipfrac=pg_clipfrac.detach(),\n                advantages=advantages.detach(),\n                advantages_mean=masked_mean(advantages, mask).detach(),\n                ratio=ratio.detach(),\n            ),\n            returns=dict(mean=return_mean.detach(), var=return_var.detach()),\n            val=dict(\n                vpred=masked_mean(vpreds, mask).detach(),\n                error=masked_mean((vpreds - returns) ** 2, mask).detach(),\n                clipfrac=vf_clipfrac.detach(),\n                mean=value_mean.detach(),\n                var=value_var.detach(),\n            ),\n        )\n        return pg_loss, self.config.vf_coef * vf_loss, flatten_dict(stats)\n\n    def record_step_stats(self, kl_coef: float, **data):\n        \"\"\"\n        Record training step statistics.\n\n\n        Args:\n            kl_coef (`float`):\n                KL coefficient\n            data (`dict`):\n                Dictionary of training step data\n\n        Returns:\n            stats (`dict`):\n                Dictionary of training step statistics\n        \"\"\"\n        mask = data.pop(\"masks\")\n\n        kls = data.pop(\"kls\")\n        kl_list = ((kls) * mask).sum(axis=-1)\n        mean_kl = kl_list.mean()\n        mean_entropy = (-data[\"logprobs\"] * mask).sum(axis=-1).mean()\n\n        mean_non_score_reward = masked_mean(\n            data[\"non_score_reward\"], mask\n        )  # non_score_reward is size `batch_size`, `response_length`\n        mean_scores = data[\"scores\"].mean()  # scores is size `batch_size`\n        std_scores = data[\"scores\"].std()\n\n        if mean_kl.item() < -1.0:\n            # warn users\n            warnings.warn(\n                f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n                \" sometimes this happens because the generation kwargs are not correctly set. Please make sure\"\n                \" that the generation kwargs are set correctly, or review your training hyperparameters.\"\n            )\n\n        stats = {\n            \"objective/kl\": mean_kl,\n            \"objective/kl_dist\": kl_list,\n            \"objective/logprobs\": data[\"logprobs\"],\n            \"objective/ref_logprobs\": data[\"ref_logprobs\"],\n            \"objective/kl_coef\": kl_coef,\n            \"objective/entropy\": mean_entropy,\n            \"ppo/mean_non_score_reward\": mean_non_score_reward,\n            \"ppo/mean_scores\": mean_scores,\n            \"ppo/std_scores\": std_scores,\n        }\n\n        # Log text properties\n        query_lens = torch.tensor([len(query) for query in data[\"queries\"]], dtype=torch.float)\n        response_lens = torch.tensor([len(response) for response in data[\"responses\"]], dtype=torch.float)\n\n        stats[\"tokens/queries_len_mean\"] = torch.mean(query_lens).cpu().numpy().item()\n        stats[\"tokens/queries_len_std\"] = torch.std(query_lens).cpu().numpy().item()\n        stats[\"tokens/queries_dist\"] = query_lens.cpu().numpy()\n        stats[\"tokens/responses_len_mean\"] = torch.mean(response_lens).cpu().numpy().item()\n        stats[\"tokens/responses_len_std\"] = torch.std(response_lens).cpu().numpy().item()\n        stats[\"tokens/responses_dist\"] = response_lens.cpu().numpy()\n\n        for k, v in data[\"train_stats\"].items():\n            stats[f\"ppo/{k}\"] = torch.mean(v, axis=0)\n        stats[\"ppo/val/var_explained\"] = 1 - stats[\"ppo/val/error\"] / stats[\"ppo/returns/var\"]\n        return stats\n\n    def log_stats(\n        self,\n        stats: dict,\n        batch: dict,\n        rewards: List[torch.FloatTensor],\n        columns_to_log: typing.Iterable[str] = (\"query\", \"response\"),\n    ):\n        \"\"\"\n        A function that logs all the training stats. Call it at the end of each epoch.\n\n        Args:\n            stats (dict[str, Any]):\n                A dictionary of training stats.\n            batch (dict[str, Any]):\n                A dictionary of batch data, this contains the queries and responses.\n            rewards (`List[torch.FloatTensor]`):\n                A tensor of rewards.\n        \"\"\"\n\n        # all gather stats\n        if not isinstance(rewards, torch.Tensor):\n            rewards = torch.tensor(rewards).to(self.current_device)\n        rewards = self.accelerator.gather(rewards).flatten()\n\n        if self.config.log_with == \"wandb\":\n            import wandb\n\n            if any(column_to_log not in batch.keys() for column_to_log in columns_to_log):\n                raise ValueError(f\"Columns to log {columns_to_log} are not present in the batch {batch.keys()}.\")\n\n            batch_list = [batch[column_to_log] for column_to_log in columns_to_log]\n            if self.is_distributed:\n                gathered_batch_list = []\n                for b in batch_list:\n                    flattened = gather_object(b)\n                    gathered_batch_list.append(flattened)\n                batch_list = gathered_batch_list\n\n        # Log only if we are in the main process\n        if self.accelerator.is_main_process:\n            logs = {}\n\n            # Log stats\n            if \"query\" not in batch.keys() and \"response\" not in batch.keys():\n                # warn the user that the game logs will not be logged\n                warnings.warn(\n                    \"The game logs will not be logged because the batch does not contain the keys 'query' and \"\n                    \"'response'. \"\n                )\n            elif self.config.log_with == \"wandb\":\n                table_rows = [list(r) for r in zip(*batch_list, rewards.cpu().tolist())]\n                logs.update({\"game_log\": wandb.Table(columns=[*columns_to_log, \"reward\"], rows=table_rows)})\n\n            logs.update(stats)\n\n            # manually cast in fp32 for bf16 torch tensors\n            for k, v in logs.items():\n                if isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16:\n                    logs[k] = v.float()\n\n            logs[\"env/reward_mean\"] = torch.mean(rewards).cpu().numpy().item()\n            logs[\"env/reward_std\"] = torch.std(rewards).cpu().numpy().item()\n            logs[\"env/reward_dist\"] = rewards.cpu().numpy()\n\n            if self.config.log_with == \"tensorboard\":\n                # update the current step\n                self.current_step += 1\n\n            self.accelerator.log(\n                logs,\n                step=self.current_step if self.config.log_with == \"tensorboard\" else None,\n            )\n\n    def create_model_card(self, path: str, model_name: Optional[str] = \"TRL Model\") -> None:\n        \"\"\"Creates and saves a model card for a TRL model.\n\n        Args:\n            path (`str`): The path to save the model card to.\n            model_name (`str`, *optional*): The name of the model, defaults to `TRL Model`.\n        \"\"\"\n        try:\n            user = whoami()[\"name\"]\n        # handle the offline case\n        except Exception:\n            warnings.warn(\"Cannot retrieve user information assuming you are running in offline mode.\")\n            return\n\n        if not os.path.exists(path):\n            os.makedirs(path)\n\n        model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, model_id=f\"{user}/{path}\")\n        with open(os.path.join(path, \"README.md\"), \"w\", encoding=\"utf-8\") as f:\n            f.write(model_card_content)\n\n    def _save_pretrained(self, save_directory: str) -> None:\n        self.accelerator.unwrap_model(self.model).save_pretrained(save_directory)\n        self.tokenizer.save_pretrained(save_directory)\n        self.create_model_card(save_directory)\n\n    def _show_tokens(self, tokens, masks):\n        from rich import print\n        from rich.text import Text\n\n        text = Text()\n\n        for _i, (token, mask) in enumerate(zip(tokens, masks)):\n            if mask == 1:\n                text.append(self.tokenizer.decode(token.item()), style=\"black on deep_sky_blue1\")\n                text.append(\" \")\n            else:\n                text.append(self.tokenizer.decode(token.item()), style=\"black on cyan3\")\n                text.append(\" \")\n        print(text)\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepspeed_plugin.deepspeed_config\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport gc\nimport math\nimport os\nimport time\nfrom collections import defaultdict\nfrom functools import wraps\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import Accelerator\nfrom accelerate.utils import broadcast, gather_object\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    DataCollatorWithPadding,\n    GenerationConfig,\n    PreTrainedTokenizer,\n    Trainer,\n    TrainerCallback,\n    TrainerControl,\n)\nfrom transformers.integrations import get_reporting_integration_callbacks\nfrom transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK\nfrom transformers.trainer_callback import CallbackHandler, PrinterCallback\n\nfrom ..models.utils import unwrap_model_for_generation\nfrom ..trainer.utils import (\n    OnlineTrainerState,\n    batch_generation,\n    disable_dropout_in_model,\n    exact_div,\n    first_true_indices,\n    forward,\n    get_reward,\n    prepare_deepspeed,\n    print_rich_table,\n    truncate_response,\n)\nfrom .rloo_config import RLOOConfig\nfrom .utils import trl_sanitze_kwargs_for_tagging\n\n\nINVALID_LOGPROB = 1.0\n\n\nclass RLOOTrainer(Trainer):\n    _tag_names = [\"trl\", \"rloo\"]\n\n    def __init__(\n        self,\n        config: RLOOConfig,\n        tokenizer: PreTrainedTokenizer,\n        policy: nn.Module,\n        ref_policy: nn.Module,\n        reward_model: nn.Module,\n        train_dataset: Dataset,\n        data_collator: Optional[DataCollatorWithPadding] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        # less commonly used\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        callbacks: Optional[List[TrainerCallback]] = None,\n    ) -> None:\n        if ref_policy is policy:\n            raise ValueError(\n                \"`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the \"\n                \"same as `policy`, you must mass a copy of it, or `None` if you use peft.\"\n            )\n\n        self.args = config\n        args = config\n        self.tokenizer = tokenizer\n        self.policy = policy\n\n        self.policy.generation_config.eos_token_id = (\n            None  # disable `pad_token_id` and `eos_token_id` because we just want to\n        )\n        self.policy.generation_config.pad_token_id = None  # generate tokens without truncation / padding\n\n        self.ref_policy = ref_policy\n        self.reward_model = reward_model\n        self.train_dataset = train_dataset\n        self.train_dataset_len = len(train_dataset)\n        self.data_collator = data_collator\n        self.eval_dataset = eval_dataset\n        self.optimizer, self.lr_scheduler = optimizers\n\n        #########\n        # calculate various batch sizes\n        #########\n        if args.total_episodes is None:  # allow the users to define episodes in terms of epochs.\n            args.total_episodes = int(args.num_train_epochs * self.train_dataset_len)\n        accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)\n        self.accelerator = accelerator\n        args.world_size = accelerator.num_processes\n        args.local_batch_size = (\n            args.per_device_train_batch_size * args.gradient_accumulation_steps * args.num_mini_batches\n        )\n        args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size)\n        args.batch_size = int(args.local_batch_size * args.world_size)\n        args.mini_batch_size = exact_div(\n            args.batch_size, args.num_mini_batches, \"`batch_size` must be a multiple of `num_mini_batches`\"\n        )\n        args.local_mini_batch_size = exact_div(\n            args.local_batch_size, args.num_mini_batches, \"`local_batch_size` must be a multiple of `num_mini_batches`\"\n        )\n        args.num_total_batches = math.ceil(\n            args.total_episodes / args.batch_size\n        )  # we may train for more than `total_episodes`\n        time_tensor = torch.tensor(int(time.time()), device=accelerator.device)\n        time_int = broadcast(time_tensor, 0).item()  # avoid different timestamps across processes\n        args.run_name = f\"{args.exp_name}__{args.seed}__{time_int}\"\n        self.local_seed = args.seed + accelerator.process_index * 100003  # Prime\n        if args.num_sample_generations > 0:\n            self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations)\n        self.local_dataloader_batch_size = exact_div(\n            args.local_batch_size, args.rloo_k, \"`local_batch_size` must be a multiple of rloo_k\"\n        )  # RLOO logic: needed because RLOO repeats the same prompt args.rloo_k times\n\n        #########\n        # setup model, optimizer, and others\n        #########\n        for module in [policy, ref_policy, reward_model]:\n            disable_dropout_in_model(module)\n        if args.stop_token and args.stop_token == \"eos\":\n            args.stop_token_id = tokenizer.eos_token_id\n        self.model = policy\n        self.create_optimizer_and_scheduler(\n            num_training_steps=args.num_total_batches\n        )  # note that we are calling `self.lr_scheduler.step()` manually only at the batch level\n\n        #########\n        ### trainer specifics\n        #########\n        self.state = OnlineTrainerState(\n            is_local_process_zero=self.is_local_process_zero(),\n            is_world_process_zero=self.is_world_process_zero(),\n        )\n        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)\n        self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks\n        self.callback_handler = CallbackHandler(\n            self.callbacks, self.model, self.tokenizer, self.optimizer, self.lr_scheduler\n        )\n        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)\n        self.control = TrainerControl()\n        self.current_flos = 0\n        self.hp_search_backend = None\n        self.is_deepspeed_enabled = getattr(self.accelerator.state, \"deepspeed_plugin\", None) is not None\n        self.is_fsdp_enabled = getattr(self.accelerator.state, \"fsdp_plugin\", None) is not None\n        # Create distant repo and output directory if needed\n        self.hub_model_id = None\n        if self.args.push_to_hub:\n            self.init_hf_repo()\n        if self.args.should_save:\n            os.makedirs(self.args.output_dir, exist_ok=True)\n        self.backup_model = None\n\n        #########\n        ### setup dataloader\n        #########\n        self.dataloader = DataLoader(\n            self.train_dataset,\n            batch_size=self.local_dataloader_batch_size,\n            shuffle=True,\n            collate_fn=DataCollatorWithPadding(tokenizer),\n            drop_last=True,  # needed; otherwise the last batch will be of ragged shape\n        )\n        # sync random states for DataLoader(shuffle=True) before `accelerator.prepare`\n        # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c\n        torch.manual_seed(args.seed)\n        self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader)\n        torch.manual_seed(self.local_seed)  # reset the local seed again\n\n        self.eval_dataloader = DataLoader(\n            self.eval_dataset,\n            batch_size=args.per_device_eval_batch_size,\n            collate_fn=DataCollatorWithPadding(self.tokenizer),\n            drop_last=True,\n        )  # no need to shuffle eval dataset\n        self.eval_dataloader = accelerator.prepare(self.eval_dataloader)\n\n        if self.is_deepspeed_enabled:\n            self.reward_model = prepare_deepspeed(\n                self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16\n            )\n            self.ref_policy = prepare_deepspeed(\n                self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16\n            )\n            self.deepspeed = self.model\n        else:\n            self.ref_policy = self.ref_policy.to(self.accelerator.device)\n            self.reward_model = self.reward_model.to(self.accelerator.device)\n\n    def get_train_dataloader(self) -> DataLoader:\n        return self.dataloader\n\n    def get_eval_dataloader(self) -> DataLoader:\n        return self.eval_dataloader\n\n    def train(self):\n        args = self.args\n        accelerator = self.accelerator\n        optimizer = self.optimizer\n        model = self.model\n        self.model_wrapped = self.model\n        ref_policy = self.ref_policy\n        reward_model = self.reward_model\n        tokenizer = self.tokenizer\n        dataloader = self.dataloader\n        device = accelerator.device\n\n        def repeat_generator():\n            while True:\n                yield from dataloader\n\n        iter_dataloader = iter(repeat_generator())\n        generation_config = GenerationConfig(\n            max_new_tokens=args.response_length,\n            temperature=(args.temperature + 1e-7),\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n        )\n\n        accelerator.print(\"===training policy===\")\n        start_time = time.time()\n        stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps)\n        approxkl_stats = torch.zeros(stats_shape, device=device)\n        pg_clipfrac_stats = torch.zeros(stats_shape, device=device)\n        pg_loss_stats = torch.zeros(stats_shape, device=device)\n        vf_loss_stats = torch.zeros(stats_shape, device=device)\n        vf_clipfrac_stats = torch.zeros(stats_shape, device=device)\n        entropy_stats = torch.zeros(stats_shape, device=device)\n        ratio_stats = torch.zeros(stats_shape, device=device)\n        model.train()\n\n        # trainer state initialization\n        self.state.global_step = 0\n        self.state.episode = 0\n        self.state.max_steps = args.num_total_batches * args.num_mini_batches\n        self.state.num_train_epochs = args.total_episodes / self.train_dataset_len\n        # Compute absolute values for logging, eval, and save if given as ratio\n        if args.logging_steps is not None:\n            if args.logging_steps < 1:\n                self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps)\n            else:\n                self.state.logging_steps = args.logging_steps\n        if args.eval_steps is not None:\n            if args.eval_steps < 1:\n                self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps)\n            else:\n                self.state.eval_steps = args.eval_steps\n        if args.save_steps is not None:\n            if args.save_steps < 1:\n                self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps)\n            else:\n                self.state.save_steps = args.save_steps\n        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)\n\n        for update in range(1, args.num_total_batches + 1):\n            self.state.episode += 1 * args.batch_size\n            data = next(iter_dataloader)\n            with torch.no_grad():\n                queries = data[\"input_ids\"].to(device)\n                queries = queries.repeat(args.rloo_k, 1)\n                context_length = queries.shape[1]\n                query_responses = []\n                responses = []\n                postprocessed_responses = []\n                logprobs = []\n                ref_logprobs = []\n                scores = []\n                sequence_lengths = []\n                with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n                    query_responses, logitss = batch_generation(\n                        unwrapped_model,\n                        queries,\n                        args.local_rollout_forward_batch_size,\n                        tokenizer.pad_token_id,\n                        generation_config,\n                    )\n\n                for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size):\n                    query = queries[i : i + args.local_rollout_forward_batch_size]\n                    query_response = query_responses[i : i + args.local_rollout_forward_batch_size]\n                    response = query_response[:, context_length:]\n                    logits = logitss[i : i + args.local_rollout_forward_batch_size]\n                    all_logprob = F.log_softmax(logits, dim=-1)\n                    logprob = torch.gather(all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)\n                    del logits, all_logprob\n                    torch.cuda.empty_cache()\n\n                    ref_output = forward(ref_policy, query_response, tokenizer.pad_token_id)\n                    ref_logits = ref_output.logits[:, context_length - 1 : -1]\n                    ref_logits /= args.temperature + 1e-7\n                    ref_all_logprob = F.log_softmax(ref_logits, dim=-1)\n                    ref_logprob = torch.gather(ref_all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)\n                    del ref_output, ref_logits, ref_all_logprob\n                    torch.cuda.empty_cache()\n\n                    # Response Processing 1. truncate response after the first occurrence of `stop_token_id`\n                    postprocessed_response = response\n                    if args.stop_token_id is not None:  # handle the edge case when stop_token_id exists but is 0\n                        postprocessed_response = truncate_response(\n                            args.stop_token_id, tokenizer.pad_token_id, response\n                        )\n\n                    # Response Processing 2. run reward model on the truncated responses\n                    postprocessed_query_response = torch.cat((query, postprocessed_response), 1)\n                    sequence_length = first_true_indices(postprocessed_response == tokenizer.pad_token_id) - 1\n                    _, score, _ = get_reward(\n                        reward_model, postprocessed_query_response, tokenizer.pad_token_id, context_length\n                    )\n\n                    responses.append(response)\n                    postprocessed_responses.append(postprocessed_response)\n                    logprobs.append(logprob)\n                    ref_logprobs.append(ref_logprob)\n                    sequence_lengths.append(sequence_length)\n                    scores.append(score)\n                responses = torch.cat(responses, 0)\n                postprocessed_responses = torch.cat(postprocessed_responses, 0)\n                logprobs = torch.cat(logprobs, 0)\n                ref_logprobs = torch.cat(ref_logprobs, 0)\n                sequence_lengths = torch.cat(sequence_lengths, 0)\n                scores = torch.cat(scores, 0)\n                del (logprob, ref_logprob, score)\n                torch.cuda.empty_cache()\n                gc.collect()\n\n                # Response Processing 3. filter response. Ensure that the sample contains stop_token_id\n                # responses not passing that filter will receive a low (fixed) score\n                # only query humans on responses that pass that filter\n                contain_eos_token = torch.any(postprocessed_responses == tokenizer.eos_token_id, dim=-1)\n                if args.missing_eos_penalty is not None:\n                    scores[~contain_eos_token] -= self.args.missing_eos_penalty\n                # accelerator.print(f\"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}\")\n\n                # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw\n                response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1)\n                padding_mask = response_idxs > sequence_lengths.unsqueeze(1)\n                logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB)\n                ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB)\n\n                # 4. compute rewards\n                kl = logprobs - ref_logprobs\n                non_score_reward = (-args.kl_coef * kl).sum(1)\n                rlhf_reward = scores + non_score_reward\n\n                # vectorized RLOO advantages implementation\n                rlhf_reward = rlhf_reward.reshape(args.rloo_k, -1)\n                baseline = (rlhf_reward.sum(0) - rlhf_reward) / (args.rloo_k - 1)\n                advantages = rlhf_reward - baseline\n                advantages = advantages.flatten()\n                torch.cuda.empty_cache()\n\n            # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch\n            for ppo_epoch_idx in range(args.num_ppo_epochs):\n                b_inds = np.random.permutation(args.local_batch_size)\n                minibatch_idx = 0\n                for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size):\n                    mini_batch_end = mini_batch_start + args.local_mini_batch_size\n                    mini_batch_inds = b_inds[mini_batch_start:mini_batch_end]\n                    gradient_accumulation_idx = 0\n                    for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size):\n                        with accelerator.accumulate(model):\n                            micro_batch_end = micro_batch_start + args.per_device_train_batch_size\n                            micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end]\n                            mb_advantage = advantages[micro_batch_inds]\n                            mb_responses = responses[micro_batch_inds]\n                            mb_query_responses = query_responses[micro_batch_inds]\n                            mb_logprobs = logprobs[micro_batch_inds]\n\n                            output = forward(model, mb_query_responses, tokenizer.pad_token_id)\n                            logits = output.logits[:, context_length - 1 : -1]\n                            logits /= args.temperature + 1e-7\n                            new_all_logprobs = F.log_softmax(logits, dim=-1)\n                            new_logprobs = torch.gather(new_all_logprobs, 2, mb_responses.unsqueeze(-1)).squeeze(-1)\n                            new_logprobs = torch.masked_fill(\n                                new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB\n                            )\n                            new_ratio = (new_logprobs - mb_logprobs).exp()\n                            new_logprobs = new_logprobs.sum(1)\n                            mb_logprobs = mb_logprobs.sum(1)\n                            logprobs_diff = new_logprobs - mb_logprobs\n                            ratio = torch.exp(logprobs_diff)\n                            pg_losses = -mb_advantage * ratio\n                            pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange)\n                            pg_loss_max = torch.max(pg_losses, pg_losses2)\n                            pg_loss = pg_loss_max.mean()\n                            loss = pg_loss\n                            accelerator.backward(loss)\n                            optimizer.step()\n                            optimizer.zero_grad()\n                            with torch.no_grad():\n                                pg_clipfrac = (pg_losses2 > pg_losses).float().mean()\n                                prob_dist = torch.nn.functional.softmax(logits, dim=-1)\n                                entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1)\n                                approxkl = 0.5 * (logprobs_diff**2).mean()\n                                approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl\n                                pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = (\n                                    pg_clipfrac\n                                )\n                                pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss\n                                entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean()\n                                ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = new_ratio.mean()\n                        gradient_accumulation_idx += 1\n                    minibatch_idx += 1\n                    self.state.global_step += 1\n                    # del everything and empty cache\n                    # fmt: off\n                    del (\n                        output, logits, new_all_logprobs, new_logprobs,\n                        logprobs_diff, ratio, pg_losses, pg_losses2,\n                        pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl,\n                        mb_advantage, mb_responses, mb_query_responses, mb_logprobs,\n                    )\n                    # fmt: on\n                    torch.cuda.empty_cache()\n            with torch.no_grad():\n                mean_kl = kl.sum(1).mean()\n                mean_entropy = (-logprobs).sum(1).mean()\n                mean_non_score_reward = non_score_reward.mean()\n                eps = int(self.state.episode / (time.time() - start_time))\n                metrics = {}\n                metrics[\"eps\"] = eps\n                metrics[\"objective/kl\"] = self.accelerator.gather(mean_kl).mean().item()\n                metrics[\"objective/entropy\"] = self.accelerator.gather(mean_entropy).mean().item()\n                metrics[\"objective/non_score_reward\"] = self.accelerator.gather(mean_non_score_reward).mean().item()\n                metrics[\"objective/rlhf_reward\"] = self.accelerator.gather(rlhf_reward).mean().item()\n                metrics[\"objective/scores\"] = self.accelerator.gather(scores.mean()).mean().item()\n                metrics[\"policy/approxkl_avg\"] = self.accelerator.gather(approxkl_stats).mean().item()\n                metrics[\"policy/clipfrac_avg\"] = self.accelerator.gather(pg_clipfrac_stats).mean().item()\n                metrics[\"loss/policy_avg\"] = self.accelerator.gather(pg_loss_stats).mean().item()\n                metrics[\"loss/value_avg\"] = self.accelerator.gather(vf_loss_stats).mean().item()\n                metrics[\"val/clipfrac_avg\"] = self.accelerator.gather(vf_clipfrac_stats).mean().item()\n                metrics[\"policy/entropy_avg\"] = self.accelerator.gather(entropy_stats).mean().item()\n                metrics[\"val/ratio\"] = self.accelerator.gather(ratio_stats).mean().item()\n                metrics[\"val/ratio_var\"] = self.accelerator.gather(ratio_stats).var().item()\n                metrics[\"val/num_eos_tokens\"] = (responses == tokenizer.eos_token_id).sum().item()\n                metrics[\"lr\"] = self.lr_scheduler.get_last_lr()[0]\n                metrics[\"episode\"] = self.state.episode\n                self.state.epoch = self.state.episode / self.train_dataset_len  # used by self.log\n                self.state.global_step += 1\n                self.log(metrics)\n            del kl, mean_kl, mean_entropy, scores\n\n            self.lr_scheduler.step()\n            self.control = self.callback_handler.on_step_end(args, self.state, self.control)\n            if self.control.should_save:\n                self._save_checkpoint(model, trial=None, metrics=metrics)\n                self.control = self.callback_handler.on_save(self.args, self.state, self.control)\n            torch.cuda.empty_cache()\n            gc.collect()\n\n            if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0:\n                self.generate_completions(sampling=True)\n\n        # HF trainer specifics\n        self.control = self.callback_handler.on_train_end(args, self.state, self.control)\n        if self.control.should_save:\n            self._save_checkpoint(model, trial=None, metrics=None)\n            self.control = self.callback_handler.on_save(self.args, self.state, self.control)\n\n    def generate_completions(self, sampling: bool = False):\n        args = self.args\n        tokenizer = self.tokenizer\n        generation_config = GenerationConfig(\n            max_new_tokens=self.args.response_length,\n            temperature=(0.01 + 1e-7),\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n        )\n\n        table = defaultdict(list)\n        with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:\n            for batch in self.eval_dataloader:\n                query = batch[\"input_ids\"]\n                with torch.no_grad():\n                    context_length = query.shape[1]\n                    query_response, _ = batch_generation(\n                        unwrapped_model,\n                        query,\n                        query.shape[0],\n                        tokenizer.pad_token_id,\n                        generation_config,\n                    )\n                    response = query_response[:, context_length:]\n                    postprocessed_response = response\n                    if args.stop_token_id is not None:  # handle the edge case when stop_token_id exists but is 0\n                        postprocessed_response = truncate_response(\n                            args.stop_token_id, tokenizer.pad_token_id, response\n                        )\n                    table[\"query\"].extend(gather_object(tokenizer.batch_decode(query, skip_special_tokens=True)))\n                    table[\"model response\"].extend(gather_object(tokenizer.batch_decode(postprocessed_response)))\n\n                    postprocessed_query_response = torch.cat((query, postprocessed_response), 1)\n                    _, score, _ = get_reward(\n                        self.reward_model, postprocessed_query_response, tokenizer.pad_token_id, context_length\n                    )\n                    table[\"score\"].extend(self.accelerator.gather(score).float().cpu().numpy())\n\n                if sampling:\n                    break\n        df = pd.DataFrame(table)\n\n        if self.accelerator.is_main_process:\n            print_rich_table(df.iloc[0 : 0 + 5])\n            if \"wandb\" in args.report_to:\n                import wandb\n\n                if wandb.run is not None:\n                    wandb.log({\"completions\": wandb.Table(dataframe=df)})\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"rloo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport random\nimport warnings\nfrom copy import deepcopy\nfrom typing import Any, Dict, Optional, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate.utils import is_deepspeed_available\nfrom transformers import AutoModelForCausalLM, GenerationConfig, PreTrainedModel\n\nfrom ..import_utils import is_liger_kernel_available\nfrom ..models import PreTrainedModelWrapper\nfrom ..models.utils import unwrap_model_for_generation\nfrom .gkd_config import GKDConfig\nfrom .sft_trainer import SFTTrainer\nfrom .utils import DataCollatorForChatML, disable_dropout_in_model, empty_cache\n\n\nif is_deepspeed_available():\n    import deepspeed\n\nif is_liger_kernel_available():\n    from liger_kernel.transformers import AutoLigerKernelForCausalLM\n\n\nclass GKDTrainer(SFTTrainer):\n    _tag_names = [\"trl\", \"gkd\"]\n\n    def __init__(\n        self,\n        teacher_model: Union[PreTrainedModel, nn.Module, str],\n        args: Optional[GKDConfig] = None,\n        *sft_args,\n        **kwargs,\n    ):\n        # add remove_unused_columns=False to the the dataclass args\n        args.remove_unused_columns = False\n        kwargs[\"data_collator\"] = DataCollatorForChatML(tokenizer=kwargs[\"tokenizer\"], max_length=args.max_seq_length)\n\n        super().__init__(*sft_args, args=args, **kwargs)\n\n        if args.teacher_model_init_kwargs is None:\n            teacher_model_init_kwargs = {}\n        elif not isinstance(teacher_model, str):\n            raise ValueError(\n                \"You passed teacher_model_init_kwargs to the GKDConfig, but your teacher_model is already instantiated.\"\n            )\n        else:\n            teacher_model_init_kwargs = args.teacher_model_init_kwargs\n            teacher_model_init_kwargs[\"torch_dtype\"] = (\n                teacher_model_init_kwargs[\"torch_dtype\"]\n                if teacher_model_init_kwargs[\"torch_dtype\"] in [\"auto\", None]\n                else getattr(torch, teacher_model_init_kwargs[\"torch_dtype\"])\n            )\n\n        if isinstance(teacher_model, str):\n            warnings.warn(\n                \"You passed a teacher model_id to the GKDTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM`\"\n            )\n            if args.use_liger:\n                teacher_model = AutoLigerKernelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs)\n            else:\n                teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs)\n\n        if args.disable_dropout:\n            disable_dropout_in_model(self.model)\n\n        if self.is_deepspeed_enabled:\n            self.teacher_model = self._prepare_deepspeed(teacher_model)\n        else:\n            self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True)\n\n        self.lmbda = args.lmbda\n        self.beta = args.beta\n        self.temperature = args.temperature\n\n        self.generation_config = GenerationConfig(\n            max_new_tokens=args.max_new_tokens,\n            temperature=args.temperature,\n            do_sample=True,\n            top_k=0,\n            use_cache=False if args.gradient_checkpointing else True,\n        )\n\n    @staticmethod\n    def generalized_jsd_loss(\n        student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction=\"batchmean\"\n    ):\n        \"\"\"\n        Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) of https://arxiv.org/abs/2306.13649 for the definition.\n\n        Args:\n            student_logits: Tensor of shape (batch_size, sequence_length, vocab_size)\n            teacher_logits: Tensor of shape (batch_size, sequence_length, vocab_size)\n            labels: Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing loss\n            beta: Interpolation coefficient between 0 and 1 (default: 0.5)\n            temperature: Softmax temperature (default: 1.0)\n            reduction: Specifies the reduction to apply to the output (default: 'batchmean')\n\n        Returns:\n            loss: Scalar tensor with the generalized JSD loss\n        \"\"\"\n\n        # Apply temperature scaling\n        student_logits = student_logits / temperature\n        teacher_logits = teacher_logits / temperature\n\n        # Compute log probabilities for student and probabilities for teacher\n        student_log_probs = F.log_softmax(student_logits, dim=-1)\n        teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)\n\n        # Compute the interpolated log probabilities\n        interpolated_log_probs = beta * student_log_probs + (1 - beta) * teacher_log_probs\n\n        # Compute KL divergences using F.kl_div\n        # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper.\n        kl_teacher = F.kl_div(interpolated_log_probs, teacher_log_probs, reduction=\"none\", log_target=True)\n        kl_student = F.kl_div(interpolated_log_probs, student_log_probs, reduction=\"none\", log_target=True)\n\n        # Compute the Generalized Jensen-Shannon Divergence\n        jsd = beta * kl_teacher + (1 - beta) * kl_student\n\n        # Masking\n        if labels is not None:\n            mask = labels != -100\n            jsd = jsd[mask]\n\n        # Apply reduction\n        if reduction == \"batchmean\":\n            return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / (jsd.size(0) * jsd.size(1))\n        elif reduction == \"sum\":\n            return jsd.sum()\n        elif reduction == \"mean\":\n            return jsd.mean()\n        else:\n            return jsd\n\n    def compute_loss(self, model, inputs, return_outputs=False):\n        # compute student output\n        outputs_student = model(\n            input_ids=inputs[\"input_ids\"],\n            attention_mask=inputs[\"attention_mask\"],\n        )\n\n        # compute teacher output in eval mode\n        self.teacher_model.eval()\n        with torch.no_grad():\n            outputs_teacher = self.teacher_model(\n                input_ids=inputs[\"input_ids\"],\n                attention_mask=inputs[\"attention_mask\"],\n            )\n\n        # slice the logits for the generated tokens using the inputs[\"prompts\"] lengths\n        prompt_lengths = inputs[\"prompts\"].shape[1]\n        shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :]\n        shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :]\n        shifted_labels = inputs[\"labels\"][:, prompt_lengths:]\n\n        # compute loss\n        loss = self.generalized_jsd_loss(\n            student_logits=shifted_student_logits,\n            teacher_logits=shifted_teacher_logits,\n            labels=shifted_labels,\n            beta=self.beta,\n        )\n\n        # empty cache\n        empty_cache()\n\n        # Return loss\n        return (loss, outputs_student) if return_outputs else loss\n\n    @staticmethod\n    def generate_on_policy_outputs(model, inputs, generation_config, pad_token_id=None):\n        # Generate output with respect to the prompt only\n        generated_outputs = model.generate(\n            input_ids=inputs[\"prompts\"],\n            attention_mask=inputs.get(\"prompt_attention_mask\", None),\n            generation_config=generation_config,\n            return_dict_in_generate=True,\n        )\n\n        # Get the generated token IDs\n        generated_tokens = generated_outputs.sequences\n        # Calculate new attention mask\n        new_attention_mask = torch.ones_like(generated_tokens)\n        new_labels = generated_tokens.clone()\n\n        # If there's pad_token_id, set attention mask to 0 for padding tokens\n        if pad_token_id is not None:\n            new_labels[new_labels == pad_token_id] = -100\n            new_attention_mask[generated_tokens == pad_token_id] = 0\n\n        return generated_tokens, new_attention_mask, new_labels\n\n    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:\n        \"\"\"\n        Perform a training step for the Generalized Knowledge Distillation (GKD) model.\n\n        This method implements the on-policy learning approach described in the GKD paper.\n        With probability `self.lmbda`, it generates new responses using the student model,\n        which are then used for training instead of the original inputs.\n        \"\"\"\n        if random.random() <= self.lmbda:\n            with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n                new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs(\n                    unwrapped_model, inputs, self.generation_config, self.tokenizer.pad_token_id\n                )\n            inputs[\"input_ids\"] = new_input_ids\n            inputs[\"attention_mask\"] = new_attention_mask\n            inputs[\"labels\"] = new_labels\n\n        loss = super().training_step(model, inputs)\n        return loss\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)\n\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport dataclasses\nimport json\nimport random\nimport warnings\nfrom collections import deque\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Literal, Optional, Tuple, Union\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState, PartialState\nfrom rich.console import Console\nfrom rich.table import Table\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torch.utils.data import IterableDataset\nfrom transformers import (\n    BitsAndBytesConfig,\n    DataCollatorForLanguageModeling,\n    GenerationConfig,\n    PreTrainedTokenizerBase,\n    TrainerState,\n    TrainingArguments,\n)\nfrom transformers.utils import (\n    is_peft_available,\n    is_torch_mlu_available,\n    is_torch_npu_available,\n    is_torch_xpu_available,\n)\n\nfrom ..import_utils import is_unsloth_available\nfrom ..trainer.model_config import ModelConfig\n\n\nif is_peft_available():\n    from peft import LoraConfig, PeftConfig\n\n\nclass AdaptiveKLController:\n    \"\"\"\n    Adaptive KL controller described in the paper:\n    https://huggingface.co/papers/1909.08593\n    \"\"\"\n\n    def __init__(self, init_kl_coef, target, horizon):\n        self.value = init_kl_coef\n        self.target = target\n        self.horizon = horizon\n\n    def update(self, current, n_steps):\n        target = self.target\n        proportional_error = np.clip(current / target - 1, -0.2, 0.2)\n        mult = 1 + proportional_error * n_steps / self.horizon\n        self.value *= mult\n\n\nclass FixedKLController:\n    \"\"\"Fixed KL controller.\"\"\"\n\n    def __init__(self, kl_coef):\n        self.value = kl_coef\n\n    def update(self, current, n_steps):\n        pass\n\n\nclass DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling):\n    \"\"\"\n    Data collator used for completion tasks. It ensures that all the tokens of the labels are set to an 'ignore_index'\n    when they do not come from the assistant. This ensure that the loss is only\n    calculated on the completion made by the assistant.\n\n    Args:\n        response_template (`Union[str, List[int]]`): the template form that indicates the start of the response, typically something like\n            '### Response:\\n'. It can also be passed as tokenized ids, which can be useful when using a tokenizer that encodes the response\n            differently if it does not have proper context.\n        instruction_template (`Union[str, List[int]]`): the template form that indicates the start of the human instruction, typically something like\n            '### Human:\\n'. Useful for assistant-style conversation datasets. It can also be passed as tokenized ids.\n        mlm (`bool`, *optional*, defaults to `False`): Whether or not to use masked language modeling in the underlying\n            `DataCollatorForLanguageModeling` class. Note that this option currently has no effect but is present\n             for flexibility and backwards-compatibility.\n        ignore_index (`int`, *optional*, defaults to `-100`):\n            The index to use to ignore the initial tokens with\n    \"\"\"\n\n    def __init__(\n        self,\n        response_template: Union[str, List[int]],\n        instruction_template: Optional[Union[str, List[int]]] = None,\n        *args,\n        mlm: bool = False,\n        ignore_index: int = -100,\n        padding_free: bool = False,\n        **kwargs,\n    ):\n        super().__init__(*args, mlm=mlm, **kwargs)\n\n        self.instruction_template = instruction_template\n        if isinstance(instruction_template, str):\n            # The user provides a string, must tokenize\n            self.instruction_token_ids = self.tokenizer.encode(self.instruction_template, add_special_tokens=False)\n        else:\n            # The user already provides the token ids\n            self.instruction_token_ids = instruction_template\n\n        self.response_template = response_template\n        if isinstance(response_template, str):\n            # The user provides a string, must tokenize\n            self.response_token_ids = self.tokenizer.encode(self.response_template, add_special_tokens=False)\n        else:\n            # The user already provides the token ids\n            self.response_token_ids = response_template\n\n        if not self.mlm and self.instruction_template and self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:\n            warnings.warn(\n                \"The pad_token_id and eos_token_id values of this tokenizer are identical. \"\n                \"If you are planning for multi-turn training, \"\n                \"it can result in the model continuously generating questions and answers without eos token. \"\n                \"To avoid this, set the pad_token_id to a different value.\"\n            )\n\n        self.ignore_index = ignore_index\n        self.padding_free = padding_free\n\n    def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:\n        batch = super().torch_call(examples)\n\n        if self.instruction_template is None:\n            for i in range(len(examples)):\n                response_token_ids_start_idx = None\n\n                for idx in np.where(batch[\"labels\"][i] == self.response_token_ids[0])[0]:\n                    # `response_token_ids` is `'### Response:\\n'`, here we are just making sure that the token IDs match\n                    if (\n                        self.response_token_ids\n                        == batch[\"labels\"][i][idx : idx + len(self.response_token_ids)].tolist()\n                    ):\n                        response_token_ids_start_idx = idx\n\n                if response_token_ids_start_idx is None:\n                    warnings.warn(\n                        f\"Could not find response key `{self.response_template}` in the \"\n                        f'following instance: {self.tokenizer.decode(batch[\"input_ids\"][i])} '\n                        f\"This instance will be ignored in loss calculation. \"\n                        f\"Note, if this happens often, consider increasing the `max_seq_length`.\"\n                    )\n                    batch[\"labels\"][i, :] = self.ignore_index\n                else:\n                    response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids)\n\n                    # Make pytorch loss function ignore all tokens up through the end of the response key\n                    batch[\"labels\"][i, :response_token_ids_end_idx] = self.ignore_index\n\n        else:\n            for i in range(len(examples)):\n                response_token_ids_idxs = []\n                human_token_ids_idxs = []\n\n                for assistant_idx in np.where(batch[\"labels\"][i] == self.response_token_ids[0])[0]:\n                    # find the indexes of the start of a response.\n                    if (\n                        self.response_token_ids\n                        == batch[\"labels\"][i][assistant_idx : assistant_idx + len(self.response_token_ids)].tolist()\n                    ):\n                        response_token_ids_idxs.append(assistant_idx + len(self.response_token_ids))\n\n                if len(response_token_ids_idxs) == 0:\n                    warnings.warn(\n                        f\"Could not find response key `{self.response_template}` in the \"\n                        f'following instance: {self.tokenizer.decode(batch[\"input_ids\"][i])} '\n                        f\"This instance will be ignored in loss calculation. \"\n                        f\"Note, if this happens often, consider increasing the `max_seq_length`.\"\n                    )\n                    batch[\"labels\"][i, :] = self.ignore_index\n\n                human_token_ids = self.instruction_token_ids\n                for human_idx in np.where(batch[\"labels\"][i] == human_token_ids[0])[0]:\n                    # find the indexes of the start of a human answer.\n                    if human_token_ids == batch[\"labels\"][i][human_idx : human_idx + len(human_token_ids)].tolist():\n                        human_token_ids_idxs.append(human_idx)\n\n                if len(human_token_ids_idxs) == 0:\n                    warnings.warn(\n                        f\"Could not find instruction key `{self.instruction_template}` in the \"\n                        f'following instance: {self.tokenizer.decode(batch[\"input_ids\"][i])} '\n                        f\"This instance will be ignored in loss calculation. \"\n                        f\"Note, if this happens often, consider increasing the `max_seq_length`.\"\n                    )\n                    batch[\"labels\"][i, :] = self.ignore_index\n\n                if (\n                    len(human_token_ids_idxs) > 0\n                    and len(response_token_ids_idxs) > 0\n                    and human_token_ids_idxs[0] > response_token_ids_idxs[0]\n                ):\n                    human_token_ids_idxs = [0] + human_token_ids_idxs\n\n                for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)):\n                    # Make pytorch loss function ignore all non response tokens\n                    if idx != 0:\n                        batch[\"labels\"][i, start:end] = self.ignore_index\n                    else:\n                        batch[\"labels\"][i, :end] = self.ignore_index\n\n                if len(response_token_ids_idxs) < len(human_token_ids_idxs):\n                    batch[\"labels\"][i, human_token_ids_idxs[-1] :] = self.ignore_index\n\n        if self.padding_free:\n            # remove padding, `attention_mask` and add `position_ids`\n            attn_mask = batch.pop(\"attention_mask\")\n            batch[\"input_ids\"] = batch[\"input_ids\"][attn_mask.bool()].unsqueeze(0)\n            batch[\"position_ids\"] = attn_mask.cumsum(1)[attn_mask.bool()].unsqueeze(0) - 1\n            batch[\"labels\"] = batch[\"labels\"][attn_mask.bool()].unsqueeze(0)\n            batch[\"labels\"][batch[\"position_ids\"] == 0] = self.ignore_index\n\n        return batch\n\n\n@dataclass\nclass DataCollatorForChatML:\n    \"\"\"\n    Data collator for ChatML format datasets.\n    \"\"\"\n\n    tokenizer: PreTrainedTokenizerBase\n    ignore_index: int = -100\n    max_length: int = None\n    messages_key: str = \"messages\"\n\n    def __post_init__(self):\n        if self.tokenizer.pad_token_id is None:\n            raise ValueError(\"The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.\")\n        if self.max_length is None:\n            # set a sensible default\n            self.max_length = min(self.tokenizer.model_max_length, 1024)\n\n    def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:\n        prompts = []\n        completions = []\n\n        for example in examples:\n            messages = example[self.messages_key]\n            formatted_chat = self.tokenizer.apply_chat_template(messages, tokenize=False)\n\n            # Split the formatted chat into prompt and completion\n            assistant_messages = [msg for msg in messages if msg[\"role\"] == \"assistant\"]\n            last_assistant_message = assistant_messages[-1][\"content\"]\n            prompt = formatted_chat.rsplit(last_assistant_message, 1)[0]\n            completion = last_assistant_message\n\n            prompts.append(prompt)\n            completions.append(completion)\n\n        # Tokenize prompts and completions\n        tokenized_prompts = self.tokenizer(\n            prompts, truncation=True, max_length=self.max_length, padding=False, return_tensors=None\n        )\n        tokenized_completions = self.tokenizer(\n            completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None\n        )\n\n        # Combine prompts and completions\n        input_ids = []\n        attention_mask = []\n        labels = []\n\n        for prompt, completion in zip(tokenized_prompts[\"input_ids\"], tokenized_completions[\"input_ids\"]):\n            combined_input_ids = prompt + completion\n            combined_attention_mask = [1] * len(combined_input_ids)\n\n            # Create labels for one-token ahead task, masking the prompt\n            combined_labels = [self.ignore_index] * len(prompt) + completion[:-1]\n            combined_labels.append(self.tokenizer.eos_token_id)  # Add EOS token as final target\n\n            input_ids.append(combined_input_ids)\n            attention_mask.append(combined_attention_mask)\n            labels.append(combined_labels)\n\n        # first convert to list of tensors\n        input_ids = [torch.tensor(ids) for ids in input_ids]\n        attention_mask = [torch.tensor(mask) for mask in attention_mask]\n        labels = [torch.tensor(label) for label in labels]\n\n        # pad the input_ids, attention_mask and labels to the same length across the batch\n        input_ids = pad(input_ids, padding_side=\"left\", padding_value=self.tokenizer.pad_token_id)\n        attention_mask = pad(attention_mask, padding_side=\"left\", padding_value=0)\n        labels = pad(labels, padding_side=\"left\", padding_value=self.ignore_index)\n\n        # pad the tokenized_prompts on the left to the same length convert to tensor first\n        prompts_input_ids = [torch.tensor(ids) for ids in tokenized_prompts[\"input_ids\"]]\n        prompts_input_ids = pad(prompts_input_ids, padding_side=\"left\", padding_value=self.tokenizer.pad_token_id)\n\n        # prompt attention mask\n        prompt_attention_mask = pad(\n            [torch.tensor([1] * len(ids)) for ids in tokenized_prompts[\"input_ids\"]],\n            padding_side=\"left\",\n            padding_value=0,\n        )\n\n        return {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n            \"labels\": labels,\n            \"prompts\": prompts_input_ids,\n            \"prompt_attention_mask\": prompt_attention_mask,\n        }\n\n\n@dataclass\nclass RewardDataCollatorWithPadding:\n    r\"\"\"\n    Reward DataCollator class that pads the inputs to the maximum length of the batch.\n\n    Args:\n        tokenizer (`PreTrainedTokenizerBase`):\n            The tokenizer used for encoding the data.\n        padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`):\n            padding_strategy to pass to the tokenizer.\n        max_length (`Optional[int]`, `optional`, defaults to `None`):\n            The maximum length of the sequence to be processed.\n        pad_to_multiple_of (`Optional[int]`, `optional`, defaults to `None`):\n            If set will pad the sequence to a multiple of the provided value.\n        return_tensors (`str`, `optional`, defaults to `\"pt\"`):\n            The tensor type to use.\n    \"\"\"\n\n    tokenizer: PreTrainedTokenizerBase\n    padding: Union[bool, str] = True\n    max_length: Optional[int] = None\n    pad_to_multiple_of: Optional[int] = None\n    return_tensors: str = \"pt\"\n\n    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:\n        features_chosen = []\n        features_rejected = []\n        margin = []\n        # check if we have a margin. If we do, we need to batch it as well\n        has_margin = \"margin\" in features[0]\n        for feature in features:\n            # check if the keys are named as expected\n            if (\n                \"input_ids_chosen\" not in feature\n                or \"input_ids_rejected\" not in feature\n                or \"attention_mask_chosen\" not in feature\n                or \"attention_mask_rejected\" not in feature\n            ):\n                raise ValueError(\n                    \"The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`\"\n                )\n\n            features_chosen.append(\n                {\n                    \"input_ids\": feature[\"input_ids_chosen\"],\n                    \"attention_mask\": feature[\"attention_mask_chosen\"],\n                }\n            )\n            features_rejected.append(\n                {\n                    \"input_ids\": feature[\"input_ids_rejected\"],\n                    \"attention_mask\": feature[\"attention_mask_rejected\"],\n                }\n            )\n            if has_margin:\n                margin.append(feature[\"margin\"])\n        batch_chosen = self.tokenizer.pad(\n            features_chosen,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=self.return_tensors,\n        )\n        batch_rejected = self.tokenizer.pad(\n            features_rejected,\n            padding=self.padding,\n            max_length=self.max_length,\n            pad_to_multiple_of=self.pad_to_multiple_of,\n            return_tensors=self.return_tensors,\n        )\n        batch = {\n            \"input_ids_chosen\": batch_chosen[\"input_ids\"],\n            \"attention_mask_chosen\": batch_chosen[\"attention_mask\"],\n            \"input_ids_rejected\": batch_rejected[\"input_ids\"],\n            \"attention_mask_rejected\": batch_rejected[\"attention_mask\"],\n            \"return_loss\": True,\n        }\n        if has_margin:\n            margin = torch.tensor(margin, dtype=torch.float)\n            batch[\"margin\"] = margin\n        return batch\n\n\ndef pad(tensors: List[torch.Tensor], padding_value: int = 0, padding_side: str = \"right\") -> torch.Tensor:\n    \"\"\"\n    Pads a list of tensors to the same shape along the first dimension.\n\n    Args:\n        tensors (`List[torch.Tensor]`):\n            List of input tensors to pad.\n        padding_value (`int`):\n            Value to use for padding. Default is 0.\n        padding_side (`str`):\n            Side on which to add padding. Must be 'left' or 'right'. Default is 'right'.\n\n    Returns:\n        `torch.Tensor`:\n            A single tensor containing the padded tensors.\n\n    Examples:\n        >>> import torch\n        >>> pad([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])\n        tensor([[1, 2, 3],\n                [4, 5, 0]])\n        >>> pad([torch.tensor([[1, 2], [3, 4]]), torch.tensor([[5, 6]])])\n        tensor([[[1, 2],\n                [3, 4]],\n\n                [[5, 6],\n                [0, 0]]])\n    \"\"\"\n    # Determine the maximum shape for each dimension\n    output_shape = np.max([t.shape for t in tensors], 0).tolist()\n\n    # Create an output tensor filled with the padding value\n    output = torch.full((len(tensors), *output_shape), padding_value, dtype=tensors[0].dtype, device=tensors[0].device)\n\n    for i, t in enumerate(tensors):\n        # Determine the slice for the sequence dimension\n        if padding_side == \"left\":\n            seq_slice = slice(output_shape[0] - t.shape[0], output_shape[0])\n        elif padding_side == \"right\":\n            seq_slice = slice(0, t.shape[0])\n        else:\n            raise ValueError(\"padding_side must be 'left' or 'right'\")\n\n        slices = (seq_slice,) + tuple(slice(0, s) for s in t.shape[1:])\n        output[i][slices] = t\n\n    return output\n\n\n@dataclass\nclass DPODataCollatorWithPadding:\n    r\"\"\"\n    DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch.\n\n    Args:\n        pad_token_id (`int` defaults to 0):\n            The tokenizer's pad_token_id.\n        label_pad_token_id (`int`, defaults to -100):\n            The label used for masking.\n        is_encoder_decoder (`Optional[bool]`, `optional`, defaults to `None`):\n            Whether or not you model has an encoder_decoder architecture.\n    \"\"\"\n\n    pad_token_id: int = 0\n    label_pad_token_id: int = -100\n    is_encoder_decoder: Optional[bool] = False\n\n    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:\n        # first, pad everything to the same length\n        padded_batch = {}\n        for k in features[0].keys():\n            if k.endswith((\"_input_ids\", \"_attention_mask\", \"_labels\", \"_pixel_values\")):\n                if self.is_encoder_decoder:\n                    to_pad = [torch.LongTensor(ex[k]) for ex in features]\n\n                    if (k.startswith(\"prompt\")) and (k.endswith(\"input_ids\")):\n                        if self.pad_token_id is None:\n                            raise ValueError(\n                                \"Padding is enabled, but the tokenizer is not configured with a padding token.\"\n                                \" Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)\"\n                                \" before calling the trainer.\"\n                            )\n                        padding_value = self.pad_token_id\n                    elif k.endswith(\"_attention_mask\"):\n                        padding_value = 0\n                    elif k.startswith((\"chosen\", \"rejected\", \"completion\")) or (\"decoder\" in k):\n                        padding_value = self.label_pad_token_id\n                    else:\n                        raise ValueError(f\"Unexpected key in batch '{k}'\")\n                    padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value)\n                else:\n                    # Set padding value based on the key\n                    if k.endswith(\"_input_ids\"):\n                        if self.pad_token_id is None:\n                            raise ValueError(\n                                \"Padding is enabled, but the tokenizer is not configured with a padding token.\"\n                                \" Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)\"\n                                \" before calling the trainer.\"\n                            )\n                        padding_value = self.pad_token_id\n                    elif k.endswith(\"_labels\"):\n                        padding_value = self.label_pad_token_id\n                    elif k.endswith(\"_attention_mask\"):\n                        padding_value = 0\n                    elif k.endswith(\"_pixel_values\"):\n                        padding_value = 0  # TODO: check if this is correct\n                    else:\n                        raise ValueError(f\"Unexpected key in batch '{k}'\")\n\n                    # Set padding side based on the key\n                    if k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                        padding_side = \"left\"\n                    else:\n                        padding_side = \"right\"\n\n                    # Set the dtype\n                    if k.endswith(\"_pixel_values\"):\n                        dtype = torch.float32  # will be downcasted if necessary by the Trainer\n                    else:\n                        dtype = torch.int64\n\n                    # Convert to tensor and pad\n                    to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features]\n                    padded_batch[k] = pad(to_pad, padding_value=padding_value, padding_side=padding_side)\n            elif k.endswith(\"_logps\"):\n                # the cached reference model logprobs\n                padded_batch[k] = torch.tensor([ex[k] for ex in features])\n            else:\n                padded_batch[k] = [ex[k] for ex in features]\n\n        return padded_batch\n\n\nclass ConstantLengthDataset(IterableDataset):\n    \"\"\"\n    Iterable dataset that returns constant length chunks of tokens from stream of text files.\n    The dataset also formats the text before tokenization with a specific format that is provided\n    by the user.\n\n    Args:\n        tokenizer (`transformers.PreTrainedTokenizer`):\n            The processor used for processing the data.\n        dataset (`dataset.Dataset`):\n            Dataset with text files.\n        dataset_text_field (`Optional[str]`, *optional*, defaults to `None`):\n            Name of the field in the dataset that contains the text. Used only if `formatting_func` is `None`.\n        formatting_func (`Callable`, *optional*):\n            Function that formats the text before tokenization. Usually it is recommended to have follows a certain\n            pattern such as `\"### Question: {question} ### Answer: {answer}\"`\n        infinite (`bool`, *optional*, defaults to `False`):\n            If True the iterator is reset after dataset reaches end else stops.\n        seq_length (`int`, *optional*, defaults to `1024`):\n            Length of token sequences to return.\n        num_of_sequences (`int`, *optional*, defaults to `1024`):\n            Number of token sequences to keep in buffer.\n        chars_per_token (`int`, *optional*, defaults to `3.6`):\n            Number of characters per token used to estimate number of tokens in text buffer.\n        eos_token_id (`int`, *optional*, defaults to `0`):\n            Id of the end of sequence token if the passed tokenizer does not have an EOS token.\n        shuffle (`bool`, *optional*, defaults to `True`)\n            Shuffle the examples before they are returned\n        append_concat_token (`bool`, *optional*, defaults to `True`)\n            If true, appends `eos_token_id` at the end of each sample being packed.\n        add_special_tokens (`bool`, *optional*, defaults to `True`)\n            If true, tokenizers adds special tokens to each sample being packed.\n    \"\"\"\n\n    def __init__(\n        self,\n        tokenizer,\n        dataset,\n        dataset_text_field=None,\n        formatting_func=None,\n        infinite=False,\n        seq_length=1024,\n        num_of_sequences=1024,\n        chars_per_token=3.6,\n        eos_token_id=0,\n        shuffle=True,\n        append_concat_token=True,\n        add_special_tokens=True,\n    ):\n        self.tokenizer = tokenizer\n\n        if tokenizer.eos_token_id is None:\n            warnings.warn(\n                \"The passed tokenizer does not have an EOS token. We will use the passed eos_token_id instead which corresponds\"\n                f\" to {eos_token_id}. If this is not the correct EOS token, make sure to pass the correct eos_token_id.\"\n            )\n\n        self.concat_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id else eos_token_id\n        self.dataset = dataset\n        self.seq_length = seq_length\n        self.infinite = infinite\n        self.current_size = 0\n        self.max_buffer_size = seq_length * chars_per_token * num_of_sequences\n        self.shuffle = shuffle\n        self.append_concat_token = append_concat_token\n        self.add_special_tokens = add_special_tokens\n        if formatting_func is None:\n            self.formatting_func = lambda x: x[dataset_text_field]\n        else:\n            self.formatting_func = formatting_func\n\n        if formatting_func is not None:\n            if formatting_func.__code__.co_argcount > 1:\n                warnings.warn(\n                    \"The passed formatting_func has more than one argument. Usually that function should have a single argument `example`\"\n                    \" which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing.\"\n                )\n\n    def __len__(self):\n        return len(self.dataset)\n\n    def __iter__(self):\n        iterator = iter(self.dataset)\n        more_examples = True\n        while more_examples:\n            buffer, buffer_len = [], 0\n            while True:\n                if buffer_len >= self.max_buffer_size:\n                    break\n                try:\n                    buffer.append(self.formatting_func(next(iterator)))\n                    buffer_len += len(buffer[-1])\n                except StopIteration:\n                    if self.infinite:\n                        iterator = iter(self.dataset)\n                        warnings.warn(\"The dataset reached end and the iterator is reset to the start.\")\n                    else:\n                        more_examples = False\n                        break\n            if self.shuffle:\n                random.shuffle(buffer)\n            tokenized_inputs = self.tokenizer(buffer, add_special_tokens=self.add_special_tokens, truncation=False)[\n                \"input_ids\"\n            ]\n            all_token_ids = []\n            for tokenized_input in tokenized_inputs:\n                if self.append_concat_token:\n                    tokenized_input = tokenized_input + [self.concat_token_id]\n                all_token_ids.extend(tokenized_input)\n            examples = []\n            for i in range(0, len(all_token_ids), self.seq_length):\n                input_ids = all_token_ids[i : i + self.seq_length]\n                if len(input_ids) == self.seq_length:\n                    examples.append(input_ids)\n            if self.shuffle:\n                # Shuffle again, otherwise split examples occur in consecutive tensors.\n                random.shuffle(examples)\n            for example in examples:\n                self.current_size += 1\n                yield {\n                    \"input_ids\": torch.LongTensor(example),\n                    \"labels\": torch.LongTensor(example),\n                }\n\n\n@dataclass\nclass RunningMoments:\n    \"\"\"\n    Calculates the running mean and standard deviation of a data stream. Reference:\n    https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L75\n    \"\"\"\n\n    accelerator: Accelerator\n    mean: float = 0\n    std: float = 1\n    var: float = 1\n    count: float = 1e-24\n\n    @torch.no_grad()\n    def update(self, xs: torch.Tensor) -> Tuple[float, float]:\n        \"\"\"\n        Updates running moments from batch's moments computed across ranks\n        \"\"\"\n        if self.accelerator.use_distributed:\n            xs_mean, xs_var, xs_count = get_global_statistics(self.accelerator, xs)\n        else:\n            xs_count = xs.numel()\n            xs_var, xs_mean = torch.var_mean(xs, unbiased=False)\n        xs_mean, xs_var = xs_mean.float(), xs_var.float()\n\n        delta = xs_mean - self.mean\n        tot_count = self.count + xs_count\n\n        new_sum = xs_var * xs_count\n        # correct old_sum deviation accounting for the new mean\n        old_sum = self.var * self.count + delta**2 * self.count * xs_count / tot_count\n        tot_sum = old_sum + new_sum\n\n        self.mean += (delta * xs_count / tot_count).item()\n        new_var = tot_sum / tot_count\n        self.std = (new_var * tot_count / (tot_count - 1)).float().sqrt().item()\n        self.var = new_var.item()\n        self.count = tot_count\n\n        return xs_mean.item(), (xs_var * xs_count / (xs_count - 1)).float().sqrt().item()\n\n    def save_to_json(self, json_path: str):\n        \"\"\"Save the content of this instance in JSON format inside `json_path`.\"\"\"\n        # save everything except accelerator\n        if self.accelerator.is_main_process:\n            save_dict = dataclasses.asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if k != \"accelerator\"})\n            json_string = json.dumps(save_dict, indent=2, sort_keys=True) + \"\\n\"\n            with open(json_path, \"w\", encoding=\"utf-8\") as f:\n                f.write(json_string)\n\n    @classmethod\n    def load_from_json(cls, accelerator: Accelerator, json_path: str):\n        \"\"\"Create an instance from the content of `json_path`.\"\"\"\n        # load everything except accelerator\n        with open(json_path, encoding=\"utf-8\") as f:\n            text = f.read()\n        return cls(accelerator=accelerator, **json.loads(text))\n\n\n@torch.no_grad()\ndef get_global_statistics(\n    accelerator, xs: torch.Tensor, mask=None, device=\"cpu\"\n) -> Tuple[torch.Tensor, torch.Tensor, int]:\n    \"\"\"\n    Computes element-wise mean and variance of the tensor across processes. Reference:\n    https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L57C1-L73C75\n    \"\"\"\n    xs = xs.to(accelerator.device)\n    sum_and_count = torch.tensor([xs.sum(), (xs.numel() if mask is None else mask.sum())], device=xs.device)\n    sum_and_count = accelerator.reduce(sum_and_count)\n    global_sum, count = sum_and_count\n    global_mean = global_sum / count\n\n    sum_var = torch.sum(((xs - global_mean) ** 2).mul(1 if mask is None else mask))\n    sum_var = accelerator.reduce(sum_var)\n    global_var = sum_var / count\n\n    return global_mean.to(device), global_var.to(device), count.item()\n\n\ndef compute_accuracy(eval_pred) -> Dict[str, float]:\n    predictions, labels = eval_pred\n    # Here, predictions is rewards_chosen and rewards_rejected.\n    # We want to see how much of the time rewards_chosen > rewards_rejected.\n    if np.array(predictions[:, 0] == predictions[:, 1], dtype=float).sum() > 0:\n        warnings.warn(\n            f\"There are {np.array(predictions[:, 0] == predictions[:, 1]).sum()} out of {len(predictions[:, 0])} instances where the predictions for both options are equal. As a consequence the accuracy can be misleading.\"\n        )\n    predictions = np.argmax(predictions, axis=1)\n\n    accuracy = np.array(predictions == labels, dtype=float).mean().item()\n    return {\"accuracy\": accuracy}\n\n\ndef pad_to_length(tensor: torch.Tensor, length: int, pad_value: Union[int, float], dim: int = -1) -> torch.Tensor:\n    if tensor.size(dim) >= length:\n        return tensor\n    else:\n        pad_size = list(tensor.shape)\n        pad_size[dim] = length - tensor.size(dim)\n        return torch.cat(\n            [\n                tensor,\n                pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device),\n            ],\n            dim=dim,\n        )\n\n\ndef disable_dropout_in_model(model: torch.nn.Module) -> None:\n    for module in model.modules():\n        if isinstance(module, torch.nn.Dropout):\n            module.p = 0\n\n\ndef exact_div(a, b, custom_error_message=\"\"):\n    q = a // b\n    if a != q * b:\n        raise ValueError(f\"{custom_error_message}, inexact division: {a} / {b} = {a / b}\")\n    return q\n\n\n# copied from https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/stat_tracking.py#L5\nclass PerPromptStatTracker:\n    r\"\"\"\n    Class for tracking statistics per prompt. Mainly used to calculate advantage for the DPPO algorithm\n\n    Args:\n        buffer_size (`int`):\n            Size of the buffer to keep for each prompt.\n        min_count (`int`):\n            Minimum number of samples to keep in the buffer before calculating the mean and std.\n    \"\"\"\n\n    def __init__(self, buffer_size, min_count):\n        self.buffer_size = buffer_size\n        self.min_count = min_count\n        self.stats = {}\n\n    def update(self, prompts, rewards):\n        prompts = np.array(prompts)\n        rewards = np.array(rewards)\n        unique = np.unique(prompts)\n        advantages = np.empty_like(rewards)\n        for prompt in unique:\n            prompt_rewards = rewards[prompts == prompt]\n            if prompt not in self.stats:\n                self.stats[prompt] = deque(maxlen=self.buffer_size)\n            self.stats[prompt].extend(prompt_rewards)\n\n            if len(self.stats[prompt]) < self.min_count:\n                mean = np.mean(rewards)\n                std = np.std(rewards) + 1e-6\n            else:\n                mean = np.mean(self.stats[prompt])\n                std = np.std(self.stats[prompt]) + 1e-6\n            advantages[prompts == prompt] = (prompt_rewards - mean) / std\n\n        return advantages\n\n    def get_stats(self):\n        return {k: {\"mean\": np.mean(v), \"std\": np.std(v), \"count\": len(v)} for k, v in self.stats.items()}\n\n\ndef peft_module_casting_to_bf16(model):\n    for name, module in model.named_modules():\n        if isinstance(module, torch.nn.LayerNorm) or \"norm\" in name:\n            module = module.to(torch.float32)\n        elif any(x in name for x in [\"lm_head\", \"embed_tokens\", \"wte\", \"wpe\"]):\n            if hasattr(module, \"weight\"):\n                if module.weight.dtype == torch.float32:\n                    module = module.to(torch.bfloat16)\n\n\ndef trl_sanitze_kwargs_for_tagging(model, tag_names, kwargs=None):\n    if is_unsloth_available():\n        # Unsloth adds a new attribute in the model config `unsloth_version`\n        # to keep track of models that have been patched with unsloth.\n        if hasattr(model, \"config\") and getattr(model.config, \"unsloth_version\", None) is not None:\n            tag_names.append(\"unsloth\")\n\n    if kwargs is not None:\n        if \"tags\" not in kwargs:\n            kwargs[\"tags\"] = tag_names\n        elif \"tags\" in kwargs and isinstance(kwargs[\"tags\"], list):\n            kwargs[\"tags\"].extend(tag_names)\n        elif \"tags\" in kwargs and isinstance(kwargs[\"tags\"], str):\n            tag_names.append(kwargs[\"tags\"])\n            kwargs[\"tags\"] = tag_names\n    return kwargs\n\n\ndef get_quantization_config(model_config: ModelConfig) -> Optional[BitsAndBytesConfig]:\n    if model_config.load_in_4bit:\n        quantization_config = BitsAndBytesConfig(\n            load_in_4bit=True,\n            bnb_4bit_compute_dtype=model_config.torch_dtype,  # For consistency with model weights, we use the same value as `torch_dtype`\n            bnb_4bit_quant_type=model_config.bnb_4bit_quant_type,\n            bnb_4bit_use_double_quant=model_config.use_bnb_nested_quant,\n            bnb_4bit_quant_storage=model_config.torch_dtype,\n        )\n    elif model_config.load_in_8bit:\n        quantization_config = BitsAndBytesConfig(\n            load_in_8bit=True,\n        )\n    else:\n        quantization_config = None\n\n    return quantization_config\n\n\ndef get_kbit_device_map() -> Optional[Dict[str, int]]:\n    if is_torch_xpu_available():\n        return {\"\": f\"xpu:{PartialState().local_process_index}\"}\n    elif torch.cuda.is_available():\n        return {\"\": PartialState().local_process_index}\n    else:\n        return None\n\n\ndef get_peft_config(model_config: ModelConfig) -> \"Optional[PeftConfig]\":\n    if model_config.use_peft is False:\n        return None\n\n    if not is_peft_available():\n        raise ValueError(\n            \"You need to have PEFT library installed in your environment, make sure to install `peft`. \"\n            \"Make sure to run `pip install -U peft`.\"\n        )\n\n    peft_config = LoraConfig(\n        task_type=model_config.lora_task_type,\n        r=model_config.lora_r,\n        target_modules=model_config.lora_target_modules,\n        lora_alpha=model_config.lora_alpha,\n        lora_dropout=model_config.lora_dropout,\n        bias=\"none\",\n        use_rslora=model_config.use_rslora,\n        modules_to_save=model_config.lora_modules_to_save,\n    )\n\n    return peft_config\n\n\ndef get_exp_cap(value, decimal=4):\n    \"\"\"\n    Get the exponent cap of a value. This is used to cap the exponent of a value to avoid overflow.\n    The formula is : log(value.dtype.max)\n    E.g.\n      For float32 data type, the maximum exponent value is 88.7228 to 4 decimal points.\n    ```\n\n    Args:\n        value (`torch.Tensor`):\n            The input tensor to obtain the data type\n        decimal (`int`):\n            The number of decimal points of the output exponent cap.\n            eg: direct calling exp(log(torch.float32.max)) will result in inf\n            so we cap the exponent to 88.7228 to avoid overflow.\n    \"\"\"\n    vdtype_max = torch.zeros([1]).to(value.dtype) + torch.finfo(value.dtype).max\n    vdtype_log_max = torch.log(vdtype_max).to(value.device)\n    return torch.floor(vdtype_log_max * 10**decimal) / 10**decimal if decimal > 0 else vdtype_log_max\n\n\ndef cap_exp(value, cap=-1):\n    # Cap the exponent value below the upper-bound to avoid overflow, before calling torch.exp\n    cap = get_exp_cap(value) if cap < 0 else cap\n    return torch.exp(torch.clamp(value, max=cap))\n\n\ndef print_rich_table(df: pd.DataFrame) -> Table:\n    console = Console()\n    table = Table(show_lines=True)\n    for column in df.columns:\n        table.add_column(column)\n    for _, row in df.iterrows():\n        table.add_row(*row.astype(str).tolist())\n    console.print(table)\n\n\nSIMPLE_SFT_CHAT_TEMPLATE = \"{% for message in messages %}{{' ' + message['content']}}{% endfor %}{{ eos_token }}\"\n# SIMPLE_SFT_CHAT_TEMPLATE simply ends things with an EOS token, this helps the SFT model learn to end the completions with EOS tokens\n\nSIMPLE_CHAT_TEMPLATE = \"{% for message in messages %}{{message['role'].capitalize() + ': ' + message['content'] + '\\n\\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}\"\n\n\n@dataclass\nclass OnlineTrainerState(TrainerState):\n    episode: int = 0\n\n\n@dataclass\nclass OnPolicyConfig(TrainingArguments):\n    r\"\"\"\n    Base configuration class for on-policy trainers.\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        run_name (`Optional[str]`, *optional*, defaults to `None`):\n            Name of the run.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n        num_mini_batches (`int`, *optional*, defaults to `1`):\n            Number of minibatches to split a batch into.\n        total_episodes (`Optional[int]`, *optional*, defaults to `None`):\n            Total number of episodes in the dataset.\n        local_rollout_forward_batch_size (`int`, *optional*, defaults to `64`):\n            Per rank no grad forward pass in the rollout phase.\n        num_sample_generations (`int`, *optional*, defaults to `10`):\n            Number of debugging samples generations (i.e., `generate_completions` calls) throughout training.\n        response_length (`int`, *optional*, defaults to `53`):\n            Length of the response.\n        stop_token (`Optional[str]`, *optional*, defaults to `None`):\n            Stop token.\n        stop_token_id (`Optional[int]`, *optional*, defaults to `None`):\n            Truncation token id.\n        temperature (`float`, *optional*, defaults to `0.7`):\n            Sampling temperature.\n        missing_eos_penalty (`Optional[float]`, *optional*, defaults to `None`):\n            Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage\n            to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive\n            value.\n        sft_model_path (`str`, *optional*, defaults to `\"EleutherAI/pythia-160m\"`):\n            Path to the SFT model.\n        world_size (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes (GPUs) to use for the training.\n        num_total_batches (`Optional[int]`, *optional*, defaults to `None`):\n            Number of total batches to train.\n        micro_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Micro batch size across devices (HF's `per_device_train_batch_size` * `world_size`).\n        local_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Batch size per GPU (HF's `per_device_train_batch_size` * `gradient_accumulation_steps`).\n        batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Batch size across devices (HF's `per_device_train_batch_size` * `world_size` * `gradient_accumulation_steps`).\n        local_mini_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Mini batch size per GPU.\n        mini_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Mini batch size across GPUs.\n    \"\"\"\n\n    run_name: Optional[str] = None\n    dataset_num_proc: Optional[int] = None\n    num_mini_batches: int = 1\n    total_episodes: Optional[int] = None\n    local_rollout_forward_batch_size: int = 64\n    num_sample_generations: int = 10\n    response_length: int = 53\n    stop_token: Optional[Literal[\"eos\"]] = None\n    stop_token_id: Optional[int] = None\n    temperature: float = 0.7\n    missing_eos_penalty: Optional[float] = None\n    sft_model_path: str = \"EleutherAI/pythia-160m\"\n    world_size: Optional[int] = None\n    num_total_batches: Optional[int] = None\n    micro_batch_size: Optional[int] = None\n    local_batch_size: Optional[int] = None\n    batch_size: Optional[int] = None\n    local_mini_batch_size: Optional[int] = None\n    mini_batch_size: Optional[int] = None\n\n\ndef first_true_indices(bools: torch.Tensor, dtype=torch.long):\n    \"\"\"\n    Takes an N-dimensional bool tensor and returns an (N-1)-dimensional tensor of integers giving\n    the position of the first True in each \"row\".\n\n    Returns the length of the rows (bools.size(-1)) if no element is True in a given row.\n\n    Args:\n        bools (`torch.Tensor`):\n            An N-dimensional boolean tensor.\n        dtype (`torch.dtype`, optional):\n            The desired data type of the output tensor. Defaults to `torch.long`.\n\n    Returns:\n        `torch.Tensor`:\n            An (N-1)-dimensional tensor of integers indicating the position of the first True\n            in each row. If no True value is found in a row, returns the length of the row.\n    \"\"\"\n    row_len = bools.size(-1)\n    zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device)\n    return torch.min(zero_or_index, dim=-1).values\n\n\ndef get_reward(\n    model: torch.nn.Module, query_responses: torch.Tensor, pad_token_id: int, context_length: int\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Computes the reward logits and the rewards for a given model and query responses.\n\n    Args:\n        model (`torch.nn.Module`):\n            The model used to compute the reward logits.\n        query_responses (`torch.Tensor`):\n            The tensor containing the query responses.\n        pad_token_id (`int`):\n            The token ID representing the pad token.\n        context_length (`int`):\n            The length of the context in the query responses.\n\n    Returns:\n        tuple:\n            - `reward_logits` (`torch.Tensor`):\n                The logits for the reward model.\n            - `final_rewards` (`torch.Tensor`):\n                The final rewards for each query response.\n            - `sequence_lengths` (`torch.Tensor`):\n                The lengths of the sequences in the query responses.\n    \"\"\"\n    attention_mask = query_responses != pad_token_id\n    position_ids = attention_mask.cumsum(1) - attention_mask.long()  # exclusive cumsum\n    lm_backbone = getattr(model, model.base_model_prefix)\n    input_ids = torch.masked_fill(query_responses, ~attention_mask, 0)\n    output = lm_backbone(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        return_dict=True,\n        output_hidden_states=True,\n        use_cache=False,  # otherwise mistral-based RM would error out\n    )\n    reward_logits = model.score(output.hidden_states[-1])\n    sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length\n    # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454\n    return (\n        reward_logits,\n        reward_logits[\n            torch.arange(reward_logits.size(0), device=reward_logits.device),\n            sequence_lengths,\n        ].squeeze(-1),\n        sequence_lengths,\n    )\n\n\ndef forward(\n    model: torch.nn.Module,\n    query_responses: torch.Tensor,\n    pad_token_id: int,\n) -> torch.nn.Module:\n    \"\"\"\n    Performs a forward pass through the model with the given query responses and pad token ID.\n\n    Args:\n        model (`torch.nn.Module`):\n            The model to perform the forward pass.\n        query_responses (`torch.Tensor`):\n            The tensor containing the query responses.\n        pad_token_id (`int`):\n            The token ID representing the pad token.\n\n    Returns:\n        `torch.nn.Module`:\n            The output of the model, including hidden states.\n    \"\"\"\n    attention_mask = query_responses != pad_token_id\n    position_ids = attention_mask.cumsum(1) - attention_mask.long()\n    input_ids = torch.masked_fill(query_responses, ~attention_mask, 0)\n    return model(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        return_dict=True,\n        output_hidden_states=True,\n    )\n\n\ndef prepare_deepspeed(\n    model: torch.nn.Module, per_device_train_batch_size: int, fp16: bool = False, bf16: bool = False\n):\n    \"\"\"\n    Prepares the model for training with DeepSpeed (both for stage 2 and 3), configuring the appropriate settings based on the model and\n    batch size.\n\n    Args:\n        model (`torch.nn.Module`):\n            The model to be prepared for DeepSpeed training.\n        per_device_train_batch_size (`int`):\n            The training batch size per device.\n\n    Returns:\n        `torch.nn.Module`:\n            The model initialized and configured with DeepSpeed for training.\n    \"\"\"\n    import deepspeed\n\n    deepspeed_plugin = AcceleratorState().deepspeed_plugin\n    config_kwargs = deepspeed_plugin.deepspeed_config\n    if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n        config_kwargs[\"train_micro_batch_size_per_gpu\"] = per_device_train_batch_size\n        config_kwargs = {\n            \"train_micro_batch_size_per_gpu\": config_kwargs[\"train_micro_batch_size_per_gpu\"],\n            \"prescale_gradients\": False,\n            \"wall_clock_breakdown\": False,\n        }\n        if bf16:\n            config_kwargs[\"bf16\"] = {\"enabled\": True}\n        elif fp16:\n            config_kwargs[\"fp16\"] = {\"enabled\": True}\n    else:\n        if hasattr(model, \"config\"):\n            hidden_size = (\n                max(model.config.hidden_sizes)\n                if getattr(model.config, \"hidden_sizes\", None)\n                else getattr(model.config, \"hidden_size\", None)\n            )\n            if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                config_kwargs.update(\n                    {\n                        \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                        \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                        \"zero_optimization.stage3_prefetch_bucket_size\": 0,\n                    }\n                )\n    model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n    model.eval()\n    return model\n\n\ndef truncate_response(stop_token_id: int, pad_token_id: int, responses: torch.Tensor):\n    \"\"\"\n    Truncates the responses at the first occurrence of the stop token, filling the rest with pad tokens.\n\n    Args:\n        stop_token_id (`int`):\n            The token ID representing the stop token where truncation occurs.\n        pad_token_id (`int`):\n            The token ID representing the pad token used to fill the truncated responses.\n        responses (`torch.Tensor`):\n            The tensor containing the responses to be truncated.\n\n    Returns:\n        `torch.Tensor`:\n            The truncated responses tensor with pad tokens filled after the stop token.\n    \"\"\"\n    trunc_idxs = first_true_indices(responses == stop_token_id).unsqueeze(-1)\n    new_size = [1] * (len(responses.size()) - 1) + [responses.shape[1]]\n    idxs = torch.arange(responses.shape[1], device=responses.device).view(*new_size)\n    postprocessed_responses = torch.masked_fill(responses, idxs > trunc_idxs, pad_token_id)\n    return postprocessed_responses\n\n\ndef generate(\n    lm_backbone: torch.nn.Module, queries: torch.Tensor, pad_token_id: int, generation_config: GenerationConfig\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Generates sequences from the language model backbone in a way that does not affect padding tokens.\n\n    Args:\n        lm_backbone (`torch.nn.Module`):\n            The language model backbone used for generation.\n        queries (`torch.Tensor`):\n            The tensor containing the input queries.\n        pad_token_id (`int`):\n            The token ID representing the pad token.\n        generation_config (`GenerationConfig`):\n            The configuration for the generation process.\n\n    Returns:\n        tuple:\n            - `generated_sequences` (`torch.Tensor`):\n                The concatenated tensor of input queries and generated sequences.\n            - `logits` (`torch.Tensor`):\n                The logits output from the generation process.\n    \"\"\"\n    context_length = queries.shape[1]\n    attention_mask = queries != pad_token_id\n    input_ids = torch.masked_fill(queries, ~attention_mask, 0)\n    output = lm_backbone.generate(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        # position_ids=attention_mask.cumsum(1) - attention_mask.long(), # not needed: already adjusted in generations\n        # https://github.com/huggingface/transformers/blob/ac33aeeeee2a7a89b89c93c2962e6feb90daef0a/src/transformers/models/gpt2/modeling_gpt2.py#L1227-L1250\n        generation_config=generation_config,\n        return_dict_in_generate=True,\n        output_scores=True,\n    )\n    logits = torch.stack(output.scores, 1)\n    return torch.cat((queries, output.sequences[:, context_length:]), dim=1), logits\n\n\n@torch.no_grad()\ndef batch_generation(\n    model: torch.nn.Module,\n    queries: torch.Tensor,\n    local_rollout_forward_batch_size: int,\n    pad_token_id: int,\n    generation_config: GenerationConfig,\n):\n    query_responses = []\n    logitss = []\n    for i in range(0, queries.shape[0], local_rollout_forward_batch_size):\n        query = queries[i : i + local_rollout_forward_batch_size]\n        query_response, logits = generate(\n            model,\n            query,\n            pad_token_id,\n            generation_config,\n        )\n        query_responses.append(query_response)\n        logitss.append(logits)\n    return torch.cat(query_responses, 0), torch.cat(logitss, 0)\n\n\ndef add_bos_token_if_needed(\n    bos_token_id: Optional[int],\n    prompt_len_input_ids: int,\n    prompt_tokens: Dict[str, List[int]],\n    chosen_prompt_len_input_ids: int,\n    chosen_tokens: Dict[str, List[int]],\n    rejected_prompt_len_input_ids: int,\n    rejected_tokens: Dict[str, List[int]],\n):\n    if bos_token_id is not None:\n        if prompt_len_input_ids == 0 or bos_token_id != prompt_tokens[\"prompt_input_ids\"][0]:\n            prompt_tokens[\"prompt_input_ids\"] = [bos_token_id] + prompt_tokens[\"prompt_input_ids\"]\n            prompt_tokens[\"prompt_attention_mask\"] = [1] + prompt_tokens[\"prompt_attention_mask\"]\n        if chosen_prompt_len_input_ids == 0 or bos_token_id != chosen_tokens[\"prompt_input_ids\"][0]:\n            chosen_tokens[\"prompt_input_ids\"] = [bos_token_id] + chosen_tokens[\"prompt_input_ids\"]\n            chosen_tokens[\"prompt_attention_mask\"] = [1] + chosen_tokens[\"prompt_attention_mask\"]\n        if rejected_prompt_len_input_ids == 0 or bos_token_id != rejected_tokens[\"prompt_input_ids\"][0]:\n            rejected_tokens[\"prompt_input_ids\"] = [bos_token_id] + rejected_tokens[\"prompt_input_ids\"]\n            rejected_tokens[\"prompt_attention_mask\"] = [1] + rejected_tokens[\"prompt_attention_mask\"]\n    return prompt_tokens, chosen_tokens, rejected_tokens\n\n\ndef add_eos_token_if_needed(\n    eos_token_id: int, chosen_tokens: Dict[str, List[int]], rejected_tokens: Dict[str, List[int]]\n):\n    if len(chosen_tokens[\"input_ids\"]) == 0 or eos_token_id != chosen_tokens[\"input_ids\"][-1]:\n        chosen_tokens[\"input_ids\"].append(eos_token_id)\n        chosen_tokens[\"attention_mask\"].append(1)\n    if len(rejected_tokens[\"input_ids\"]) == 0 or eos_token_id != rejected_tokens[\"input_ids\"][-1]:\n        rejected_tokens[\"input_ids\"].append(eos_token_id)\n        rejected_tokens[\"attention_mask\"].append(1)\n    return chosen_tokens, rejected_tokens\n\n\ndef truncate_right(\n    input_ids: torch.Tensor, stop_token_id: int, pad_token_id: int\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Truncates the input tensor from the right side after the first occurrence of the stop token.\n\n    Args:\n        input_ids (`torch.Tensor`):\n            The tensor containing the responses to be truncated\n        stop_token_id (`int`):\n            The token ID representing the stop token where truncation occurs\n        pad_token_id (`int`):\n            The token ID representing the pad token used to fill the truncated responses\n\n    Returns:\n        tuple:\n            - `output_ids` (`torch.Tensor`):\n                The truncated responses tensor with pad tokens filled after the stop token\n            - `mask` (`torch.Tensor`):\n                The mask tensor to indicate the padding tokens\n    \"\"\"\n    trunc_idxs = first_true_indices(input_ids == stop_token_id).unsqueeze(-1)\n    new_size = [1] * (len(input_ids.size()) - 1) + [input_ids.shape[1]]\n    idxs = torch.arange(input_ids.shape[1], device=input_ids.device).view(*new_size)\n    output_ids = torch.masked_fill(input_ids, idxs > trunc_idxs, pad_token_id)\n    mask = torch.masked_fill(torch.ones_like(input_ids), idxs > trunc_idxs, 0)\n    return output_ids, mask\n\n\ndef empty_cache() -> None:\n    \"\"\"Empties the cache of the available torch device.\n\n    This function checks for the availability of different torch devices (XPU, MLU, NPU, CUDA)\n    and empties the cache of the first available device it finds.\n\n    If none of the specific devices are available, it defaults to emptying the CUDA cache.\n    \"\"\"\n    if is_torch_xpu_available():\n        torch.xpu.empty_cache()\n    elif is_torch_mlu_available():\n        torch.mlu.empty_cache()\n    elif is_torch_npu_available():\n        torch.npu.empty_cache()\n    else:\n        torch.cuda.empty_cache()\n\n\ndef decode_and_strip_padding(inputs: torch.Tensor, tokenizer: PreTrainedTokenizerBase) -> List[str]:\n    \"\"\"\n    Decodes the input tensor and strips the padding tokens.\n\n    Args:\n        inputs (`torch.Tensor`):\n            The input tensor to be decoded.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer used to decode the input tensor.\n\n    Returns:\n        `List[str]`:\n            The list of decoded strings with padding tokens stripped.\n    \"\"\"\n    decoded = tokenizer.batch_decode(inputs, skip_special_tokens=False)\n    return [d.replace(tokenizer.pad_token, \"\") for d in decoded]\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport warnings\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport datasets\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data\nfrom accelerate import PartialState\nfrom datasets import Dataset\nfrom packaging import version\nfrom torch.utils.data import DataLoader, IterableDataset\nfrom transformers import (\n    DataCollator,\n    GenerationConfig,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainerCallback,\n    is_apex_available,\n)\nfrom transformers.modeling_utils import PreTrainedModel\nfrom transformers.trainer_utils import EvalPrediction, seed_worker\nfrom transformers.training_args import OptimizerNames\nfrom transformers.utils import is_peft_available, is_sagemaker_mp_enabled, logging\n\nfrom ..data_utils import maybe_apply_chat_template\nfrom ..models import create_reference_model\nfrom ..models.utils import unwrap_model_for_generation\nfrom .judges import BasePairwiseJudge\nfrom .online_dpo_config import OnlineDPOConfig\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    disable_dropout_in_model,\n    empty_cache,\n    get_reward,\n    prepare_deepspeed,\n    trl_sanitze_kwargs_for_tagging,\n    truncate_right,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model\n\nif is_apex_available():\n    from apex import amp\n\n\nif is_sagemaker_mp_enabled():\n    from smdistributed.modelparallel import __version__ as SMP_VERSION\n\n    IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse(\"1.10\")\n\nelse:\n    IS_SAGEMAKER_MP_POST_1_10 = False\n\nlogger = logging.get_logger(__name__)\n\n\nclass OnlineDPOTrainer(Trainer):\n    r\"\"\"\n    Initialize OnlineDPOTrainer.\n\n    Args:\n        model (`transformers.PreTrainedModel` or `torch.nn.Module`):\n            The model to train, preferably an `AutoModelForCausalLM`.\n        ref_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`):\n            The reference model to use for training. If None is specified, the reference model will be created from\n            the model.\n        reward_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`):\n            The reward model to score completions with, preferably an `AutoModelForSequenceClassification`.\n        judge (`BasePairwiseJudge`):\n            The judge to use for pairwise comparison of model completions.\n        args (`OnlineDPOConfig`):\n            The online DPO config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        peft_config (`Dict`):\n            The peft config to use for training.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"online-dpo\"]\n\n    def __init__(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        ref_model: Union[PreTrainedModel, nn.Module, None] = None,\n        reward_model: Union[PreTrainedModel, nn.Module, None] = None,\n        judge: Optional[BasePairwiseJudge] = None,\n        args: Optional[OnlineDPOConfig] = None,\n        data_collator: Optional[DataCollator] = None,\n        train_dataset: Optional[Union[Dataset, IterableDataset, \"datasets.Dataset\"]] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], \"datasets.Dataset\"]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n    ) -> None:\n        if ref_model is model:\n            raise ValueError(\n                \"`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the \"\n                \"same as `model`, either omit the `ref_model` argument or pass `None`.\"\n            )\n\n        self.ref_model = ref_model\n\n        if reward_model is not None and judge is not None:\n            warnings.warn(\n                \"Both `reward_model` and `judge` are provided. Please choose provide only one of them. \"\n                \"Ignoring `judge` and using `reward_model`.\"\n            )\n        elif reward_model is None and judge is None:\n            raise ValueError(\"Either `reward_model` or `judge` must be provided.\")\n        elif reward_model is None and judge is not None:\n            raise NotImplementedError(\"Using `judge` is not yet supported.\")\n\n        self.reward_model = reward_model\n        self.judge = judge\n\n        if args is None:\n            raise ValueError(\"`args` must be provided.\")\n\n        # Check that the tokenizer is provided\n        if tokenizer is None:\n            raise ValueError(\"`tokenizer` must be provided.\")\n\n        # Convert to PEFT model if peft_config is provided\n        if peft_config is not None:\n            # Check if PEFT is available\n            if not is_peft_available():\n                raise ImportError(\n                    \"PEFT is not available and passed `peft_config`. Please install PEFT with \"\n                    \"`pip install peft` to use it.\"\n                )\n\n            # If the model is already a PeftModel, we need to merge and unload it.\n            # Further information here: https://huggingface.co/docs/trl/dpo_trainer#reference-model-considerations-with-peft\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            # Get peft model with the given config\n            model = get_peft_model(model, peft_config)\n\n        # Disable dropout in the model if specified\n        if args.disable_dropout:\n            disable_dropout_in_model(model)\n\n        # Handle the ref_model\n        # Usually, the user wants the ref model to be the initial version of the model. When using PEFT, it's easy to\n        # get the ref model, as it's just the model with a disabled adapter. When not using PEFT, we need to create\n        # the ref model from the model by copying it and disable the gradients and set it in evaluation mode.\n        if ref_model is None:  # No ref model provided, the most common case\n            if peft_config is None:\n                self.ref_model = create_reference_model(model)  # copy, disable gradients, set eval mode\n            else:\n                self.ref_model = None  # we don't need a ref model here, we can just disable the adapter.\n        else:  # rare case, the user provided a ref model\n            self.ref_model = ref_model\n            self.ref_model.eval()\n\n        # Disable the gradient and set the reward model in eval mode\n        if self.reward_model is not None:\n            self.reward_model.eval()\n\n        # Define the collator is not provided\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(pad_token_id=tokenizer.pad_token_id)\n\n        # Compute that only on the main process for faster data processing.\n        # see: https://github.com/huggingface/trl/pull/1255\n        with PartialState().local_main_process_first():\n            # Apply the chat template if needed\n            train_dataset = train_dataset.map(\n                maybe_apply_chat_template, fn_kwargs={\"tokenizer\": tokenizer}, num_proc=args.dataset_num_proc\n            )\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(\n                    maybe_apply_chat_template, fn_kwargs={\"tokenizer\": tokenizer}, num_proc=args.dataset_num_proc\n                )\n\n            # Tokenize the dataset\n            fn_kwargs = {\"is_encoder_decoder\": model.config.is_encoder_decoder, \"tokenizer\": tokenizer}\n            train_dataset = train_dataset.map(self.tokenize_row, fn_kwargs=fn_kwargs, num_proc=args.dataset_num_proc)\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(self.tokenize_row, fn_kwargs=fn_kwargs, num_proc=args.dataset_num_proc)\n\n        self.stats = {\n            \"objective/kl\": [],\n            \"objective/entropy\": [],\n            \"objective/non_score_reward\": [],\n            \"objective/rlhf_reward\": [],\n            \"objective/scores\": [],\n            \"objective/scores_margin\": [],\n            \"rewards/chosen\": [],\n            \"rewards/rejected\": [],\n            \"rewards/accuracies\": [],\n            \"rewards/margins\": [],\n            \"logps/chosen\": [],\n            \"logps/rejected\": [],\n            \"val/contain_eos_token\": [],\n            \"beta\": [],\n        }\n\n        self.generation_config = GenerationConfig(\n            max_new_tokens=args.max_new_tokens,\n            temperature=args.temperature,\n            top_k=0,\n            top_p=1.0,\n            do_sample=True,\n            use_cache=False if args.gradient_checkpointing else True,\n        )\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        self._beta = args.beta\n\n        # Placed after the super().__init__ because we need self.is_deepspeed_enabled and self.accelerator\n        if self.is_deepspeed_enabled:\n            if self.reward_model is not None:\n                self.reward_model = prepare_deepspeed(\n                    self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16\n                )\n            self.ref_model = prepare_deepspeed(self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16)\n        else:\n            if self.ref_model is not None:\n                self.ref_model = self.ref_model.to(self.accelerator.device)\n            if self.reward_model is not None:\n                self.reward_model = self.reward_model.to(self.accelerator.device)\n\n    @property\n    def beta(self):\n        if isinstance(self._beta, list):\n            epoch = self.state.epoch\n            return self._beta[epoch] if epoch < len(self._beta) else self._beta[-1]\n        else:\n            return self._beta\n\n    @staticmethod\n    def tokenize_row(feature, is_encoder_decoder: bool, tokenizer: PreTrainedTokenizerBase) -> Dict[str, Any]:\n        \"\"\"Tokenize a single row from a DPO specific dataset.\"\"\"\n        if not is_encoder_decoder:\n            batch = tokenizer(feature[\"prompt\"], add_special_tokens=False)\n            # Add BOS token to head of prompt. Avoid adding if it's already there\n            if tokenizer.bos_token_id is not None:\n                prompt_len_input_ids = len(batch[\"input_ids\"])\n                if prompt_len_input_ids == 0 or tokenizer.bos_token_id != batch[\"input_ids\"][0]:\n                    batch[\"input_ids\"] = [tokenizer.bos_token_id] + batch[\"input_ids\"]\n                    batch[\"attention_mask\"] = [1] + batch[\"attention_mask\"]\n        else:\n            batch = tokenizer(feature[\"prompt\"], add_special_tokens=True)\n        batch = {f\"prompt_{key}\": value for key, value in batch.items()}\n        return batch\n\n    # Same as Trainer.get_train_dataloader but skip the \"remove_unused_columns\".\n    @wraps(Trainer.get_train_dataloader)\n    def get_train_dataloader(self) -> DataLoader:\n        if self.train_dataset is None:\n            raise ValueError(\"Trainer: training requires a train_dataset.\")\n\n        train_dataset = self.train_dataset\n        data_collator = self.data_collator\n        dataloader_params = {\n            \"batch_size\": self._train_batch_size,\n            \"collate_fn\": data_collator,\n            \"num_workers\": self.args.dataloader_num_workers,\n            \"pin_memory\": self.args.dataloader_pin_memory,\n            \"persistent_workers\": self.args.dataloader_persistent_workers,\n        }\n\n        if not isinstance(train_dataset, torch.utils.data.IterableDataset):\n            dataloader_params[\"sampler\"] = self._get_train_sampler()\n            dataloader_params[\"drop_last\"] = self.args.dataloader_drop_last\n            dataloader_params[\"worker_init_fn\"] = seed_worker\n            dataloader_params[\"prefetch_factor\"] = self.args.dataloader_prefetch_factor\n\n        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))\n\n    # Same as Trainer.get_eval_dataloader but skip the \"remove_unused_columns\".\n    @wraps(Trainer.get_eval_dataloader)\n    def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader:\n        if eval_dataset is None and self.eval_dataset is None:\n            raise ValueError(\"Trainer: evaluation requires an eval_dataset.\")\n\n        # If we have persistent workers, don't do a fork bomb especially as eval datasets\n        # don't change during training\n        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else \"eval\"\n        if (\n            hasattr(self, \"_eval_dataloaders\")\n            and dataloader_key in self._eval_dataloaders\n            and self.args.dataloader_persistent_workers\n        ):\n            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])\n\n        eval_dataset = (\n            self.eval_dataset[eval_dataset]\n            if isinstance(eval_dataset, str)\n            else eval_dataset\n            if eval_dataset is not None\n            else self.eval_dataset\n        )\n        data_collator = self.data_collator\n\n        dataloader_params = {\n            \"batch_size\": self.args.eval_batch_size,\n            \"collate_fn\": data_collator,\n            \"num_workers\": self.args.dataloader_num_workers,\n            \"pin_memory\": self.args.dataloader_pin_memory,\n            \"persistent_workers\": self.args.dataloader_persistent_workers,\n        }\n\n        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):\n            dataloader_params[\"sampler\"] = self._get_eval_sampler(eval_dataset)\n            dataloader_params[\"drop_last\"] = self.args.dataloader_drop_last\n            dataloader_params[\"prefetch_factor\"] = self.args.dataloader_prefetch_factor\n\n        # accelerator.free_memory() will destroy the references, so\n        # we need to store the non-prepared version\n        eval_dataloader = DataLoader(eval_dataset, **dataloader_params)\n        if self.args.dataloader_persistent_workers:\n            if hasattr(self, \"_eval_dataloaders\"):\n                self._eval_dataloaders[dataloader_key] = eval_dataloader\n            else:\n                self._eval_dataloaders = {dataloader_key: eval_dataloader}\n\n        return self.accelerator.prepare(eval_dataloader)\n\n    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:\n        model.train()\n\n        # Sample 2 completations per prompt of size `max_new_tokens` from the model\n        inputs = self._prepare_inputs(inputs)\n        num_examples, context_length = inputs[\"prompt_input_ids\"].shape\n        prompt_ids = inputs[\"prompt_input_ids\"].repeat(2, 1)\n        prompt_mask = inputs[\"prompt_attention_mask\"].repeat(2, 1)\n        with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n            output = unwrapped_model.generate(\n                input_ids=prompt_ids,\n                attention_mask=prompt_mask,\n                generation_config=self.generation_config,\n            )\n        del inputs\n\n        completion_ids = output[:, context_length:]\n        completion_ids, completion_mask = truncate_right(\n            completion_ids, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id\n        )\n        prompt_completion_ids = torch.cat((prompt_ids, completion_ids), dim=1)\n        prompt_completion_mask = torch.cat((prompt_mask, completion_mask), dim=1)\n\n        # Get the logprobs of the completions from the model\n        output = model(prompt_completion_ids, attention_mask=prompt_completion_mask)\n        # There is 1 offset, because the model predict the next token\n        logits = output.logits[:, context_length - 1 : -1]\n        # Turn logits into logprobs\n        all_logprobs = F.log_softmax(logits, dim=-1)\n        # Take the completion tokens logprob\n        logprobs = torch.take_along_dim(all_logprobs, completion_ids.unsqueeze(-1), dim=2).squeeze(-1)\n        del output, logits, all_logprobs  # free memory\n\n        # Same for the reference model\n        with torch.no_grad():\n            if self.ref_model is not None:\n                ref_output = self.ref_model(prompt_completion_ids, attention_mask=prompt_completion_mask)\n            else:  # peft case: we just need to disable the adapter\n                with self.model.disable_adapter():\n                    ref_output = self.model(prompt_completion_ids, attention_mask=prompt_completion_mask)\n            ref_logits = ref_output.logits[:, context_length - 1 : -1]\n            ref_all_logprobs = F.log_softmax(ref_logits, dim=-1)\n            ref_logprobs = torch.take_along_dim(ref_all_logprobs, completion_ids.unsqueeze(-1), dim=2).squeeze(-1)\n            del ref_output, ref_logits, ref_all_logprobs  # free memory\n\n            # Get the reward from the reward model\n            _, scores, _ = get_reward(\n                self.reward_model, prompt_completion_ids, self.tokenizer.pad_token_id, context_length\n            )\n\n        # Filter completion. Ensure that the sample contains stop_token_id\n        # Completions not passing that filter will receive a lower score.\n        contain_eos_token = torch.any(completion_ids == self.tokenizer.eos_token_id, dim=-1)\n        if self.args.missing_eos_penalty is not None:\n            scores[~contain_eos_token] -= self.args.missing_eos_penalty\n\n        # Split the scores in 2 (the prompts of the first half are the same as the second half)\n        first_half, second_half = scores.split(num_examples)\n\n        # Get the indices of the chosen and rejected examples\n        num_examples_range = torch.arange(num_examples, device=scores.device)\n        mask = first_half >= second_half\n        chosen_indices = num_examples_range + (~mask * num_examples)\n        rejected_indices = num_examples_range + (mask * num_examples)\n\n        # Build tensor so that the first half is the chosen examples and the second half the rejected examples\n        cr_indices = torch.cat((chosen_indices, rejected_indices), dim=0)  # cr = chosen and rejected\n        cr_logprobs = logprobs[cr_indices]\n        cr_ref_logprobs = ref_logprobs[cr_indices]\n\n        # mask out the padding tokens\n        padding_mask = ~completion_mask.bool()\n        cr_padding_mask = padding_mask[cr_indices]\n\n        cr_logprobs_sum = (cr_logprobs * ~cr_padding_mask).sum(1)\n        cr_ref_logprobs_sum = (cr_ref_logprobs * ~cr_padding_mask).sum(1)\n\n        # Split the chosen and rejected examples\n        chosen_logprobs_sum, rejected_logprobs_sum = torch.split(cr_logprobs_sum, num_examples)\n        chosen_ref_logprobs_sum, rejected_ref_logprobs_sum = torch.split(cr_ref_logprobs_sum, num_examples)\n        pi_logratios = chosen_logprobs_sum - rejected_logprobs_sum\n        ref_logratios = chosen_ref_logprobs_sum - rejected_ref_logprobs_sum\n\n        logits = pi_logratios - ref_logratios\n\n        if self.args.loss_type == \"sigmoid\":\n            losses = -F.logsigmoid(self.beta * logits)\n        elif self.args.loss_type == \"ipo\":\n            losses = (logits - 1 / (2 * self.beta)) ** 2\n        else:\n            raise NotImplementedError(f\"invalid loss type {self.loss_type}\")\n\n        loss = losses.mean()\n\n        # Log everything\n        self.stats[\"val/contain_eos_token\"].append(contain_eos_token.float().mean().item())\n        self.stats[\"logps/chosen\"].append(self.accelerator.gather(chosen_logprobs_sum).mean().item())\n        self.stats[\"logps/rejected\"].append(self.accelerator.gather(rejected_logprobs_sum).mean().item())\n        self.stats[\"objective/scores\"].append(self.accelerator.gather(scores.mean()).mean().item())\n        kl = logprobs - ref_logprobs\n        mean_kl = kl.sum(1).mean()\n        self.stats[\"objective/kl\"].append(self.accelerator.gather(mean_kl).mean().item())\n        non_score_reward = (-self.beta * kl).sum(1)\n        mean_non_score_reward = non_score_reward.mean()\n        self.stats[\"objective/non_score_reward\"].append(self.accelerator.gather(mean_non_score_reward).mean().item())\n        rlhf_reward = scores + non_score_reward\n        self.stats[\"objective/rlhf_reward\"].append(self.accelerator.gather(rlhf_reward).mean().item())\n        mean_entropy = -logprobs.sum(1).mean()\n        self.stats[\"objective/entropy\"].append(self.accelerator.gather(mean_entropy).mean().item())\n        scores_margin = scores[chosen_indices] - scores[rejected_indices]\n        self.stats[\"objective/scores_margin\"].append(self.accelerator.gather(scores_margin.mean()).mean().item())\n        chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum)\n        gathered_chosen_rewards = self.accelerator.gather(chosen_rewards)\n        self.stats[\"rewards/chosen\"].append(gathered_chosen_rewards.mean().item())\n        rejected_rewards = self.beta * (rejected_logprobs_sum - rejected_ref_logprobs_sum)\n        gathered_rejected_rewards = self.accelerator.gather(rejected_rewards)\n        self.stats[\"rewards/rejected\"].append(gathered_rejected_rewards.mean().item())\n        margin = gathered_chosen_rewards - gathered_rejected_rewards\n        self.stats[\"rewards/margins\"].append(margin.mean().item())\n        accuracy = margin > 0\n        self.stats[\"rewards/accuracies\"].append(accuracy.float().mean().item())\n        self.stats[\"beta\"].append(self.beta)\n\n        if (\n            self.args.torch_empty_cache_steps is not None\n            and self.state.global_step % self.args.torch_empty_cache_steps == 0\n        ):\n            empty_cache()\n\n        kwargs = {}\n\n        # For LOMO optimizers you need to explicitly use the learnign rate\n        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:\n            kwargs[\"learning_rate\"] = self._get_learning_rate()\n\n        if self.args.n_gpu > 1:\n            loss = loss.mean()  # mean() to average on multi-gpu parallel training\n\n        if self.use_apex:\n            with amp.scale_loss(loss, self.optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            self.accelerator.backward(loss, **kwargs)\n\n        return loss.detach() / self.args.gradient_accumulation_steps\n\n    # Same as Trainer.evaluate but log our metrics\n    def _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval):\n        if self.control.should_log and self.state.global_step > self._globalstep_last_logged:\n            logs: Dict[str, float] = {}\n\n            # all_gather + mean() to get average loss over all processes\n            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()\n\n            # reset tr_loss to zero\n            tr_loss -= tr_loss\n\n            logs[\"loss\"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)\n            if grad_norm is not None:\n                logs[\"grad_norm\"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm\n            logs[\"learning_rate\"] = self._get_learning_rate()\n\n            # Add our metrics\n            for key, val in self.stats.items():\n                logs[key] = sum(val) / len(val)\n            self.stats = {key: [] for key in self.stats}  # reset stats\n\n            self._total_loss_scalar += tr_loss_scalar\n            self._globalstep_last_logged = self.state.global_step\n            self.store_flos()\n\n            self.log(logs)\n\n        metrics = None\n        if self.control.should_evaluate:\n            metrics = self._evaluate(trial, ignore_keys_for_eval)\n\n        if self.control.should_save:\n            self._save_checkpoint(model, trial, metrics=metrics)\n            self.control = self.callback_handler.on_save(self.args, self.state, self.control)\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"online-dpo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass BCOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`BCOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        max_prompt_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the prompt. This argument is required if you want to use the default data collator.\n        max_completion_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the completion. This argument is required if you want to use the default data collator\n            and your model is an encoder-decoder.\n        beta (`float`, *optional*, defaults to `0.1`):\n            Parameter controlling the deviation from the reference model. Higher β means less deviation from the\n            reference model.\n        label_pad_token_id (`int`,  *optional*, defaults to `-100`):\n            Label pad token id. This argument is required if you want to use the default data collator.\n        padding_value (`Optional[int]`, *optional*, defaults to `None`):\n            Padding value to use. If `None`, the padding value of the tokenizer is used.\n        truncation_mode (`str`, *optional*, defaults to `\"keep_end\"`):\n            Truncation mode to use when the prompt is too long. Possible values are `\"keep_end\"` or `\"keep_start\"`.\n            This argument is required if you want to use the default data collator.\n        generate_during_eval (`bool`, *optional*, defaults to `False`):\n            If `True`, generates and logs completions from both the model and the reference model to W&B during\n            evaluation.\n        is_encoder_decoder (`Optional[bool]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):\n            Whether to precompute reference model log probabilities for training and evaluation datasets. This is\n            useful when training without the reference model to reduce the total GPU memory needed.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        ref_model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model\n            from a string.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n        prompt_sample_size (`int`, *optional*, defaults to `1024`):\n            Number of prompts that are fed to density ratio classifier.\n        min_density_ratio (`float`, *optional*, defaults to `0.5`):\n            Minimum value of the density ratio. The estimated density ratio is clamped to this value.\n        max_density_ratio (`float`, *optional*, defaults to `10.0`):\n            Maximum value of the density ratio. The estimated density ratio is clamped to this value.\n    \"\"\"\n\n    max_length: Optional[int] = None\n    max_prompt_length: Optional[int] = None\n    max_completion_length: Optional[int] = None\n    beta: float = 0.1\n    label_pad_token_id: int = -100\n    padding_value: Optional[int] = None\n    truncation_mode: str = \"keep_end\"\n    generate_during_eval: bool = False\n    is_encoder_decoder: Optional[bool] = None\n    precompute_ref_log_probs: bool = False\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    ref_model_init_kwargs: Optional[Dict[str, Any]] = None\n    dataset_num_proc: Optional[int] = None\n    prompt_sample_size: int = 1024\n    min_density_ratio: float = 0.5\n    max_density_ratio: float = 10.0\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass RewardConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`RewardTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        dataset_num_proc (`int`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n        center_rewards_coefficient (`float`, *optional*, defaults to `None`):\n            Coefficient to incentivize the reward model to output mean-zero rewards (proposed by\n            https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`.\n    \"\"\"\n\n    max_length: Optional[int] = None\n    dataset_num_proc: Optional[int] = None\n    center_rewards_coefficient: Optional[float] = None\n\n\n# Copyright 2023 AlignProp-pytorch authors (Mihir Prabhudesai), metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport warnings\nfrom collections import defaultdict\nfrom typing import Any, Callable, Optional, Tuple\nfrom warnings import warn\n\nimport torch\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import ProjectConfiguration, set_seed\nfrom huggingface_hub import whoami\n\nfrom ..models import DDPOStableDiffusionPipeline\nfrom . import AlignPropConfig, BaseTrainer\n\n\nlogger = get_logger(__name__)\n\nMODEL_CARD_TEMPLATE = \"\"\"---\nlicense: apache-2.0\nlibrary_name: transformers\ntags:\n- trl\n- alignprop\n- diffusers\n- reinforcement-learning\n- text-to-image\n- stable-diffusion\n---\n\n# {model_name}\n\nThis is a pipeline that finetunes a diffusion model with reward backpropagation while using randomized truncation (https://huggingface.co/papers/2310.03739). The model can be used for image generation conditioned with text.\n\n\"\"\"\n\n\nclass AlignPropTrainer(BaseTrainer):\n    \"\"\"\n    The AlignPropTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models.\n    Note, this trainer is heavily inspired by the work here: https://github.com/mihirp1998/AlignProp/\n    As of now only Stable Diffusion based pipelines are supported\n\n    Attributes:\n        config (`AlignPropConfig`):\n            Configuration object for AlignPropTrainer. Check the documentation of `PPOConfig` for more details.\n        reward_function (`Callable[[torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor]`):\n            Reward function to be used\n        prompt_function (`Callable[[], Tuple[str, Any]]`):\n            Function to generate prompts to guide model\n        sd_pipeline (`DDPOStableDiffusionPipeline`):\n            Stable Diffusion pipeline to be used for training.\n        image_samples_hook (`Optional[Callable[[Any, Any, Any], Any]]`):\n            Hook to be called to log images\n    \"\"\"\n\n    _tag_names = [\"trl\", \"alignprop\"]\n\n    def __init__(\n        self,\n        config: AlignPropConfig,\n        reward_function: Callable[[torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor],\n        prompt_function: Callable[[], Tuple[str, Any]],\n        sd_pipeline: DDPOStableDiffusionPipeline,\n        image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None,\n    ):\n        if image_samples_hook is None:\n            warn(\"No image_samples_hook provided; no images will be logged\")\n\n        self.prompt_fn = prompt_function\n        self.reward_fn = reward_function\n        self.config = config\n        self.image_samples_callback = image_samples_hook\n\n        accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs)\n\n        if self.config.resume_from:\n            self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from))\n            if \"checkpoint_\" not in os.path.basename(self.config.resume_from):\n                # get the most recent checkpoint in this directory\n                checkpoints = list(\n                    filter(\n                        lambda x: \"checkpoint_\" in x,\n                        os.listdir(self.config.resume_from),\n                    )\n                )\n                if len(checkpoints) == 0:\n                    raise ValueError(f\"No checkpoints found in {self.config.resume_from}\")\n                checkpoint_numbers = sorted([int(x.split(\"_\")[-1]) for x in checkpoints])\n                self.config.resume_from = os.path.join(\n                    self.config.resume_from,\n                    f\"checkpoint_{checkpoint_numbers[-1]}\",\n                )\n\n                accelerator_project_config.iteration = checkpoint_numbers[-1] + 1\n\n        self.accelerator = Accelerator(\n            log_with=self.config.log_with,\n            mixed_precision=self.config.mixed_precision,\n            project_config=accelerator_project_config,\n            # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the\n            # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get\n            # the total number of optimizer steps to accumulate across.\n            gradient_accumulation_steps=self.config.train_gradient_accumulation_steps,\n            **self.config.accelerator_kwargs,\n        )\n\n        is_using_tensorboard = config.log_with is not None and config.log_with == \"tensorboard\"\n\n        if self.accelerator.is_main_process:\n            self.accelerator.init_trackers(\n                self.config.tracker_project_name,\n                config=dict(alignprop_trainer_config=config.to_dict())\n                if not is_using_tensorboard\n                else config.to_dict(),\n                init_kwargs=self.config.tracker_kwargs,\n            )\n\n        logger.info(f\"\\n{config}\")\n\n        set_seed(self.config.seed, device_specific=True)\n\n        self.sd_pipeline = sd_pipeline\n\n        self.sd_pipeline.set_progress_bar_config(\n            position=1,\n            disable=not self.accelerator.is_local_main_process,\n            leave=False,\n            desc=\"Timestep\",\n            dynamic_ncols=True,\n        )\n\n        # For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision\n        # as these weights are only used for inference, keeping weights in full precision is not required.\n        if self.accelerator.mixed_precision == \"fp16\":\n            inference_dtype = torch.float16\n        elif self.accelerator.mixed_precision == \"bf16\":\n            inference_dtype = torch.bfloat16\n        else:\n            inference_dtype = torch.float32\n\n        self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype)\n        self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype)\n        self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype)\n\n        trainable_layers = self.sd_pipeline.get_trainable_layers()\n\n        self.accelerator.register_save_state_pre_hook(self._save_model_hook)\n        self.accelerator.register_load_state_pre_hook(self._load_model_hook)\n\n        # Enable TF32 for faster training on Ampere GPUs,\n        # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\n        if self.config.allow_tf32:\n            torch.backends.cuda.matmul.allow_tf32 = True\n\n        self.optimizer = self._setup_optimizer(\n            trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers\n        )\n\n        self.neg_prompt_embed = self.sd_pipeline.text_encoder(\n            self.sd_pipeline.tokenizer(\n                [\"\"] if self.config.negative_prompts is None else self.config.negative_prompts,\n                return_tensors=\"pt\",\n                padding=\"max_length\",\n                truncation=True,\n                max_length=self.sd_pipeline.tokenizer.model_max_length,\n            ).input_ids.to(self.accelerator.device)\n        )[0]\n\n        # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses\n        # more memory\n        self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast\n\n        if hasattr(self.sd_pipeline, \"use_lora\") and self.sd_pipeline.use_lora:\n            unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer)\n            self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters()))\n        else:\n            self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer)\n\n        if config.resume_from:\n            logger.info(f\"Resuming from {config.resume_from}\")\n            self.accelerator.load_state(config.resume_from)\n            self.first_epoch = int(config.resume_from.split(\"_\")[-1]) + 1\n        else:\n            self.first_epoch = 0\n\n    def compute_rewards(self, prompt_image_pairs):\n        reward, reward_metadata = self.reward_fn(\n            prompt_image_pairs[\"images\"], prompt_image_pairs[\"prompts\"], prompt_image_pairs[\"prompt_metadata\"]\n        )\n        return reward\n\n    def step(self, epoch: int, global_step: int):\n        \"\"\"\n        Perform a single step of training.\n\n        Args:\n            epoch (int): The current epoch.\n            global_step (int): The current global step.\n\n        Side Effects:\n            - Model weights are updated\n            - Logs the statistics to the accelerator trackers.\n            - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, and the accelerator tracker.\n\n        Returns:\n            global_step (int): The updated global step.\n        \"\"\"\n        info = defaultdict(list)\n\n        self.sd_pipeline.unet.train()\n\n        for _ in range(self.config.train_gradient_accumulation_steps):\n            with self.accelerator.accumulate(self.sd_pipeline.unet), self.autocast(), torch.enable_grad():\n                prompt_image_pairs = self._generate_samples(\n                    batch_size=self.config.train_batch_size,\n                )\n\n                rewards = self.compute_rewards(prompt_image_pairs)\n\n                prompt_image_pairs[\"rewards\"] = rewards\n\n                rewards_vis = self.accelerator.gather(rewards).detach().cpu().numpy()\n\n                loss = self.calculate_loss(rewards)\n\n                self.accelerator.backward(loss)\n\n                if self.accelerator.sync_gradients:\n                    self.accelerator.clip_grad_norm_(\n                        self.trainable_layers.parameters()\n                        if not isinstance(self.trainable_layers, list)\n                        else self.trainable_layers,\n                        self.config.train_max_grad_norm,\n                    )\n\n                self.optimizer.step()\n                self.optimizer.zero_grad()\n\n            info[\"reward_mean\"].append(rewards_vis.mean())\n            info[\"reward_std\"].append(rewards_vis.std())\n            info[\"loss\"].append(loss.item())\n\n        # Checks if the accelerator has performed an optimization step behind the scenes\n        if self.accelerator.sync_gradients:\n            # log training-related stuff\n            info = {k: torch.mean(torch.tensor(v)) for k, v in info.items()}\n            info = self.accelerator.reduce(info, reduction=\"mean\")\n            info.update({\"epoch\": epoch})\n            self.accelerator.log(info, step=global_step)\n            global_step += 1\n            info = defaultdict(list)\n        else:\n            raise ValueError(\n                \"Optimization step should have been performed by this point. Please check calculated gradient accumulation settings.\"\n            )\n        # Logs generated images\n        if self.image_samples_callback is not None and global_step % self.config.log_image_freq == 0:\n            self.image_samples_callback(prompt_image_pairs, global_step, self.accelerator.trackers[0])\n\n        if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process:\n            self.accelerator.save_state()\n\n        return global_step\n\n    def calculate_loss(self, rewards):\n        \"\"\"\n        Calculate the loss for a batch of an unpacked sample\n\n        Args:\n            rewards (torch.Tensor):\n                Differentiable reward scalars for each generated image, shape: [batch_size]\n\n        Returns:\n            loss (torch.Tensor)\n            (all of these are of shape (1,))\n        \"\"\"\n        #  Loss is specific to Aesthetic Reward function used in AlignProp (https://huggingface.co/papers/2310.03739)\n        loss = 10.0 - (rewards).mean()\n        return loss\n\n    def loss(\n        self,\n        advantages: torch.Tensor,\n        clip_range: float,\n        ratio: torch.Tensor,\n    ):\n        unclipped_loss = -advantages * ratio\n        clipped_loss = -advantages * torch.clamp(\n            ratio,\n            1.0 - clip_range,\n            1.0 + clip_range,\n        )\n        return torch.mean(torch.maximum(unclipped_loss, clipped_loss))\n\n    def _setup_optimizer(self, trainable_layers_parameters):\n        if self.config.train_use_8bit_adam:\n            import bitsandbytes\n\n            optimizer_cls = bitsandbytes.optim.AdamW8bit\n        else:\n            optimizer_cls = torch.optim.AdamW\n\n        return optimizer_cls(\n            trainable_layers_parameters,\n            lr=self.config.train_learning_rate,\n            betas=(self.config.train_adam_beta1, self.config.train_adam_beta2),\n            weight_decay=self.config.train_adam_weight_decay,\n            eps=self.config.train_adam_epsilon,\n        )\n\n    def _save_model_hook(self, models, weights, output_dir):\n        self.sd_pipeline.save_checkpoint(models, weights, output_dir)\n        weights.pop()  # ensures that accelerate doesn't try to handle saving of the model\n\n    def _load_model_hook(self, models, input_dir):\n        self.sd_pipeline.load_checkpoint(models, input_dir)\n        models.pop()  # ensures that accelerate doesn't try to handle loading of the model\n\n    def _generate_samples(self, batch_size, with_grad=True, prompts=None):\n        \"\"\"\n        Generate samples from the model\n\n        Args:\n            batch_size (int): Batch size to use for sampling\n            with_grad (bool): Whether the generated RGBs should have gradients attached to it.\n\n        Returns:\n            prompt_image_pairs (Dict[Any])\n        \"\"\"\n        prompt_image_pairs = {}\n\n        sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1)\n\n        if prompts is None:\n            prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)])\n        else:\n            prompt_metadata = [{} for _ in range(batch_size)]\n\n        prompt_ids = self.sd_pipeline.tokenizer(\n            prompts,\n            return_tensors=\"pt\",\n            padding=\"max_length\",\n            truncation=True,\n            max_length=self.sd_pipeline.tokenizer.model_max_length,\n        ).input_ids.to(self.accelerator.device)\n\n        prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0]\n\n        if with_grad:\n            sd_output = self.sd_pipeline.rgb_with_grad(\n                prompt_embeds=prompt_embeds,\n                negative_prompt_embeds=sample_neg_prompt_embeds,\n                num_inference_steps=self.config.sample_num_steps,\n                guidance_scale=self.config.sample_guidance_scale,\n                eta=self.config.sample_eta,\n                truncated_backprop_rand=self.config.truncated_backprop_rand,\n                truncated_backprop_timestep=self.config.truncated_backprop_timestep,\n                truncated_rand_backprop_minmax=self.config.truncated_rand_backprop_minmax,\n                output_type=\"pt\",\n            )\n        else:\n            sd_output = self.sd_pipeline(\n                prompt_embeds=prompt_embeds,\n                negative_prompt_embeds=sample_neg_prompt_embeds,\n                num_inference_steps=self.config.sample_num_steps,\n                guidance_scale=self.config.sample_guidance_scale,\n                eta=self.config.sample_eta,\n                output_type=\"pt\",\n            )\n\n        images = sd_output.images\n\n        prompt_image_pairs[\"images\"] = images\n        prompt_image_pairs[\"prompts\"] = prompts\n        prompt_image_pairs[\"prompt_metadata\"] = prompt_metadata\n\n        return prompt_image_pairs\n\n    def train(self, epochs: Optional[int] = None):\n        \"\"\"\n        Train the model for a given number of epochs\n        \"\"\"\n        global_step = 0\n        if epochs is None:\n            epochs = self.config.num_epochs\n        for epoch in range(self.first_epoch, epochs):\n            global_step = self.step(epoch, global_step)\n\n    def create_model_card(self, path: str, model_name: Optional[str] = \"TRL AlignProp Model\") -> None:\n        \"\"\"Creates and saves a model card for a TRL model.\n\n        Args:\n            path (`str`): The path to save the model card to.\n            model_name (`str`, *optional*): The name of the model, defaults to `TRL AlignProp Model`.\n        \"\"\"\n        try:\n            user = whoami()[\"name\"]\n        # handle the offline case\n        except Exception:\n            warnings.warn(\"Cannot retrieve user information assuming you are running in offline mode.\")\n            return\n\n        if not os.path.exists(path):\n            os.makedirs(path)\n\n        model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, model_id=f\"{user}/{path}\")\n        with open(os.path.join(path, \"README.md\"), \"w\", encoding=\"utf-8\") as f:\n            f.write(model_card_content)\n\n    def _save_pretrained(self, save_directory):\n        self.sd_pipeline.save_pretrained(save_directory)\n        self.create_model_card(save_directory)\n\n\n# Copyright 2023 DDPO-pytorch authors (Kevin Black), metric-space, The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport warnings\nfrom collections import defaultdict\nfrom concurrent import futures\nfrom typing import Any, Callable, Optional, Tuple\nfrom warnings import warn\n\nimport torch\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import ProjectConfiguration, set_seed\nfrom huggingface_hub import whoami\n\nfrom ..models import DDPOStableDiffusionPipeline\nfrom . import BaseTrainer, DDPOConfig\nfrom .utils import PerPromptStatTracker\n\n\nlogger = get_logger(__name__)\n\nMODEL_CARD_TEMPLATE = \"\"\"---\nlicense: apache-2.0\nlibrary_name: transformers\ntags:\n- trl\n- ddpo\n- diffusers\n- reinforcement-learning\n- text-to-image\n- stable-diffusion\n---\n\n# {model_name}\n\nThis is a diffusion model that has been fine-tuned with reinforcement learning to\n guide the model outputs according to a value, function, or human feedback. The model can be used for image generation conditioned with text.\n\n\"\"\"\n\n\nclass DDPOTrainer(BaseTrainer):\n    \"\"\"\n    The DDPOTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models.\n    Note, this trainer is heavily inspired by the work here: https://github.com/kvablack/ddpo-pytorch\n    As of now only Stable Diffusion based pipelines are supported\n\n    Attributes:\n        **config** (`DDPOConfig`) -- Configuration object for DDPOTrainer. Check the documentation of `PPOConfig` for more\n         details.\n        **reward_function** (Callable[[torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor]) -- Reward function to be used\n        **prompt_function** (Callable[[], Tuple[str, Any]]) -- Function to generate prompts to guide model\n        **sd_pipeline** (`DDPOStableDiffusionPipeline`) -- Stable Diffusion pipeline to be used for training.\n        **image_samples_hook** (Optional[Callable[[Any, Any, Any], Any]]) -- Hook to be called to log images\n    \"\"\"\n\n    _tag_names = [\"trl\", \"ddpo\"]\n\n    def __init__(\n        self,\n        config: DDPOConfig,\n        reward_function: Callable[[torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor],\n        prompt_function: Callable[[], Tuple[str, Any]],\n        sd_pipeline: DDPOStableDiffusionPipeline,\n        image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None,\n    ):\n        if image_samples_hook is None:\n            warn(\"No image_samples_hook provided; no images will be logged\")\n\n        self.prompt_fn = prompt_function\n        self.reward_fn = reward_function\n        self.config = config\n        self.image_samples_callback = image_samples_hook\n\n        accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs)\n\n        if self.config.resume_from:\n            self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from))\n            if \"checkpoint_\" not in os.path.basename(self.config.resume_from):\n                # get the most recent checkpoint in this directory\n                checkpoints = list(\n                    filter(\n                        lambda x: \"checkpoint_\" in x,\n                        os.listdir(self.config.resume_from),\n                    )\n                )\n                if len(checkpoints) == 0:\n                    raise ValueError(f\"No checkpoints found in {self.config.resume_from}\")\n                checkpoint_numbers = sorted([int(x.split(\"_\")[-1]) for x in checkpoints])\n                self.config.resume_from = os.path.join(\n                    self.config.resume_from,\n                    f\"checkpoint_{checkpoint_numbers[-1]}\",\n                )\n\n                accelerator_project_config.iteration = checkpoint_numbers[-1] + 1\n\n        # number of timesteps within each trajectory to train on\n        self.num_train_timesteps = int(self.config.sample_num_steps * self.config.train_timestep_fraction)\n\n        self.accelerator = Accelerator(\n            log_with=self.config.log_with,\n            mixed_precision=self.config.mixed_precision,\n            project_config=accelerator_project_config,\n            # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the\n            # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get\n            # the total number of optimizer steps to accumulate across.\n            gradient_accumulation_steps=self.config.train_gradient_accumulation_steps * self.num_train_timesteps,\n            **self.config.accelerator_kwargs,\n        )\n\n        is_okay, message = self._config_check()\n        if not is_okay:\n            raise ValueError(message)\n\n        is_using_tensorboard = config.log_with is not None and config.log_with == \"tensorboard\"\n\n        if self.accelerator.is_main_process:\n            self.accelerator.init_trackers(\n                self.config.tracker_project_name,\n                config=dict(ddpo_trainer_config=config.to_dict()) if not is_using_tensorboard else config.to_dict(),\n                init_kwargs=self.config.tracker_kwargs,\n            )\n\n        logger.info(f\"\\n{config}\")\n\n        set_seed(self.config.seed, device_specific=True)\n\n        self.sd_pipeline = sd_pipeline\n\n        self.sd_pipeline.set_progress_bar_config(\n            position=1,\n            disable=not self.accelerator.is_local_main_process,\n            leave=False,\n            desc=\"Timestep\",\n            dynamic_ncols=True,\n        )\n\n        # For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision\n        # as these weights are only used for inference, keeping weights in full precision is not required.\n        if self.accelerator.mixed_precision == \"fp16\":\n            inference_dtype = torch.float16\n        elif self.accelerator.mixed_precision == \"bf16\":\n            inference_dtype = torch.bfloat16\n        else:\n            inference_dtype = torch.float32\n\n        self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype)\n        self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype)\n        self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype)\n\n        trainable_layers = self.sd_pipeline.get_trainable_layers()\n\n        self.accelerator.register_save_state_pre_hook(self._save_model_hook)\n        self.accelerator.register_load_state_pre_hook(self._load_model_hook)\n\n        # Enable TF32 for faster training on Ampere GPUs,\n        # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\n        if self.config.allow_tf32:\n            torch.backends.cuda.matmul.allow_tf32 = True\n\n        self.optimizer = self._setup_optimizer(\n            trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers\n        )\n\n        self.neg_prompt_embed = self.sd_pipeline.text_encoder(\n            self.sd_pipeline.tokenizer(\n                [\"\"] if self.config.negative_prompts is None else self.config.negative_prompts,\n                return_tensors=\"pt\",\n                padding=\"max_length\",\n                truncation=True,\n                max_length=self.sd_pipeline.tokenizer.model_max_length,\n            ).input_ids.to(self.accelerator.device)\n        )[0]\n\n        if config.per_prompt_stat_tracking:\n            self.stat_tracker = PerPromptStatTracker(\n                config.per_prompt_stat_tracking_buffer_size,\n                config.per_prompt_stat_tracking_min_count,\n            )\n\n        # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses\n        # more memory\n        self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast\n\n        if hasattr(self.sd_pipeline, \"use_lora\") and self.sd_pipeline.use_lora:\n            unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer)\n            self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters()))\n        else:\n            self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer)\n\n        if self.config.async_reward_computation:\n            self.executor = futures.ThreadPoolExecutor(max_workers=config.max_workers)\n\n        if config.resume_from:\n            logger.info(f\"Resuming from {config.resume_from}\")\n            self.accelerator.load_state(config.resume_from)\n            self.first_epoch = int(config.resume_from.split(\"_\")[-1]) + 1\n        else:\n            self.first_epoch = 0\n\n    def compute_rewards(self, prompt_image_pairs, is_async=False):\n        if not is_async:\n            rewards = []\n            for images, prompts, prompt_metadata in prompt_image_pairs:\n                reward, reward_metadata = self.reward_fn(images, prompts, prompt_metadata)\n                rewards.append(\n                    (\n                        torch.as_tensor(reward, device=self.accelerator.device),\n                        reward_metadata,\n                    )\n                )\n        else:\n            rewards = self.executor.map(lambda x: self.reward_fn(*x), prompt_image_pairs)\n            rewards = [\n                (torch.as_tensor(reward.result(), device=self.accelerator.device), reward_metadata.result())\n                for reward, reward_metadata in rewards\n            ]\n\n        return zip(*rewards)\n\n    def step(self, epoch: int, global_step: int):\n        \"\"\"\n        Perform a single step of training.\n\n        Args:\n            epoch (int): The current epoch.\n            global_step (int): The current global step.\n\n        Side Effects:\n            - Model weights are updated\n            - Logs the statistics to the accelerator trackers.\n            - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, and the accelerator tracker.\n\n        Returns:\n            global_step (int): The updated global step.\n\n        \"\"\"\n        samples, prompt_image_data = self._generate_samples(\n            iterations=self.config.sample_num_batches_per_epoch,\n            batch_size=self.config.sample_batch_size,\n        )\n\n        # collate samples into dict where each entry has shape (num_batches_per_epoch * sample.batch_size, ...)\n        samples = {k: torch.cat([s[k] for s in samples]) for k in samples[0].keys()}\n        rewards, rewards_metadata = self.compute_rewards(\n            prompt_image_data, is_async=self.config.async_reward_computation\n        )\n\n        for i, image_data in enumerate(prompt_image_data):\n            image_data.extend([rewards[i], rewards_metadata[i]])\n\n        if self.image_samples_callback is not None:\n            self.image_samples_callback(prompt_image_data, global_step, self.accelerator.trackers[0])\n\n        rewards = torch.cat(rewards)\n        rewards = self.accelerator.gather(rewards).cpu().numpy()\n\n        self.accelerator.log(\n            {\n                \"reward\": rewards,\n                \"epoch\": epoch,\n                \"reward_mean\": rewards.mean(),\n                \"reward_std\": rewards.std(),\n            },\n            step=global_step,\n        )\n\n        if self.config.per_prompt_stat_tracking:\n            # gather the prompts across processes\n            prompt_ids = self.accelerator.gather(samples[\"prompt_ids\"]).cpu().numpy()\n            prompts = self.sd_pipeline.tokenizer.batch_decode(prompt_ids, skip_special_tokens=True)\n            advantages = self.stat_tracker.update(prompts, rewards)\n        else:\n            advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)\n\n        # ungather advantages;  keep the entries corresponding to the samples on this process\n        samples[\"advantages\"] = (\n            torch.as_tensor(advantages)\n            .reshape(self.accelerator.num_processes, -1)[self.accelerator.process_index]\n            .to(self.accelerator.device)\n        )\n\n        del samples[\"prompt_ids\"]\n\n        total_batch_size, num_timesteps = samples[\"timesteps\"].shape\n\n        for inner_epoch in range(self.config.train_num_inner_epochs):\n            # shuffle samples along batch dimension\n            perm = torch.randperm(total_batch_size, device=self.accelerator.device)\n            samples = {k: v[perm] for k, v in samples.items()}\n\n            # shuffle along time dimension independently for each sample\n            # still trying to understand the code below\n            perms = torch.stack(\n                [torch.randperm(num_timesteps, device=self.accelerator.device) for _ in range(total_batch_size)]\n            )\n\n            for key in [\"timesteps\", \"latents\", \"next_latents\", \"log_probs\"]:\n                samples[key] = samples[key][\n                    torch.arange(total_batch_size, device=self.accelerator.device)[:, None],\n                    perms,\n                ]\n\n            original_keys = samples.keys()\n            original_values = samples.values()\n            # rebatch them as user defined train_batch_size is different from sample_batch_size\n            reshaped_values = [v.reshape(-1, self.config.train_batch_size, *v.shape[1:]) for v in original_values]\n\n            # Transpose the list of original values\n            transposed_values = zip(*reshaped_values)\n            # Create new dictionaries for each row of transposed values\n            samples_batched = [dict(zip(original_keys, row_values)) for row_values in transposed_values]\n\n            self.sd_pipeline.unet.train()\n            global_step = self._train_batched_samples(inner_epoch, epoch, global_step, samples_batched)\n            # ensure optimization step at the end of the inner epoch\n            if not self.accelerator.sync_gradients:\n                raise ValueError(\n                    \"Optimization step should have been performed by this point. Please check calculated gradient accumulation settings.\"\n                )\n\n        if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process:\n            self.accelerator.save_state()\n\n        return global_step\n\n    def calculate_loss(self, latents, timesteps, next_latents, log_probs, advantages, embeds):\n        \"\"\"\n        Calculate the loss for a batch of an unpacked sample\n\n        Args:\n            latents (torch.Tensor):\n                The latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width]\n            timesteps (torch.Tensor):\n                The timesteps sampled from the diffusion model, shape: [batch_size]\n            next_latents (torch.Tensor):\n                The next latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width]\n            log_probs (torch.Tensor):\n                The log probabilities of the latents, shape: [batch_size]\n            advantages (torch.Tensor):\n                The advantages of the latents, shape: [batch_size]\n            embeds (torch.Tensor):\n                The embeddings of the prompts, shape: [2*batch_size or batch_size, ...]\n                Note: the \"or\" is because if train_cfg is True, the expectation is that negative prompts are concatenated to the embeds\n\n        Returns:\n            loss (torch.Tensor), approx_kl (torch.Tensor), clipfrac (torch.Tensor)\n            (all of these are of shape (1,))\n        \"\"\"\n        with self.autocast():\n            if self.config.train_cfg:\n                noise_pred = self.sd_pipeline.unet(\n                    torch.cat([latents] * 2),\n                    torch.cat([timesteps] * 2),\n                    embeds,\n                ).sample\n                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n                noise_pred = noise_pred_uncond + self.config.sample_guidance_scale * (\n                    noise_pred_text - noise_pred_uncond\n                )\n            else:\n                noise_pred = self.sd_pipeline.unet(\n                    latents,\n                    timesteps,\n                    embeds,\n                ).sample\n            # compute the log prob of next_latents given latents under the current model\n\n            scheduler_step_output = self.sd_pipeline.scheduler_step(\n                noise_pred,\n                timesteps,\n                latents,\n                eta=self.config.sample_eta,\n                prev_sample=next_latents,\n            )\n\n            log_prob = scheduler_step_output.log_probs\n\n        advantages = torch.clamp(\n            advantages,\n            -self.config.train_adv_clip_max,\n            self.config.train_adv_clip_max,\n        )\n\n        ratio = torch.exp(log_prob - log_probs)\n\n        loss = self.loss(advantages, self.config.train_clip_range, ratio)\n\n        approx_kl = 0.5 * torch.mean((log_prob - log_probs) ** 2)\n\n        clipfrac = torch.mean((torch.abs(ratio - 1.0) > self.config.train_clip_range).float())\n\n        return loss, approx_kl, clipfrac\n\n    def loss(\n        self,\n        advantages: torch.Tensor,\n        clip_range: float,\n        ratio: torch.Tensor,\n    ):\n        unclipped_loss = -advantages * ratio\n        clipped_loss = -advantages * torch.clamp(\n            ratio,\n            1.0 - clip_range,\n            1.0 + clip_range,\n        )\n        return torch.mean(torch.maximum(unclipped_loss, clipped_loss))\n\n    def _setup_optimizer(self, trainable_layers_parameters):\n        if self.config.train_use_8bit_adam:\n            import bitsandbytes\n\n            optimizer_cls = bitsandbytes.optim.AdamW8bit\n        else:\n            optimizer_cls = torch.optim.AdamW\n\n        return optimizer_cls(\n            trainable_layers_parameters,\n            lr=self.config.train_learning_rate,\n            betas=(self.config.train_adam_beta1, self.config.train_adam_beta2),\n            weight_decay=self.config.train_adam_weight_decay,\n            eps=self.config.train_adam_epsilon,\n        )\n\n    def _save_model_hook(self, models, weights, output_dir):\n        self.sd_pipeline.save_checkpoint(models, weights, output_dir)\n        weights.pop()  # ensures that accelerate doesn't try to handle saving of the model\n\n    def _load_model_hook(self, models, input_dir):\n        self.sd_pipeline.load_checkpoint(models, input_dir)\n        models.pop()  # ensures that accelerate doesn't try to handle loading of the model\n\n    def _generate_samples(self, iterations, batch_size):\n        \"\"\"\n        Generate samples from the model\n\n        Args:\n            iterations (int): Number of iterations to generate samples for\n            batch_size (int): Batch size to use for sampling\n\n        Returns:\n            samples (List[Dict[str, torch.Tensor]]), prompt_image_pairs (List[List[Any]])\n        \"\"\"\n        samples = []\n        prompt_image_pairs = []\n        self.sd_pipeline.unet.eval()\n\n        sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1)\n\n        for _ in range(iterations):\n            prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)])\n\n            prompt_ids = self.sd_pipeline.tokenizer(\n                prompts,\n                return_tensors=\"pt\",\n                padding=\"max_length\",\n                truncation=True,\n                max_length=self.sd_pipeline.tokenizer.model_max_length,\n            ).input_ids.to(self.accelerator.device)\n            prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0]\n\n            with self.autocast():\n                sd_output = self.sd_pipeline(\n                    prompt_embeds=prompt_embeds,\n                    negative_prompt_embeds=sample_neg_prompt_embeds,\n                    num_inference_steps=self.config.sample_num_steps,\n                    guidance_scale=self.config.sample_guidance_scale,\n                    eta=self.config.sample_eta,\n                    output_type=\"pt\",\n                )\n\n                images = sd_output.images\n                latents = sd_output.latents\n                log_probs = sd_output.log_probs\n\n            latents = torch.stack(latents, dim=1)  # (batch_size, num_steps + 1, ...)\n            log_probs = torch.stack(log_probs, dim=1)  # (batch_size, num_steps, 1)\n            timesteps = self.sd_pipeline.scheduler.timesteps.repeat(batch_size, 1)  # (batch_size, num_steps)\n\n            samples.append(\n                {\n                    \"prompt_ids\": prompt_ids,\n                    \"prompt_embeds\": prompt_embeds,\n                    \"timesteps\": timesteps,\n                    \"latents\": latents[:, :-1],  # each entry is the latent before timestep t\n                    \"next_latents\": latents[:, 1:],  # each entry is the latent after timestep t\n                    \"log_probs\": log_probs,\n                    \"negative_prompt_embeds\": sample_neg_prompt_embeds,\n                }\n            )\n            prompt_image_pairs.append([images, prompts, prompt_metadata])\n\n        return samples, prompt_image_pairs\n\n    def _train_batched_samples(self, inner_epoch, epoch, global_step, batched_samples):\n        \"\"\"\n        Train on a batch of samples. Main training segment\n\n        Args:\n            inner_epoch (int): The current inner epoch\n            epoch (int): The current epoch\n            global_step (int): The current global step\n            batched_samples (List[Dict[str, torch.Tensor]]): The batched samples to train on\n\n        Side Effects:\n            - Model weights are updated\n            - Logs the statistics to the accelerator trackers.\n\n        Returns:\n            global_step (int): The updated global step\n        \"\"\"\n        info = defaultdict(list)\n        for _i, sample in enumerate(batched_samples):\n            if self.config.train_cfg:\n                # concat negative prompts to sample prompts to avoid two forward passes\n                embeds = torch.cat([sample[\"negative_prompt_embeds\"], sample[\"prompt_embeds\"]])\n            else:\n                embeds = sample[\"prompt_embeds\"]\n\n            for j in range(self.num_train_timesteps):\n                with self.accelerator.accumulate(self.sd_pipeline.unet):\n                    loss, approx_kl, clipfrac = self.calculate_loss(\n                        sample[\"latents\"][:, j],\n                        sample[\"timesteps\"][:, j],\n                        sample[\"next_latents\"][:, j],\n                        sample[\"log_probs\"][:, j],\n                        sample[\"advantages\"],\n                        embeds,\n                    )\n                    info[\"approx_kl\"].append(approx_kl)\n                    info[\"clipfrac\"].append(clipfrac)\n                    info[\"loss\"].append(loss)\n\n                    self.accelerator.backward(loss)\n                    if self.accelerator.sync_gradients:\n                        self.accelerator.clip_grad_norm_(\n                            self.trainable_layers.parameters()\n                            if not isinstance(self.trainable_layers, list)\n                            else self.trainable_layers,\n                            self.config.train_max_grad_norm,\n                        )\n                    self.optimizer.step()\n                    self.optimizer.zero_grad()\n\n                # Checks if the accelerator has performed an optimization step behind the scenes\n                if self.accelerator.sync_gradients:\n                    # log training-related stuff\n                    info = {k: torch.mean(torch.stack(v)) for k, v in info.items()}\n                    info = self.accelerator.reduce(info, reduction=\"mean\")\n                    info.update({\"epoch\": epoch, \"inner_epoch\": inner_epoch})\n                    self.accelerator.log(info, step=global_step)\n                    global_step += 1\n                    info = defaultdict(list)\n        return global_step\n\n    def _config_check(self) -> Tuple[bool, str]:\n        samples_per_epoch = (\n            self.config.sample_batch_size * self.accelerator.num_processes * self.config.sample_num_batches_per_epoch\n        )\n        total_train_batch_size = (\n            self.config.train_batch_size\n            * self.accelerator.num_processes\n            * self.config.train_gradient_accumulation_steps\n        )\n\n        if not self.config.sample_batch_size >= self.config.train_batch_size:\n            return (\n                False,\n                f\"Sample batch size ({self.config.sample_batch_size}) must be greater than or equal to the train batch size ({self.config.train_batch_size})\",\n            )\n        if not self.config.sample_batch_size % self.config.train_batch_size == 0:\n            return (\n                False,\n                f\"Sample batch size ({self.config.sample_batch_size}) must be divisible by the train batch size ({self.config.train_batch_size})\",\n            )\n        if not samples_per_epoch % total_train_batch_size == 0:\n            return (\n                False,\n                f\"Number of samples per epoch ({samples_per_epoch}) must be divisible by the total train batch size ({total_train_batch_size})\",\n            )\n        return True, \"\"\n\n    def train(self, epochs: Optional[int] = None):\n        \"\"\"\n        Train the model for a given number of epochs\n        \"\"\"\n        global_step = 0\n        if epochs is None:\n            epochs = self.config.num_epochs\n        for epoch in range(self.first_epoch, epochs):\n            global_step = self.step(epoch, global_step)\n\n    def create_model_card(self, path: str, model_name: Optional[str] = \"TRL DDPO Model\") -> None:\n        \"\"\"Creates and saves a model card for a TRL model.\n\n        Args:\n            path (`str`): The path to save the model card to.\n            model_name (`str`, *optional*): The name of the model, defaults to `TRL DDPO Model`.\n        \"\"\"\n        try:\n            user = whoami()[\"name\"]\n        # handle the offline case\n        except Exception:\n            warnings.warn(\"Cannot retrieve user information assuming you are running in offline mode.\")\n            return\n\n        if not os.path.exists(path):\n            os.makedirs(path)\n\n        model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, model_id=f\"{user}/{path}\")\n        with open(os.path.join(path, \"README.md\"), \"w\", encoding=\"utf-8\") as f:\n            f.write(model_card_content)\n\n    def _save_pretrained(self, save_directory):\n        self.sd_pipeline.save_pretrained(save_directory)\n        self.create_model_card(save_directory)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom dataclasses import dataclass\n\nfrom ..trainer.utils import OnPolicyConfig\n\n\n@dataclass\nclass RLOOConfig(OnPolicyConfig):\n    r\"\"\"\n    Configuration class for the [`RLOOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[: -len(\".py\")]`):\n            Name of this experiment.\n        reward_model_path (`str`, *optional*, defaults to `\"EleutherAI/pythia-160m\"`):\n            Path to the reward model.\n        num_ppo_epochs (`int`, *optional*, defaults to `4`):\n            Number of epochs to train.\n        whiten_rewards (`bool`, *optional*, defaults to `False`):\n            Whether to whiten the rewards.\n        kl_coef (`float`, *optional*, defaults to `0.05`):\n            KL coefficient.\n        cliprange (`float`, *optional*, defaults to `0.2`):\n            Clip range.\n        rloo_k (`int`, *optional*, defaults to `2`):\n            REINFORCE Leave-One-Out (RLOO) number of online samples per prompt.\n    \"\"\"\n\n    exp_name: str = os.path.basename(__file__)[: -len(\".py\")]\n    reward_model_path: str = \"EleutherAI/pythia-160m\"\n    num_ppo_epochs: int = 4\n    whiten_rewards: bool = False\n    kl_coef: float = 0.05\n    cliprange: float = 0.2\n    rloo_k: int = 2\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import List, Literal, Optional, Union\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass OnlineDPOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`OnlineDPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        learning_rate (`float`, *optional*, defaults to `5e-7`):\n            Initial learning rate for [`AdamW`] optimizer. The default value replaces that of\n            [`~transformers.TrainingArguments`].\n        reward_model_path (`Optional[str]`, *optional*, defaults to `None`):\n            Path to the reward model.\n        max_new_tokens (`int`, *optional*, defaults to `64`):\n            Maximum number of tokens to generate per completion.\n        temperature (`float`, *optional*, defaults to `0.9`):\n            Temperature for sampling. The higher the temperature, the more random the completions.\n        missing_eos_penalty (`Optional[float]`, *optional*, defaults to `None`):\n            Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage\n            to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive\n            value.\n        beta (`float` or `list[float]`, *optional*, defaults to `0.1`):\n            Parameter controlling the deviation from the reference model. Higher β means less deviation from the\n            reference model. For the IPO loss (`loss_type=\"ipo\"`), β is the regularization parameter denoted by τ in\n            the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β is\n            selected for each new epoch and the last β is used for the rest of the epochs.\n        loss_type (`str`, *optional*, defaults to `\"sigmoid\"`):\n            Type of loss to use. Possible values are:\n\n                - `\"sigmoid\"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper.\n                - `\"ipo\"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper.\n\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n        disable_dropout (`bool`, *optional*, defaults to `True`):\n            Whether to disable dropout in the model.\n    \"\"\"\n\n    learning_rate: float = 5e-7\n    reward_model_path: Optional[str] = None\n    max_new_tokens: int = 64\n    temperature: float = 0.9\n    missing_eos_penalty: Optional[float] = None\n    beta: Union[float, List[float]] = 0.1\n    loss_type: Literal[\"sigmoid\", \"ipo\"] = \"sigmoid\"\n    dataset_num_proc: Optional[int] = None\n    disable_dropout: bool = True\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport gc\nimport math\nimport os\nimport time\nfrom collections import defaultdict\nfrom functools import wraps\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import Accelerator\nfrom accelerate.utils import broadcast, gather_object\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    DataCollatorWithPadding,\n    GenerationConfig,\n    PreTrainedTokenizer,\n    Trainer,\n    TrainerCallback,\n    TrainerControl,\n)\nfrom transformers.integrations import get_reporting_integration_callbacks\nfrom transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK\nfrom transformers.trainer_callback import CallbackHandler, ExportableState, PrinterCallback\n\nfrom ..core import masked_mean, masked_whiten\nfrom ..models.utils import unwrap_model_for_generation\nfrom ..trainer.utils import (\n    OnlineTrainerState,\n    batch_generation,\n    disable_dropout_in_model,\n    exact_div,\n    first_true_indices,\n    forward,\n    get_reward,\n    prepare_deepspeed,\n    print_rich_table,\n    truncate_response,\n)\nfrom .ppov2_config import PPOv2Config\nfrom .utils import trl_sanitze_kwargs_for_tagging\n\n\nINVALID_LOGPROB = 1.0\n\n\n# taken from https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/ppo/ppo_trainer.py#L29\n# we did this we can do a single `model = accelerator.prepare(model)`\nclass PolicyAndValueWrapper(nn.Module):\n    def __init__(self, policy, value_model) -> None:\n        super().__init__()\n        self.policy = policy\n        self.value_model = value_model\n        self.critic_backbone = getattr(value_model, value_model.base_model_prefix)\n\n    def forward(self, **kwargs):\n        output = self.critic_backbone(\n            **kwargs,\n        )\n        logits = self.value_model.score(output.hidden_states[-1])\n        return self.policy(**kwargs), logits\n\n\nclass PPOv2Trainer(Trainer):\n    _tag_names = [\"trl\", \"ppo\"]\n\n    def __init__(\n        self,\n        config: PPOv2Config,\n        tokenizer: PreTrainedTokenizer,\n        policy: nn.Module,\n        ref_policy: nn.Module,\n        reward_model: nn.Module,\n        train_dataset: Dataset,\n        value_model: Optional[nn.Module] = None,\n        data_collator: Optional[DataCollatorWithPadding] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        # less commonly used\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        callbacks: Optional[List[TrainerCallback]] = None,\n    ) -> None:\n        if ref_policy is policy:\n            raise ValueError(\n                \"`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the \"\n                \"same as `policy`, you must mass a copy of it, or `None` if you use peft.\"\n            )\n\n        self.args = config\n        args = config\n        self.tokenizer = tokenizer\n        self.policy = policy\n\n        self.policy.generation_config.eos_token_id = (\n            None  # disable `pad_token_id` and `eos_token_id` because we just want to\n        )\n        self.policy.generation_config.pad_token_id = None  # generate tokens without truncation / padding\n\n        self.ref_policy = ref_policy\n        self.reward_model = reward_model\n        self.train_dataset = train_dataset\n        self.train_dataset_len = len(train_dataset)\n        self.value_model = value_model\n        self.data_collator = data_collator\n        self.eval_dataset = eval_dataset\n        self.optimizer, self.lr_scheduler = optimizers\n\n        #########\n        # calculate various batch sizes\n        #########\n        if args.total_episodes is None:  # allow the users to define episodes in terms of epochs.\n            args.total_episodes = int(args.num_train_epochs * self.train_dataset_len)\n        accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)\n        self.accelerator = accelerator\n        args.world_size = accelerator.num_processes\n        args.local_batch_size = (\n            args.per_device_train_batch_size * args.gradient_accumulation_steps * args.num_mini_batches\n        )\n        args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size)\n        args.batch_size = int(args.local_batch_size * args.world_size)\n        args.mini_batch_size = exact_div(\n            args.batch_size, args.num_mini_batches, \"`batch_size` must be a multiple of `num_mini_batches`\"\n        )\n        args.local_mini_batch_size = exact_div(\n            args.local_batch_size, args.num_mini_batches, \"`local_batch_size` must be a multiple of `num_mini_batches`\"\n        )\n        if args.whiten_rewards:\n            assert (\n                args.local_mini_batch_size >= 8\n            ), f\"Per-rank minibatch size {args.local_mini_batch_size} is insufficient for whitening\"\n        # `per_rank_rollout_batch_size` is our `args.local_batch_size`\n        # `per_rank_minibatch_size` is our `args.local_mini_batch_size`\n        args.num_total_batches = math.ceil(\n            args.total_episodes / args.batch_size\n        )  # we may train for more than `total_episodes`\n        time_tensor = torch.tensor(int(time.time()), device=accelerator.device)\n        time_int = broadcast(time_tensor, 0).item()  # avoid different timestamps across processes\n        args.run_name = f\"{args.exp_name}__{args.seed}__{time_int}\"\n        self.local_seed = args.seed + accelerator.process_index * 100003  # Prime\n        if args.num_sample_generations > 0:\n            self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations)\n        self.local_dataloader_batch_size = args.local_batch_size\n\n        #########\n        # setup model, optimizer, and others\n        #########\n        for module in [policy, ref_policy, value_model, reward_model]:\n            disable_dropout_in_model(module)\n        if args.stop_token and args.stop_token == \"eos\":\n            args.stop_token_id = tokenizer.eos_token_id\n        self.model = PolicyAndValueWrapper(policy, value_model)\n        self.model.config = policy.config  # needed for pushing to hub\n        self.create_optimizer_and_scheduler(\n            num_training_steps=args.num_total_batches\n        )  # note that we are calling `self.lr_scheduler.step()` manually only at the batch level\n\n        #########\n        ### trainer specifics\n        #########\n        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)\n        self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks\n        self.callback_handler = CallbackHandler(\n            self.callbacks, self.model, self.tokenizer, self.optimizer, self.lr_scheduler\n        )\n        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)\n        self.control = TrainerControl()\n        self.state = OnlineTrainerState(\n            is_local_process_zero=self.is_local_process_zero(),\n            is_world_process_zero=self.is_world_process_zero(),\n            stateful_callbacks=[\n                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)\n            ],\n        )\n        self.current_flos = 0\n        self.hp_search_backend = None\n        self.is_deepspeed_enabled = getattr(self.accelerator.state, \"deepspeed_plugin\", None) is not None\n        self.is_fsdp_enabled = getattr(self.accelerator.state, \"fsdp_plugin\", None) is not None\n        # Create distant repo and output directory if needed\n        self.hub_model_id = None\n        if self.args.push_to_hub:\n            self.init_hf_repo()\n        if self.args.should_save:\n            os.makedirs(self.args.output_dir, exist_ok=True)\n\n        #########\n        ### setup dataloader\n        #########\n        self.dataloader = DataLoader(\n            self.train_dataset,\n            batch_size=self.local_dataloader_batch_size,\n            shuffle=True,\n            collate_fn=DataCollatorWithPadding(tokenizer),\n            drop_last=True,  # needed; otherwise the last batch will be of ragged shape\n        )\n        # sync random states for DataLoader(shuffle=True) before `accelerator.prepare`\n        # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c\n        torch.manual_seed(args.seed)\n        self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader)\n        torch.manual_seed(self.local_seed)  # reset the local seed again\n\n        self.eval_dataloader = DataLoader(\n            self.eval_dataset,\n            batch_size=args.per_device_eval_batch_size,\n            collate_fn=DataCollatorWithPadding(self.tokenizer),\n            drop_last=True,\n        )  # no need to shuffle eval dataset\n        self.eval_dataloader = accelerator.prepare(self.eval_dataloader)\n\n        if self.is_deepspeed_enabled:\n            self.reward_model = prepare_deepspeed(\n                self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16\n            )\n            self.ref_policy = prepare_deepspeed(\n                self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16\n            )\n        else:\n            self.ref_policy = self.ref_policy.to(self.accelerator.device)\n            self.reward_model = self.reward_model.to(self.accelerator.device)\n\n    def get_train_dataloader(self) -> DataLoader:\n        return self.dataloader\n\n    def get_eval_dataloader(self) -> DataLoader:\n        return self.eval_dataloader\n\n    def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):\n        backup_model = self.model\n        self.model = self.model.policy  # save only the policy\n\n        if self.is_deepspeed_enabled:\n            backup_deepspeed = self.deepspeed\n            self.deepspeed = self.model\n\n        super().save_model(output_dir, _internal_call)\n\n        self.model = backup_model\n\n        if self.is_deepspeed_enabled:\n            self.deepspeed = backup_deepspeed\n\n    def train(self):\n        args = self.args\n        accelerator = self.accelerator\n        optimizer = self.optimizer\n        model = self.model\n        ref_policy = self.ref_policy\n        reward_model = self.reward_model\n        tokenizer = self.tokenizer\n        dataloader = self.dataloader\n        device = accelerator.device\n\n        def repeat_generator():\n            while True:\n                yield from dataloader\n\n        iter_dataloader = iter(repeat_generator())\n        generation_config = GenerationConfig(\n            max_new_tokens=args.response_length,\n            temperature=(args.temperature + 1e-7),\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n        )\n\n        accelerator.print(\"===training policy===\")\n        start_time = time.time()\n        stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps)\n        approxkl_stats = torch.zeros(stats_shape, device=device)\n        pg_clipfrac_stats = torch.zeros(stats_shape, device=device)\n        pg_loss_stats = torch.zeros(stats_shape, device=device)\n        vf_loss_stats = torch.zeros(stats_shape, device=device)\n        vf_clipfrac_stats = torch.zeros(stats_shape, device=device)\n        entropy_stats = torch.zeros(stats_shape, device=device)\n        ratio_stats = torch.zeros(stats_shape, device=device)\n        model.train()\n\n        # trainer state initialization\n        self.state.global_step = 0\n        self.state.episode = 0\n        self.state.max_steps = args.num_total_batches * args.num_mini_batches\n        self.state.num_train_epochs = args.total_episodes / self.train_dataset_len\n        # Compute absolute values for logging, eval, and save if given as ratio\n        if args.logging_steps is not None:\n            if args.logging_steps < 1:\n                self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps)\n            else:\n                self.state.logging_steps = args.logging_steps\n        if args.eval_steps is not None:\n            if args.eval_steps < 1:\n                self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps)\n            else:\n                self.state.eval_steps = args.eval_steps\n        if args.save_steps is not None:\n            if args.save_steps < 1:\n                self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps)\n            else:\n                self.state.save_steps = args.save_steps\n        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)\n\n        # backward compatibility\n        if self.is_deepspeed_enabled:\n            self.deepspeed = self.model\n            self.model_wrapped = self.model\n\n        for update in range(1, args.num_total_batches + 1):\n            self.state.episode += 1 * args.batch_size\n            data = next(iter_dataloader)\n            with torch.no_grad():\n                queries = data[\"input_ids\"].to(device)\n                context_length = queries.shape[1]\n                responses = []\n                postprocessed_responses = []\n                logprobs = []\n                ref_logprobs = []\n                scores = []\n                sequence_lengths = []\n                values = []\n                with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n                    query_responses, logitss = batch_generation(\n                        unwrapped_model.policy,\n                        queries,\n                        args.local_rollout_forward_batch_size,\n                        tokenizer.pad_token_id,\n                        generation_config,\n                    )\n\n                for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size):\n                    query = queries[i : i + args.local_rollout_forward_batch_size]\n                    query_response = query_responses[i : i + args.local_rollout_forward_batch_size]\n                    response = query_response[:, context_length:]\n                    logits = logitss[i : i + args.local_rollout_forward_batch_size]\n                    all_logprob = F.log_softmax(logits, dim=-1)\n                    logprob = torch.gather(all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)\n                    del logits, all_logprob\n                    torch.cuda.empty_cache()\n\n                    ref_output = forward(ref_policy, query_response, tokenizer.pad_token_id)\n                    ref_logits = ref_output.logits[:, context_length - 1 : -1]\n                    ref_logits /= args.temperature + 1e-7\n                    ref_all_logprob = F.log_softmax(ref_logits, dim=-1)\n                    ref_logprob = torch.gather(ref_all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)\n                    del ref_output, ref_logits, ref_all_logprob\n                    torch.cuda.empty_cache()\n\n                    # Response Processing 1. truncate response after the first occurrence of `stop_token_id`\n                    postprocessed_response = response\n                    if args.stop_token_id is not None:  # handle the edge case when stop_token_id exists but is 0\n                        postprocessed_response = truncate_response(\n                            args.stop_token_id, tokenizer.pad_token_id, response\n                        )\n\n                    # Response Processing 2. run reward model on the truncated responses\n                    postprocessed_query_response = torch.cat((query, postprocessed_response), 1)\n                    sequence_length = first_true_indices(postprocessed_response == tokenizer.pad_token_id) - 1\n                    unwrapped_value_model = accelerator.unwrap_model(model).value_model\n                    full_value, _, _ = get_reward(\n                        unwrapped_value_model, query_response, tokenizer.pad_token_id, context_length\n                    )\n                    value = full_value[:, context_length - 1 : -1].squeeze(-1)\n                    _, score, _ = get_reward(\n                        reward_model, postprocessed_query_response, tokenizer.pad_token_id, context_length\n                    )\n\n                    responses.append(response)\n                    postprocessed_responses.append(postprocessed_response)\n                    logprobs.append(logprob)\n                    ref_logprobs.append(ref_logprob)\n                    sequence_lengths.append(sequence_length)\n                    scores.append(score)\n                    values.append(value)\n                responses = torch.cat(responses, 0)\n                postprocessed_responses = torch.cat(postprocessed_responses, 0)\n                logprobs = torch.cat(logprobs, 0)\n                ref_logprobs = torch.cat(ref_logprobs, 0)\n                sequence_lengths = torch.cat(sequence_lengths, 0)\n                scores = torch.cat(scores, 0)\n                values = torch.cat(values, 0)\n                del (logprob, ref_logprob, full_value, value, score, unwrapped_model)\n                torch.cuda.empty_cache()\n                gc.collect()\n\n                # Response Processing 3. Filter completion. Ensure that the sample contains stop_token_id\n                # Completions not passing that filter will receive a lower score.\n                contain_eos_token = torch.any(postprocessed_responses == self.tokenizer.eos_token_id, dim=-1)\n                if self.args.missing_eos_penalty is not None:\n                    scores[~contain_eos_token] -= self.args.missing_eos_penalty\n                # accelerator.print(f\"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}\")\n\n                # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw\n                response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1)\n                padding_mask = response_idxs > sequence_lengths.unsqueeze(1)\n                logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB)\n                ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB)\n                sequence_lengths_p1 = sequence_lengths + 1\n                padding_mask_p1 = response_idxs > (sequence_lengths_p1.unsqueeze(1))\n                values = torch.masked_fill(values, padding_mask_p1, 0)\n\n                # 4. compute rewards\n                kl = logprobs - ref_logprobs\n                non_score_reward = -args.kl_coef * kl\n                rewards = non_score_reward.clone()\n                actual_start = torch.arange(rewards.size(0), device=rewards.device)\n                actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths)\n                rewards[[actual_start, actual_end]] += scores\n\n                # 5. whiten rewards\n                if args.whiten_rewards:\n                    rewards = masked_whiten(rewards, mask=~padding_mask_p1, shift_mean=False)\n                    rewards = torch.masked_fill(rewards, padding_mask_p1, 0)\n\n                # 6. compute advantages and returns\n                lastgaelam = 0\n                advantages_reversed = []\n                gen_length = responses.shape[1]\n                for t in reversed(range(gen_length)):\n                    nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0\n                    delta = rewards[:, t] + args.gamma * nextvalues - values[:, t]\n                    lastgaelam = delta + args.gamma * args.lam * lastgaelam\n                    advantages_reversed.append(lastgaelam)\n                advantages = torch.stack(advantages_reversed[::-1], axis=1)\n                returns = advantages + values\n                advantages = masked_whiten(advantages, ~padding_mask)\n                advantages = torch.masked_fill(advantages, padding_mask, 0)\n                torch.cuda.empty_cache()\n\n            # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch\n            for ppo_epoch_idx in range(args.num_ppo_epochs):\n                b_inds = np.random.permutation(args.local_batch_size)\n                minibatch_idx = 0\n                for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size):\n                    mini_batch_end = mini_batch_start + args.local_mini_batch_size\n                    mini_batch_inds = b_inds[mini_batch_start:mini_batch_end]\n                    gradient_accumulation_idx = 0\n                    for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size):\n                        with accelerator.accumulate(model):\n                            micro_batch_end = micro_batch_start + args.per_device_train_batch_size\n                            micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end]\n                            mb_advantage = advantages[micro_batch_inds]\n                            mb_responses = responses[micro_batch_inds]\n                            mb_query_responses = query_responses[micro_batch_inds]\n                            mb_logprobs = logprobs[micro_batch_inds]\n                            mb_return = returns[micro_batch_inds]\n                            mb_values = values[micro_batch_inds]\n\n                            output, vpred_temp = forward(model, mb_query_responses, tokenizer.pad_token_id)\n                            logits = output.logits[:, context_length - 1 : -1]\n                            logits /= args.temperature + 1e-7\n                            new_all_logprobs = F.log_softmax(logits, dim=-1)\n                            new_logprobs = torch.gather(new_all_logprobs, 2, mb_responses.unsqueeze(-1)).squeeze(-1)\n                            new_logprobs = torch.masked_fill(\n                                new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB\n                            )\n                            vpred = vpred_temp[:, context_length - 1 : -1].squeeze(-1)\n                            vpred = torch.masked_fill(vpred, padding_mask_p1[micro_batch_inds], 0)\n                            vpredclipped = torch.clamp(\n                                vpred,\n                                mb_values - args.cliprange_value,\n                                mb_values + args.cliprange_value,\n                            )\n                            vf_losses1 = torch.square(vpred - mb_return)\n                            vf_losses2 = torch.square(vpredclipped - mb_return)\n                            vf_loss_max = torch.max(vf_losses1, vf_losses2)\n                            vf_loss = 0.5 * masked_mean(vf_loss_max, ~padding_mask_p1[micro_batch_inds])\n                            vf_clipfrac = masked_mean(\n                                (vf_losses2 > vf_losses1).float(), ~padding_mask_p1[micro_batch_inds]\n                            )\n                            logprobs_diff = new_logprobs - mb_logprobs\n                            ratio = torch.exp(logprobs_diff)\n                            pg_losses = -mb_advantage * ratio\n                            pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange)\n                            pg_loss_max = torch.max(pg_losses, pg_losses2)\n                            pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds])\n                            loss = pg_loss + args.vf_coef * vf_loss\n                            accelerator.backward(loss)\n                            optimizer.step()\n                            optimizer.zero_grad()\n                            with torch.no_grad():\n                                pg_clipfrac = masked_mean(\n                                    (pg_losses2 > pg_losses).float(), ~padding_mask[micro_batch_inds]\n                                )\n                                prob_dist = torch.nn.functional.softmax(logits, dim=-1)\n                                entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1)\n                                approxkl = 0.5 * (logprobs_diff**2).mean()\n                                approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl\n                                pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = (\n                                    pg_clipfrac\n                                )\n                                pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss\n                                vf_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = vf_loss\n                                vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = (\n                                    vf_clipfrac\n                                )\n                                entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean()\n                                ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean()\n                        gradient_accumulation_idx += 1\n                    minibatch_idx += 1\n                    # del everything and empty cache\n                    # fmt: off\n                    del (\n                        output, vpred_temp, logits, new_all_logprobs, new_logprobs, vpred, vpredclipped,\n                        vf_losses1, vf_losses2, vf_loss, vf_clipfrac, logprobs_diff, ratio, pg_losses, pg_losses2, pg_loss_max,\n                        pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, mb_return,\n                        mb_advantage, mb_values, mb_responses, mb_query_responses, mb_logprobs,\n                    )\n                    # fmt: on\n                    torch.cuda.empty_cache()\n            with torch.no_grad():\n                mean_kl = kl.sum(1).mean()\n                mean_entropy = (-logprobs).sum(1).mean()\n                mean_non_score_reward = non_score_reward.sum(1).mean()\n                rlhf_reward = mean_non_score_reward + scores.mean()\n                eps = int(self.state.episode / (time.time() - start_time))\n                metrics = {}\n                metrics[\"eps\"] = eps\n                metrics[\"objective/kl\"] = self.accelerator.gather(mean_kl).mean().item()\n                metrics[\"objective/entropy\"] = self.accelerator.gather(mean_entropy).mean().item()\n                metrics[\"objective/non_score_reward\"] = self.accelerator.gather(mean_non_score_reward).mean().item()\n                metrics[\"objective/rlhf_reward\"] = self.accelerator.gather(rlhf_reward).mean().item()\n                metrics[\"objective/scores\"] = self.accelerator.gather(scores.mean()).mean().item()\n                metrics[\"policy/approxkl_avg\"] = self.accelerator.gather(approxkl_stats).mean().item()\n                metrics[\"policy/clipfrac_avg\"] = self.accelerator.gather(pg_clipfrac_stats).mean().item()\n                metrics[\"loss/policy_avg\"] = self.accelerator.gather(pg_loss_stats).mean().item()\n                metrics[\"loss/value_avg\"] = self.accelerator.gather(vf_loss_stats).mean().item()\n                metrics[\"val/clipfrac_avg\"] = self.accelerator.gather(vf_clipfrac_stats).mean().item()\n                metrics[\"policy/entropy_avg\"] = self.accelerator.gather(entropy_stats).mean().item()\n                metrics[\"val/ratio\"] = self.accelerator.gather(ratio_stats).mean().item()\n                metrics[\"val/ratio_var\"] = self.accelerator.gather(ratio_stats).var().item()\n                metrics[\"val/num_eos_tokens\"] = (responses == tokenizer.eos_token_id).sum().item()\n                metrics[\"lr\"] = self.lr_scheduler.get_last_lr()[0]\n                metrics[\"episode\"] = self.state.episode\n                self.state.epoch = self.state.episode / self.train_dataset_len  # used by self.log\n                self.state.global_step += 1\n                self.log(metrics)\n\n            self.lr_scheduler.step()\n            self.control = self.callback_handler.on_step_end(args, self.state, self.control)\n            if self.control.should_save:\n                self._save_checkpoint(model, trial=None, metrics=metrics)\n                self.control = self.callback_handler.on_save(self.args, self.state, self.control)\n            del kl, mean_kl, mean_entropy, mean_non_score_reward, scores, metrics, non_score_reward\n            torch.cuda.empty_cache()\n            gc.collect()\n\n            if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0:\n                self.generate_completions(sampling=True)\n                torch.cuda.empty_cache()\n            del (\n                query_responses,\n                responses,\n                postprocessed_responses,\n                logprobs,\n                ref_logprobs,\n                values,\n                sequence_lengths,\n                contain_eos_token,\n                sequence_lengths_p1,\n                response_idxs,\n                padding_mask,\n                padding_mask_p1,\n                rewards,\n                actual_start,\n                actual_end,\n                advantages,\n                returns,\n            )\n            torch.cuda.empty_cache()\n\n        # HF trainer specifics\n        self.control = self.callback_handler.on_train_end(args, self.state, self.control)\n        if self.control.should_save:\n            self._save_checkpoint(model, trial=None, metrics=None)\n            self.control = self.callback_handler.on_save(self.args, self.state, self.control)\n\n    def generate_completions(self, sampling: bool = False):\n        args = self.args\n        tokenizer = self.tokenizer\n        generation_config = GenerationConfig(\n            max_new_tokens=self.args.response_length,\n            temperature=(0.01 + 1e-7),\n            top_k=0.0,\n            top_p=1.0,\n            do_sample=True,\n        )\n\n        table = defaultdict(list)\n        with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:\n            for batch in self.eval_dataloader:\n                query = batch[\"input_ids\"]\n                with torch.no_grad():\n                    context_length = query.shape[1]\n                    query_response, _ = batch_generation(\n                        unwrapped_model.policy,\n                        query,\n                        query.shape[0],\n                        tokenizer.pad_token_id,\n                        generation_config,\n                    )\n                    response = query_response[:, context_length:]\n                    postprocessed_response = response\n                    if args.stop_token_id is not None:  # handle the edge case when stop_token_id exists but is 0\n                        postprocessed_response = truncate_response(\n                            args.stop_token_id, tokenizer.pad_token_id, response\n                        )\n                    table[\"query\"].extend(gather_object(tokenizer.batch_decode(query, skip_special_tokens=True)))\n                    table[\"model response\"].extend(gather_object(tokenizer.batch_decode(postprocessed_response)))\n\n                    postprocessed_query_response = torch.cat((query, postprocessed_response), 1)\n                    _, score, _ = get_reward(\n                        self.reward_model, postprocessed_query_response, tokenizer.pad_token_id, context_length\n                    )\n                    table[\"score\"].extend(self.accelerator.gather(score).float().cpu().numpy())\n\n                if sampling:\n                    break\n        df = pd.DataFrame(table)\n\n        if self.accelerator.is_main_process:\n            print_rich_table(df.iloc[0 : 0 + 5])\n            if \"wandb\" in args.report_to:\n                import wandb\n\n                if wandb.run is not None:\n                    wandb.log({\"completions\": wandb.Table(dataframe=df)})\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"ppo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# BCO Authors: Seungjae Jung, Gunsoo Han, Daniel Wontae Nam and Kyoung-Woon On\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport os\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager, nullcontext\nfrom copy import deepcopy\nfrom functools import wraps\nfrom operator import itemgetter\nfrom typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.amp as amp\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import PartialState\nfrom accelerate.utils import is_deepspeed_available, tqdm\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader, SequentialSampler\nfrom transformers import (\n    AutoModelForCausalLM,\n    DataCollator,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainingArguments,\n    is_sklearn_available,\n    is_wandb_available,\n)\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalLoopOutput, has_length\nfrom transformers.utils import is_peft_available\n\nfrom ..models import PreTrainedModelWrapper, create_reference_model\nfrom .bco_config import BCOConfig\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    RunningMoments,\n    disable_dropout_in_model,\n    pad_to_length,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\nif is_wandb_available():\n    import wandb\n\nif is_sklearn_available():\n    from sklearn.linear_model import LogisticRegression\n\nif is_deepspeed_available():\n    import deepspeed\n\nif TYPE_CHECKING:\n    from transformers import PreTrainedModel, PreTrainedTokenizer\n\nRUNNING_NAME = \"running.json\"\nCLF_NAME = \"clf.pt\"\n\n\ndef _tokenize(\n    batch: Dict[str, List[Any]],\n    tokenizer: \"PreTrainedTokenizer\",\n    embedding_tokenizer: Optional[\"PreTrainedTokenizer\"] = None,\n) -> Dict[str, List[Any]]:\n    \"\"\"Tokenize a batch from a BCO specific dataset.\"\"\"\n    prompt_tokenized = tokenizer(batch[\"prompt\"], add_special_tokens=False)\n    prompt_input_ids = prompt_tokenized[\"input_ids\"]\n    prompt_attention_mask = prompt_tokenized[\"attention_mask\"]\n    prompt_and_completion = [prompt + completion for prompt, completion in zip(batch[\"prompt\"], batch[\"completion\"])]\n    full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False)\n    full_input_ids = full_tokenized[\"input_ids\"]\n    full_attention_mask = full_tokenized[\"attention_mask\"]\n\n    answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids)]\n    answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask)]\n\n    # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`\n    full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids)]\n    # Prepare input tokens for token by token comparison\n    full_input_ids = [np.array(f) for f in full_input_ids]\n    for full, concat in zip(full_input_ids, full_concat_input_ids):\n        if len(full) != len(concat):\n            raise ValueError(\"Prompt input ids and answer input ids should have the same length.\")\n\n    # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens\n    # can be merged together when tokenizing prompt+answer. This could result\n    # on the last token from the prompt being different when tokenized on its own\n    # vs when done as prompt+answer.\n    response_token_ids_start_idx = [len(p) for p in prompt_input_ids]\n\n    # If tokenized prompt is different than both prompt+answer, then it means the\n    # last token has changed due to merging.\n    for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx)):\n        if not np.array_equal(p, f[:r]):\n            response_token_ids_start_idx[idx] -= 1\n\n    prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx)]\n    prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx)]\n\n    for p, m in zip(prompt_input_ids, prompt_attention_mask):\n        if len(p) != len(m):\n            raise ValueError(\"Prompt input ids and attention mask should have the same length.\")\n\n    answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx)]\n    answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx)]\n\n    output = dict(\n        prompt_input_ids=prompt_input_ids,\n        prompt_attention_mask=prompt_attention_mask,\n        answer_input_ids=answer_input_ids,\n        answer_attention_mask=answer_attention_mask,\n    )\n\n    if embedding_tokenizer is not None:\n        embedding_tokenized = embedding_tokenizer(batch[\"prompt\"], truncation=True, add_special_tokens=False)\n\n        output.update(\n            {\n                \"embedding_input_ids\": embedding_tokenized[\"input_ids\"],\n                \"embedding_attention_mask\": embedding_tokenized[\"attention_mask\"],\n            }\n        )\n\n    return output\n\n\ndef _process_tokens(example: Dict[str, Any], model: \"PreTrainedModel\" = None, **kwargs) -> Dict:\n    \"\"\"Process tokens of a BCO specific dataset.\n\n    At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation\n    in case the prompt + completion responses is/are too long. First\n    we truncate the prompt; if we're still too long, we truncate the completion.\n\n    We also create the labels for the completion responses, which are of length equal to\n    the sum of the length of the prompt and the completion response, with\n    label_pad_token_id  for the prompt tokens.\n    \"\"\"\n    prompt = example[\"prompt\"]\n    completion = example[\"completion\"]\n\n    batch = {\n        f\"{kwargs['prefix']}prompt\": prompt,\n        f\"{kwargs['prefix']}completion\": completion,\n        f\"{kwargs['prefix']}label\": example[\"label\"],\n    }\n\n    if not kwargs[\"is_encoder_decoder\"]:\n        # Check issues below for more details\n        #  1. https://github.com/huggingface/trl/issues/907\n        #  2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n        #  3. https://github.com/LianjiaTech/BELLE/issues/337\n\n        if not isinstance(prompt, str):\n            raise ValueError(f\"prompt should be an str but got {type(prompt)}\")\n\n        if not isinstance(completion, str):\n            raise ValueError(f\"completion should be an str but got {type(completion)}\")\n\n        # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer\n        all_tokens = {\n            \"prompt_input_ids\": example[\"prompt_input_ids\"],\n            \"prompt_attention_mask\": example[\"prompt_attention_mask\"],\n            \"answer_input_ids\": example[\"answer_input_ids\"],\n            \"answer_attention_mask\": example[\"answer_attention_mask\"],\n        }\n\n        # calculate max length by checking if BOS/EOS is already there\n        max_length = kwargs[\"max_length\"]\n        bos_token_id = kwargs[\"tokenizer\"].bos_token_id\n        eos_token_id = kwargs[\"tokenizer\"].eos_token_id\n        if bos_token_id != all_tokens[\"prompt_input_ids\"][0]:\n            max_length -= 1\n        if eos_token_id != all_tokens[\"answer_input_ids\"][-1]:\n            max_length -= 1\n\n        # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the prompt\n        if len(all_tokens[\"prompt_input_ids\"]) + len(all_tokens[\"answer_input_ids\"]) > max_length:\n            for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                if kwargs[\"truncation_mode\"] == \"keep_start\":\n                    all_tokens[k] = all_tokens[k][: kwargs[\"max_prompt_length\"]]\n                elif kwargs[\"truncation_mode\"] == \"keep_end\":\n                    all_tokens[k] = all_tokens[k][-kwargs[\"max_prompt_length\"] :]\n                else:\n                    raise ValueError(f\"Unknown truncation mode: {kwargs['truncation_mode']}\")\n\n        # if that's still too long, truncate the response\n        if len(all_tokens[\"prompt_input_ids\"]) + len(all_tokens[\"answer_input_ids\"]) > max_length:\n            for k in [\"answer_input_ids\", \"answer_attention_mask\"]:\n                all_tokens[k] = all_tokens[k][: max_length - kwargs[\"max_prompt_length\"]]\n\n        # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens\n        batch[f\"{kwargs['prefix']}prompt_input_ids\"] = all_tokens[\"prompt_input_ids\"]\n        batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = all_tokens[\"prompt_attention_mask\"]\n        batch[f\"{kwargs['prefix']}completion_input_ids\"] = (\n            all_tokens[\"prompt_input_ids\"] + all_tokens[\"answer_input_ids\"]\n        )\n        batch[f\"{kwargs['prefix']}completion_attention_mask\"] = (\n            all_tokens[\"prompt_attention_mask\"] + all_tokens[\"answer_attention_mask\"]\n        )\n\n        # add BOS, which affects both prompt and the full completion\n        if len(all_tokens[\"prompt_input_ids\"]) == 0 or bos_token_id != all_tokens[\"prompt_input_ids\"][0]:\n            batch[f\"{kwargs['prefix']}prompt_input_ids\"] = [bos_token_id] + batch[\n                f\"{kwargs['prefix']}prompt_input_ids\"\n            ]\n            batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = [1] + batch[f\"{kwargs['prefix']}prompt_attention_mask\"]\n            batch[f\"{kwargs['prefix']}completion_input_ids\"] = [bos_token_id] + batch[\n                f\"{kwargs['prefix']}completion_input_ids\"\n            ]\n            batch[f\"{kwargs['prefix']}completion_attention_mask\"] = [1] + batch[\n                f\"{kwargs['prefix']}completion_attention_mask\"\n            ]\n        # add EOS, which affects only the full completion\n        if len(all_tokens[\"answer_input_ids\"]) == 0 or eos_token_id != all_tokens[\"answer_input_ids\"][-1]:\n            batch[f\"{kwargs['prefix']}completion_input_ids\"] = batch[f\"{kwargs['prefix']}completion_input_ids\"] + [\n                eos_token_id\n            ]\n            batch[f\"{kwargs['prefix']}completion_attention_mask\"] = batch[\n                f\"{kwargs['prefix']}completion_attention_mask\"\n            ] + [1]\n\n        batch[f\"{kwargs['prefix']}completion_labels\"] = batch[f\"{kwargs['prefix']}completion_input_ids\"][:]\n        batch[f\"{kwargs['prefix']}completion_labels\"][: len(batch[f\"{kwargs['prefix']}prompt_input_ids\"])] = [\n            kwargs[\"label_pad_token_id\"]\n        ] * len(batch[f\"{kwargs['prefix']}prompt_input_ids\"])\n    else:\n        completion_tokens = kwargs[\"tokenizer\"](\n            completion, truncation=True, max_length=kwargs[\"max_completion_length\"], add_special_tokens=True\n        )\n        prompt_tokens = kwargs[\"tokenizer\"](\n            prompt, truncation=True, max_length=kwargs[\"max_prompt_length\"], add_special_tokens=True\n        )\n\n        batch[f\"{kwargs['prefix']}prompt_input_ids\"] = prompt_tokens[\"input_ids\"]\n        batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = prompt_tokens[\"attention_mask\"]\n\n        batch[f\"{kwargs['prefix']}completion_labels\"] = completion_tokens[\"input_ids\"]\n        batch[f\"{kwargs['prefix']}completion_attention_mask\"] = completion_tokens[\"attention_mask\"]\n        if model is not None and hasattr(model, \"prepare_decoder_input_ids_from_labels\"):\n            batch[f\"{kwargs['prefix']}completion_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                labels=torch.tensor(batch[\"completion_labels\"])\n            )\n\n    return batch\n\n\nclass BCOTrainer(Trainer):\n    r\"\"\"\n    Initialize BCOTrainer from [BCO](https://arxiv.org/abs/2404.04656) paper.\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForSequenceClassification`.\n        ref_model (`PreTrainedModelWrapper`):\n            Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no\n            reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.\n        args (`BCOConfig`):\n            The arguments to use for training.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        data_collator (`transformers.DataCollator`, *optional*, defaults to `None`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Dict`, defaults to `None`):\n            The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        disable_dropout (`bool`, defaults to `True`):\n            Whether or not to disable dropouts in `model` and `ref_model`.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n        model_adapter_name (`str`, defaults to `None`):\n            Name of the train target PEFT adapter, when using LoRA with multiple adapters.\n        ref_adapter_name (`str`, defaults to `None`):\n            Name of the reference PEFT adapter, when using LoRA with multiple adapters.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"bco\"]\n\n    def __init__(\n        self,\n        model: Union[PreTrainedModel, nn.Module, str] = None,\n        ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        args: BCOConfig = None,\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        data_collator: Optional[DataCollator] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n        model_adapter_name: Optional[str] = None,\n        ref_adapter_name: Optional[str] = None,\n        embedding_func: Optional[Callable] = None,\n        embedding_tokenizer: Optional[PreTrainedTokenizerBase] = None,\n    ):\n        if not is_sklearn_available():\n            raise ImportError(\n                \"BCOTrainer requires the scikit-learn library. Please install it with `pip install scikit-learn`.\"\n            )\n\n        if type(args) is TrainingArguments:\n            raise ValueError(\"Please use `BCOConfig` instead `TrainingArguments`.\")\n\n        if not isinstance(model, str) and ref_model is model:\n            raise ValueError(\n                \"`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the \"\n                \"same as `model`, you must mass a copy of it, or `None` if you use peft.\"\n            )\n\n        if args.model_init_kwargs is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\"You passed model_kwargs to the BCOTrainer. But your model is already instantiated.\")\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the BCOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if args.ref_model_init_kwargs is None:\n            ref_model_init_kwargs = {}\n        elif not isinstance(ref_model, str):\n            raise ValueError(\n                \"You passed ref_model_kwargs to the BCOTrainer. But your ref_model is already instantiated.\"\n            )\n        else:\n            ref_model_init_kwargs = args.ref_model_init_kwargs\n            torch_dtype = ref_model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the BCOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                ref_model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the BCOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        if isinstance(ref_model, str):\n            warnings.warn(\n                \"You passed a ref model_id to the BCOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM`\"\n            )\n            ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs)\n\n        # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`\n        # has been called in order to properly call autocast if needed.\n        self._peft_has_been_casted_to_bf16 = False\n\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it with `pip install peft` to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            # if model is a peft model and we have a peft_config, we merge and unload it first\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_loaded_in_4bit\", False):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n\n                prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                if _support_gc_kwargs:\n                    prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n            elif getattr(args, \"gradient_checkpointing\", False):\n                # For backward compatibility with older versions of transformers\n                if hasattr(model, \"enable_input_require_grads\"):\n                    model.enable_input_require_grads()\n                else:\n\n                    def make_inputs_require_grad(module, input, output):\n                        output.requires_grad_(True)\n\n                    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            # get peft model with the given config\n            model = get_peft_model(model, peft_config)\n            if args.bf16 and getattr(model, \"is_loaded_in_4bit\", False):\n                peft_module_casting_to_bf16(model)\n                # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager\n                self._peft_has_been_casted_to_bf16 = True\n\n        # For models that use gradient_checkpointing, we need to attach a hook that enables input\n        # to explicitly have `requires_grad=True`, otherwise training will either silently\n        # fail or completely fail.\n        elif getattr(args, \"gradient_checkpointing\", False):\n            # For backward compatibility with older versions of transformers\n            if hasattr(model, \"enable_input_require_grads\"):\n                model.enable_input_require_grads()\n            else:\n\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        if args.generate_during_eval and not is_wandb_available():\n            raise ValueError(\n                \"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install with `pip install wandb` to resolve.\"\n            )\n\n        if model is not None:\n            self.is_encoder_decoder = model.config.is_encoder_decoder\n        elif args.is_encoder_decoder is None:\n            raise ValueError(\"When no model is provided, you need to pass the parameter is_encoder_decoder.\")\n        else:\n            self.is_encoder_decoder = args.is_encoder_decoder\n\n        self.is_peft_model = is_peft_available() and isinstance(model, PeftModel)\n        self.model_adapter_name = model_adapter_name\n        self.ref_adapter_name = ref_adapter_name\n\n        if ref_model:\n            self.ref_model = ref_model\n        elif self.is_peft_model or args.precompute_ref_log_probs:\n            # The `model` with adapters turned off will be used as the reference model\n            self.ref_model = None\n        else:\n            self.ref_model = create_reference_model(model)\n\n        if tokenizer is None:\n            raise ValueError(\n                \"max_length or a tokenizer must be specified when using the default DPODataCollatorWithPadding\"\n            )\n        if args.max_length is None:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding, you should set `max_length` in the `BCOConfig`. \"\n                \"It will be set to `512` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_length = 512\n        if args.max_length is not None:\n            max_length = args.max_length\n\n        if args.max_prompt_length is None:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the `BCOConfig`. \"\n                \"It will be set to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_prompt_length = 128\n        if args.max_prompt_length is not None:\n            max_prompt_length = args.max_prompt_length\n\n        max_completion_length = None\n        if args.max_completion_length is None and self.is_encoder_decoder:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the BCOTrainer's init\"\n                \" it will be set to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_completion_length = 128\n        if args.max_completion_length is not None and self.is_encoder_decoder:\n            max_completion_length = args.max_completion_length\n\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(\n                pad_token_id=tokenizer.pad_token_id,\n                label_pad_token_id=args.label_pad_token_id,\n                is_encoder_decoder=self.is_encoder_decoder,\n            )\n\n            if args.remove_unused_columns:\n                args.remove_unused_columns = False\n                # warn users\n                warnings.warn(\n                    \"When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your BCOConfig\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_dpo_data_collator = True\n        else:\n            self.use_dpo_data_collator = False\n\n        # disable dropout in the model and reference model\n        disable_dropout_in_model(model)\n        if self.ref_model is not None:\n            disable_dropout_in_model(self.ref_model)\n\n        self.max_length = max_length\n        self.generate_during_eval = args.generate_during_eval\n        self.label_pad_token_id = args.label_pad_token_id\n        self.padding_value = args.padding_value if args.padding_value is not None else tokenizer.pad_token_id\n        self.max_prompt_length = max_prompt_length\n        self.truncation_mode = args.truncation_mode\n        self.max_completion_length = max_completion_length\n        self.tokenizer = tokenizer\n        self.precompute_ref_log_probs = args.precompute_ref_log_probs\n\n        # Since ref_logs are precomputed on the first call to get_train/eval_dataloader\n        # keep track of first called to avoid computation of future calls\n        self._precomputed_train_ref_log_probs = False\n        self._precomputed_eval_ref_log_probs = False\n\n        # metric\n        self._stored_metrics = defaultdict(lambda: defaultdict(list))\n\n        # BCO parameter\n        self.beta = args.beta\n        self.aux_loss_enabled = getattr(model.config, \"output_router_logits\", False)\n\n        # Underlying Distribution Matching argument\n        self.embedding_func = embedding_func\n        self.embedding_tokenizer = embedding_tokenizer\n\n        with PartialState().local_main_process_first():\n            # Shuffle the datasets\n            train_dataset = train_dataset.shuffle(seed=args.data_seed)\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.shuffle(seed=args.data_seed)\n            # Tokenize and prepare the training datasets\n            train_dataset = train_dataset.map(\n                _tokenize,\n                batched=True,\n                fn_kwargs={\"tokenizer\": self.tokenizer, \"embedding_tokenizer\": self.embedding_tokenizer},\n                num_proc=args.dataset_num_proc,\n                desc=\"Tokenizing train dataset\",\n            )\n\n            # Prepare the datasets\n            fn_kwargs = {\n                \"prefix\": \"\",\n                \"is_encoder_decoder\": self.is_encoder_decoder,\n                \"tokenizer\": self.tokenizer,\n                \"max_length\": self.max_length,\n                \"truncation_mode\": self.truncation_mode,\n                \"label_pad_token_id\": self.label_pad_token_id,\n                \"max_prompt_length\": self.max_prompt_length,\n                \"max_completion_length\": self.max_completion_length,\n            }\n            train_dataset = train_dataset.map(\n                _process_tokens,\n                fn_kwargs=fn_kwargs,\n                num_proc=args.dataset_num_proc,\n                desc=\"Processing tokenized train dataset\",\n            )\n\n            if eval_dataset is not None:\n                # Tokenize\n                eval_dataset = eval_dataset.map(\n                    _tokenize,\n                    fn_kwargs={\"tokenizer\": self.tokenizer, \"embedding_tokenizer\": self.embedding_tokenizer},\n                    batched=True,\n                    num_proc=args.dataset_num_proc,\n                    desc=\"Tokenizing eval dataset\",\n                )\n\n                # Process\n                fn_kwargs = {\n                    \"prefix\": \"\",\n                    \"is_encoder_decoder\": self.is_encoder_decoder,\n                    \"tokenizer\": self.tokenizer,\n                    \"max_length\": self.max_length,\n                    \"truncation_mode\": self.truncation_mode,\n                    \"label_pad_token_id\": self.label_pad_token_id,\n                    \"max_prompt_length\": self.max_prompt_length,\n                    \"max_completion_length\": self.max_completion_length,\n                }\n                eval_dataset = eval_dataset.map(\n                    _process_tokens,\n                    fn_kwargs=fn_kwargs,\n                    num_proc=args.dataset_num_proc,\n                    desc=\"Processing tokenized eval dataset\",\n                )\n\n            desirable = train_dataset.filter(\n                lambda x: x[\"label\"], num_proc=args.dataset_num_proc, desc=\"Filtering desirable examples\"\n            )\n            undesirable = train_dataset.filter(\n                lambda x: not x[\"label\"], num_proc=args.dataset_num_proc, desc=\"Filtering undesirable examples\"\n            )\n\n            desirable = desirable.shuffle(seed=args.data_seed)\n            undesirable = undesirable.shuffle(seed=args.data_seed)\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n        # Deepspeed Zero-3 does not support precompute_ref_log_probs\n        if self.is_deepspeed_enabled:\n            if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs:\n                raise ValueError(\n                    \"You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`.\"\n                )\n\n        if self.ref_model is None:\n            if not (self.is_peft_model or self.precompute_ref_log_probs):\n                raise ValueError(\n                    \"No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`\"\n                )\n        else:\n            if self.is_deepspeed_enabled:\n                self.ref_model = self._prepare_deepspeed(self.ref_model)\n            else:\n                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)\n\n        self.running = RunningMoments(accelerator=self.accelerator)\n\n        if self.embedding_func is None:\n            warnings.warn(\"You did not pass `embedding_func` underlying distribution matching feature is deactivated.\")\n            return\n\n        chosen_embeddings = self._get_sample_prompt_embeddings(desirable, sample_size=self.args.prompt_sample_size)\n        rejected_embeddings = self._get_sample_prompt_embeddings(undesirable, sample_size=self.args.prompt_sample_size)\n\n        embeddings = torch.cat((chosen_embeddings, rejected_embeddings), dim=0)\n        labels = torch.cat(\n            (torch.ones_like(chosen_embeddings[:, 0]), torch.zeros_like(rejected_embeddings[:, 0])), dim=0\n        )\n\n        self.clf = LogisticRegression(class_weight=\"balanced\").fit(\n            embeddings.cpu().float().numpy(), labels.cpu().numpy()\n        )\n\n    @property\n    def match_underlying_distribution(self):\n        return self.embedding_func is not None and self.embedding_tokenizer is not None\n\n    def _get_chosen_prob(self, prompt_embeddings: torch.FloatTensor) -> torch.FloatTensor:\n        \"\"\"\n        Calculates the probability if the given prompt embedding is from desirable dataset.\n        This function calculates the probability in the process and ensemble across processes.\n        \"\"\"\n        dtype = prompt_embeddings.dtype\n        device = prompt_embeddings.device\n        rank = self.accelerator.process_index\n\n        padded_prompt_embeddings = self.accelerator.pad_across_processes(\n            prompt_embeddings, pad_index=self.embedding_tokenizer.pad_token_id\n        )\n        sample_size = padded_prompt_embeddings.shape[0]\n        nonzero = padded_prompt_embeddings.mean(dim=1) != self.embedding_tokenizer.pad_token_id\n        prompt_embeddings = self.accelerator.gather(padded_prompt_embeddings)\n\n        # cannot predict for all empty values\n        if prompt_embeddings.shape[0] == 0:\n            return torch.tensor([], device=device, dtype=dtype)\n\n        prob = self.clf.predict_proba(prompt_embeddings.cpu().float().numpy())[:, 1]\n        prob = torch.as_tensor(prob, dtype=dtype, device=device)\n        prob = self.accelerator.reduce(prob, reduction=\"mean\")\n\n        prob = prob[sample_size * rank : sample_size * (rank + 1)]\n        prob = prob[nonzero]\n\n        return prob\n\n    def _vectorize_prompt(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor) -> torch.FloatTensor:\n        \"\"\"\n        Replaces tokenizer.pad_token_id to embedding_tokenizer.pad_token_id\n        and applies self.embedding_func\n        \"\"\"\n        input_ids = torch.where(\n            input_ids == self.tokenizer.pad_token_id,\n            self.embedding_tokenizer.pad_token_id,\n            input_ids,\n        )\n\n        with torch.no_grad():\n            embeddings = self.embedding_func(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n            )\n\n        return embeddings\n\n    def _get_prompt_embeddings(\n        self, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Extract embeddings from frozen embedding model\"\"\"\n\n        if not self.match_underlying_distribution:\n            return None, None\n\n        embeddings = self._vectorize_prompt(\n            input_ids=batch[\"embedding_input_ids\"],\n            attention_mask=batch[\"embedding_attention_mask\"],\n        )\n\n        chosen_idx = [i for i in range(len(batch[\"label\"])) if batch[\"label\"][i] is True]\n        rejected_idx = [i for i in range(len(batch[\"label\"])) if batch[\"label\"][i] is False]\n\n        chosen_embeddings = embeddings[chosen_idx, ...]\n        rejected_embeddings = embeddings[rejected_idx, ...]\n\n        return (chosen_embeddings, rejected_embeddings)\n\n    def _get_sample_prompt_embeddings(self, dataset: Dataset, sample_size: int = 512) -> torch.FloatTensor:\n        \"\"\"\n        Sample instances from dataset and get prompt embeddings.\n        Used for density ratio classifier training.\n        \"\"\"\n        n_samples = min(len(dataset), sample_size)\n        rand_indices = np.random.choice(len(dataset), size=(n_samples,))\n\n        embedding_dataset = dataset.select(rand_indices)\n\n        dataloader_params = {\n            \"batch_size\": self.args.per_device_train_batch_size,\n            \"collate_fn\": self.data_collator,\n            \"num_workers\": self.args.dataloader_num_workers,\n            \"pin_memory\": self.args.dataloader_pin_memory,\n            \"shuffle\": False,\n        }\n\n        # prepare dataloader\n        data_loader = self.accelerator.prepare(DataLoader(embedding_dataset, **dataloader_params))\n\n        with torch.no_grad():\n            all_embeddings = torch.empty(0)\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Building sample prompt embeddings\"):\n                embeddings = self._vectorize_prompt(\n                    input_ids=padded_batch[\"embedding_input_ids\"],\n                    attention_mask=padded_batch[\"embedding_attention_mask\"],\n                )\n                embeddings = self.accelerator.gather_for_metrics(embeddings)\n                all_embeddings = torch.cat((all_embeddings, embeddings.cpu()))\n\n        return all_embeddings\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)\n\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n    def _save_optimizer_and_scheduler(self, output_dir):\n        super()._save_optimizer_and_scheduler(output_dir)\n\n        # When saving optimizer and scheduler to checkpoint, save also the running delta object.\n        output_dir = output_dir if output_dir is not None else self.args.output_dir\n\n        self.running.save_to_json(os.path.join(output_dir, RUNNING_NAME))\n\n        if self.match_underlying_distribution:\n            torch.save(self.clf.get_params(), os.path.join(output_dir, CLF_NAME))\n\n    def _load_optimizer_and_scheduler(self, checkpoint):\n        super()._load_optimizer_and_scheduler(checkpoint)\n\n        if checkpoint is None:\n            return\n        # when loading optimizer and scheduler from checkpoint, also load the running delta object.\n        running_file = os.path.join(checkpoint, RUNNING_NAME)\n        if not os.path.isfile(running_file):\n            warnings.warn(f\"Missing file {running_file}. Will use a new running delta value for BCO loss calculation\")\n        else:\n            self.running = RunningMoments.load_from_json(self.accelerator, running_file)\n\n        if self.match_underlying_distribution:\n            clf_file = os.path.join(checkpoint, CLF_NAME)\n            if not os.path.isfile(running_file):\n                warnings.warn(f\"Missing file {clf_file}. Will use a new UDM classifier for BCO loss calculation\")\n            else:\n                self.clf.set_params(**torch.load(clf_file, weights_only=True, map_location=\"cpu\"))\n\n    @contextmanager\n    def null_ref_context(self):\n        \"\"\"Context manager for handling null reference model (that is, peft adapter manipulation).\"\"\"\n        with self.accelerator.unwrap_model(\n            self.model\n        ).disable_adapter() if self.is_peft_model and not self.ref_adapter_name else nullcontext():\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.ref_adapter_name)\n            yield\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.model_adapter_name or \"default\")\n\n    def get_train_dataloader(self) -> DataLoader:\n        \"\"\"\n        Returns the training [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`.\n        \"\"\"\n\n        if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_train_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params))\n            reference_completion_logps = []\n\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Train dataset reference log probs\"):\n                reference_completion_logp = self.compute_reference_log_probs(padded_batch)\n\n                reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp)\n                reference_completion_logps.append(reference_completion_logp.cpu())\n\n            self.train_dataset = self.train_dataset.add_column(\n                name=\"reference_logps\", column=torch.cat(reference_completion_logps).float().numpy()\n            )\n\n            self._precomputed_train_ref_log_probs = True\n\n        return super().get_train_dataloader()\n\n    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:\n        \"\"\"\n        Returns the evaluation [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`.\n\n        Args:\n            eval_dataset (`torch.utils.data.Dataset`, *optional*):\n                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted\n                by the `model.forward()` method are automatically removed. It must implement `__len__`.\n        \"\"\"\n        if eval_dataset is None and self.eval_dataset is None:\n            raise ValueError(\"Trainer: evaluation requires an eval_dataset.\")\n        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset\n\n        if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_eval_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))\n\n            reference_completion_logps = []\n\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Eval dataset reference log probs\"):\n                reference_completion_logp = self.compute_reference_log_probs(padded_batch)\n\n                reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp)\n                reference_completion_logps.append(reference_completion_logp.cpu())\n\n            eval_dataset = eval_dataset.add_column(\n                name=\"reference_logps\", column=torch.cat(reference_completion_logps).float().numpy()\n            )\n\n            # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs\n            if self.eval_dataset is not None:\n                self.eval_dataset = eval_dataset\n            self._precomputed_eval_ref_log_probs = True\n\n        return super().get_eval_dataloader(eval_dataset=eval_dataset)\n\n    def compute_reference_log_probs(self, padded_batch: Dict) -> Dict:\n        \"\"\"Computes log probabilities of the reference model for a single padded batch of a BCO specific dataset.\"\"\"\n        with torch.no_grad():\n            if self.ref_model is None:\n                with self.null_ref_context():\n                    if self.is_encoder_decoder:\n                        completion_logits = self.model(\n                            padded_batch[\"prompt_input_ids\"],\n                            attention_mask=padded_batch[\"prompt_attention_mask\"],\n                            decoder_input_ids=padded_batch.get(\"completion_decoder_input_ids\"),\n                            labels=padded_batch[\"completion_labels\"],\n                        ).logits\n\n                    else:\n                        completion_logits = self.model(\n                            padded_batch[\"completion_input_ids\"],\n                            attention_mask=padded_batch[\"completion_attention_mask\"],\n                        ).logits\n\n            else:\n                if self.is_encoder_decoder:\n                    completion_logits = self.ref_model(\n                        padded_batch[\"prompt_input_ids\"],\n                        attention_mask=padded_batch[\"prompt_attention_mask\"],\n                        decoder_input_ids=padded_batch.get(\"completion_decoder_input_ids\"),\n                        labels=padded_batch[\"completion_labels\"],\n                    ).logits\n\n                else:\n                    completion_logits = self.ref_model(\n                        padded_batch[\"completion_input_ids\"], attention_mask=padded_batch[\"completion_attention_mask\"]\n                    ).logits\n\n        completion_logps = self.get_batch_logps(\n            completion_logits,\n            padded_batch[\"completion_labels\"],\n            average_log_prob=False,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        return completion_logps\n\n    @staticmethod\n    def get_batch_logps(\n        logits: torch.FloatTensor,\n        labels: torch.LongTensor,\n        average_log_prob: bool = False,\n        label_pad_token_id: int = -100,\n        is_encoder_decoder: bool = False,\n    ) -> torch.FloatTensor:\n        \"\"\"Compute the log probabilities of the given labels under the given logits.\n\n        Args:\n            logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n            labels: Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are ignored. Shape: (batch_size, sequence_length)\n            average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the (non-masked) tokens.\n\n        Returns:\n            A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the given logits.\n        \"\"\"\n        if logits.shape[:-1] != labels.shape:\n            raise ValueError(\"Logits (batch and sequence length dim) and labels must have the same shape.\")\n\n        if not is_encoder_decoder:\n            labels = labels[:, 1:].clone()\n            logits = logits[:, :-1, :]\n        else:\n            # Fixes end-dec RuntimeError\n            labels = labels.clone()\n\n        loss_mask = labels != label_pad_token_id\n\n        # dummy token; we'll ignore the losses on these tokens later\n        labels[labels == label_pad_token_id] = 0\n\n        per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n\n        if average_log_prob:\n            return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n        else:\n            return (per_token_logps * loss_mask).sum(-1)\n\n    def forward(\n        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        model_kwargs = (\n            {\n                \"labels\": batch[\"completion_labels\"],\n                \"decoder_input_ids\": batch.get(\"completion_decoder_input_ids\"),\n            }\n            if self.is_encoder_decoder\n            else {}\n        )\n        if self.aux_loss_enabled:\n            model_kwargs[\"output_router_logits\"] = True\n\n        outputs = model(\n            batch[\"completion_input_ids\"],\n            attention_mask=batch[\"completion_attention_mask\"],\n            **model_kwargs,\n        )\n        completion_logits = outputs.logits\n\n        completion_logps = self.get_batch_logps(\n            completion_logits,\n            batch[\"completion_labels\"],\n            average_log_prob=False,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        if completion_logps.shape[0] != len(batch[\"label\"]):\n            raise ValueError(\n                \"There is a mismatch between the number of examples in this batch and the number of \"\n                \"examples for which an output sequence was predicted.\"\n            )\n\n        chosen_idx = [i for i in range(completion_logps.shape[0]) if batch[\"label\"][i] is True]\n        rejected_idx = [i for i in range(completion_logps.shape[0]) if batch[\"label\"][i] is False]\n\n        chosen_logps = completion_logps[chosen_idx, ...]\n        rejected_logps = completion_logps[rejected_idx, ...]\n\n        chosen_logits = completion_logits[chosen_idx, ...]\n        rejected_logits = completion_logits[rejected_idx, ...]\n\n        if self.aux_loss_enabled:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, outputs.aux_loss)\n        else:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits)\n\n    def _get_udm_weight(self, rejected_embeddings: torch.FloatTensor) -> torch.FloatTensor:\n        prob_desirable = self._get_chosen_prob(rejected_embeddings)\n        min_ratio = self.args.min_density_ratio\n        max_ratio = self.args.max_density_ratio\n\n        weight = (prob_desirable / (1 - prob_desirable + 1e-8)).clamp(min=min_ratio, max=max_ratio)\n\n        return weight\n\n    def bco_loss(\n        self,\n        policy_chosen_logps: torch.FloatTensor,\n        policy_rejected_logps: torch.FloatTensor,\n        reference_chosen_logps: torch.FloatTensor,\n        reference_rejected_logps: torch.FloatTensor,\n        chosen_embeddings: Optional[torch.FloatTensor],\n        rejected_embeddings: Optional[torch.FloatTensor],\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Compute the BCO loss for a batch of policy and reference model log probabilities.\n\n        Args:\n            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,)\n            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,)\n            reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,)\n            reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in batch_size,)\n            chosen_embeddings: embeddings of desirable prompts\n            rejected_embeddings: embeddings of undesirable prompts\n\n        Returns:\n            A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, delta).\n            The losses tensor contains the BCO loss for each example in the batch.\n            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.\n            The delta value contains the moving average of all implicit rewards.\n        \"\"\"\n\n        if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0:\n            chosen_logratios = policy_chosen_logps - reference_chosen_logps\n            chosen_rewards = self.beta * chosen_logratios\n        else:\n            # lists can't be empty -- if they are, then accelerate.gather will hang\n            chosen_losses = torch.Tensor([]).to(self.accelerator.device)\n            chosen_rewards = torch.Tensor([]).to(self.accelerator.device)\n\n        if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0:\n            rejected_logratios = policy_rejected_logps - reference_rejected_logps\n            rejected_rewards = self.beta * rejected_logratios\n        else:\n            # lists can't be empty -- if they are, then accelerate.gather will hang\n            rejected_losses = torch.Tensor([]).to(self.accelerator.device)\n            rejected_rewards = torch.Tensor([]).to(self.accelerator.device)\n\n        rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach()\n        self.running.update(rewards)\n        delta = self.running.mean\n\n        if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0:\n            chosen_losses = -F.logsigmoid(chosen_rewards - delta)\n\n        if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0:\n            rejected_losses = -F.logsigmoid(-(rejected_rewards - delta))\n\n        if self.match_underlying_distribution:\n            chosen_weight = torch.ones_like(chosen_losses)\n            rejected_weight = self._get_udm_weight(rejected_embeddings)\n\n            losses = torch.cat((chosen_weight * chosen_losses, rejected_weight * rejected_losses), dim=0)\n        else:\n            losses = torch.cat((chosen_losses, rejected_losses), dim=0)\n\n        return losses, chosen_rewards, rejected_rewards, torch.as_tensor(delta)\n\n    def get_batch_loss_metrics(\n        self,\n        model,\n        batch: Dict[str, Union[List, torch.LongTensor]],\n    ):\n        \"\"\"Compute the BCO loss and other metrics for the given batch of inputs for train or test.\"\"\"\n        metrics = {}\n        batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}\n\n        forward_output = self.forward(model, batch)\n        (\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_chosen_logits,\n            policy_rejected_logits,\n        ) = forward_output[:4]\n        if self.aux_loss_enabled:\n            aux_loss = forward_output[4]\n\n        # if reference_logps in batch use them, otherwise use the reference model\n        if \"reference_logps\" in batch:\n            chosen_idx = [i for i in range(batch[\"reference_logps\"].shape[0]) if batch[\"label\"][i] is True]\n            rejected_idx = [i for i in range(batch[\"reference_logps\"].shape[0]) if batch[\"label\"][i] is False]\n\n            reference_chosen_logps = batch[\"reference_logps\"][chosen_idx, ...]\n            reference_rejected_logps = batch[\"reference_logps\"][rejected_idx, ...]\n        else:\n            with torch.no_grad():\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        (\n                            reference_chosen_logps,\n                            reference_rejected_logps,\n                            _,\n                            _,\n                        ) = self.forward(self.model, batch)[:4]\n                else:\n                    (\n                        reference_chosen_logps,\n                        reference_rejected_logps,\n                        _,\n                        _,\n                    ) = self.forward(self.ref_model, batch)[:4]\n\n        chosen_embeddings, rejected_embeddings = self._get_prompt_embeddings(batch)\n\n        losses, chosen_rewards, rejected_rewards, delta = self.bco_loss(\n            policy_chosen_logps,\n            policy_rejected_logps,\n            reference_chosen_logps,\n            reference_rejected_logps,\n            chosen_embeddings,\n            rejected_embeddings,\n        )\n        metrics[\"delta\"] = delta.item()\n\n        num_chosen = torch.Tensor([len(chosen_rewards)]).to(self.accelerator.device)\n        num_rejected = torch.Tensor([len(rejected_rewards)]).to(self.accelerator.device)\n\n        all_num_chosen = self.accelerator.gather(num_chosen).sum().item()\n        all_num_rejected = self.accelerator.gather(num_rejected).sum().item()\n\n        if all_num_chosen > 0:\n            metrics[\"rewards/chosen_sum\"] = self.accelerator.gather(chosen_rewards.nansum()).nansum().item()\n            metrics[\"logps/chosen_sum\"] = self.accelerator.gather(policy_chosen_logps.nansum()).nansum().item()\n            metrics[\"count/chosen\"] = all_num_chosen\n\n        if all_num_rejected > 0:\n            metrics[\"rewards/rejected_sum\"] = self.accelerator.gather(rejected_rewards.nansum()).nansum().item()\n            metrics[\"logps/rejected_sum\"] = self.accelerator.gather(policy_rejected_logps.nansum()).nansum().item()\n            metrics[\"count/rejected\"] = all_num_rejected\n\n        loss = losses.nanmean()\n        if self.aux_loss_enabled:\n            loss += getattr(model.config, \"router_aux_loss_coef\", 0.0) * aux_loss\n\n        return loss, metrics\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        compute_loss_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with compute_loss_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs)\n\n        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:\n        loss = loss.to(self.args.device)\n        # force log the metrics\n        if self.accelerator.is_main_process:\n            self.store_metrics(metrics, train_eval=\"train\")\n\n        if return_outputs:\n            return (loss, metrics)\n        return loss\n\n    def store_metrics(self, metrics: Dict[str, float], train_eval: Literal[\"train\", \"eval\"] = \"train\") -> None:\n        for key, value in metrics.items():\n            self._stored_metrics[train_eval][key].append(value)\n\n    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:\n        if self.train_dataset is None or not has_length(self.train_dataset):\n            return None\n        return SequentialSampler(self.train_dataset)\n\n    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:\n        \"\"\"Generate samples from the model and reference model for the given batch of inputs.\"\"\"\n\n        # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with\n        # the torch cuda amp context manager as some hidden states are silently casted to full precision.\n        generate_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n        with generate_context_manager:\n            policy_output = model.generate(\n                input_ids=batch[\"prompt_input_ids\"],\n                attention_mask=batch[\"prompt_attention_mask\"],\n                max_length=self.max_length,\n                do_sample=True,\n                pad_token_id=self.tokenizer.pad_token_id,\n            )\n\n            # if reference_output in batch use that otherwise use the reference model\n            if \"reference_output\" in batch:\n                reference_output = batch[\"reference_output\"]\n            else:\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        reference_output = self.model.generate(\n                            input_ids=batch[\"prompt_input_ids\"],\n                            attention_mask=batch[\"prompt_attention_mask\"],\n                            max_length=self.max_length,\n                            do_sample=True,\n                            pad_token_id=self.tokenizer.pad_token_id,\n                        )\n                else:\n                    reference_output = self.ref_model.generate(\n                        input_ids=batch[\"prompt_input_ids\"],\n                        attention_mask=batch[\"prompt_attention_mask\"],\n                        max_length=self.max_length,\n                        do_sample=True,\n                        pad_token_id=self.tokenizer.pad_token_id,\n                    )\n\n        policy_output = pad_to_length(policy_output, self.max_length, self.tokenizer.pad_token_id)\n        policy_output_decoded = self.tokenizer.batch_decode(policy_output, skip_special_tokens=True)\n\n        reference_output = pad_to_length(reference_output, self.max_length, self.tokenizer.pad_token_id)\n        reference_output_decoded = self.tokenizer.batch_decode(reference_output, skip_special_tokens=True)\n\n        return policy_output_decoded, reference_output_decoded\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ):\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        if ignore_keys is None:\n            if hasattr(model, \"config\"):\n                ignore_keys = getattr(model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        prediction_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n        with torch.no_grad(), prediction_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs)\n\n        # force log the metrics\n        if self.accelerator.is_main_process:\n            self.store_metrics(metrics, train_eval=\"eval\")\n\n        if prediction_loss_only:\n            return (loss.detach(), None, None)\n\n        # logits for the chosen and rejected samples from model\n        logits_dict = {\n            \"eval_logits/chosen\": metrics[\"logits/chosen\"],\n            \"eval_logits/rejected\": metrics[\"logits/rejected\"],\n        }\n        logits = tuple(v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys)\n        logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)\n        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)\n\n        return (loss.detach(), logits, labels)\n\n    def evaluation_loop(\n        self,\n        dataloader: DataLoader,\n        description: str,\n        prediction_loss_only: Optional[bool] = None,\n        ignore_keys: Optional[List[str]] = None,\n        metric_key_prefix: str = \"eval\",\n    ) -> EvalLoopOutput:\n        \"\"\"\n        Overriding built-in evaluation loop to store metrics for each batch.\n        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.\n\n        Works both with or without labels.\n        \"\"\"\n\n        # Sample and save to game log if requested (for one batch to save time)\n        if self.generate_during_eval:\n            # Generate random indices within the range of the total number of samples\n            num_samples = len(dataloader.dataset)\n            random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)\n\n            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader\n            random_batch_dataset = dataloader.dataset.select(random_indices)\n            random_batch = self.data_collator(random_batch_dataset)\n            random_batch = self._prepare_inputs(random_batch)\n\n            target_indicies = [i for i in range(len(random_batch[\"delta\"])) if random_batch[\"delta\"][i] is False]\n            target_batch = {\n                \"prompt_input_ids\": itemgetter(*target_indicies)(random_batch[\"prompt_input_ids\"]),\n                \"prompt_attention_mask\": itemgetter(*target_indicies)(random_batch[\"prompt_attention_mask\"]),\n                \"prompt\": itemgetter(*target_indicies)(random_batch[\"prompt\"]),\n            }\n            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, target_batch)\n\n            self.log(\n                {\n                    \"game_log\": wandb.Table(\n                        columns=[\"Prompt\", \"Policy\", \"Ref Model\"],\n                        rows=[\n                            [prompt, pol[len(prompt) :], ref[len(prompt) :]]\n                            for prompt, pol, ref in zip(\n                                target_batch[\"prompt\"], policy_output_decoded, ref_output_decoded\n                            )\n                        ],\n                    )\n                }\n            )\n            self.state.log_history.pop()\n\n        # Base evaluation\n        initial_output = super().evaluation_loop(\n            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix\n        )\n\n        return initial_output\n\n    def log(self, logs: Dict[str, float]) -> None:\n        \"\"\"\n        Log `logs` on the various objects watching training, including stored metrics.\n\n        Args:\n            logs (`Dict[str, float]`):\n                The values to log.\n        \"\"\"\n        # logs either has 'loss' or 'eval_loss'\n        train_eval = \"train\" if \"loss\" in logs else \"eval\"\n        # train metrics should have no prefix, eval should have 'eval_'\n        prefix = \"eval_\" if train_eval == \"eval\" else \"\"\n        # accumulate average metrics from sums and lengths\n        for split in [\"chosen\", \"rejected\"]:\n            if f\"count/{split}\" in self._stored_metrics[train_eval]:\n                count_sum = torch.Tensor(self._stored_metrics[train_eval][f\"count/{split}\"]).sum().item()\n                logs[f\"{prefix}rewards/{split}\"] = (\n                    torch.Tensor(self._stored_metrics[train_eval][f\"rewards/{split}_sum\"]).sum().item() / count_sum\n                )\n                logs[f\"{prefix}logps/{split}\"] = (\n                    torch.Tensor(self._stored_metrics[train_eval][f\"logps/{split}_sum\"]).sum().item() / count_sum\n                )\n                for key in [f\"count/{split}\", f\"rewards/{split}_sum\", f\"logps/{split}_sum\"]:\n                    del self._stored_metrics[train_eval][key]\n        # calculate reward margin\n        if f\"{prefix}rewards/chosen\" in logs and f\"{prefix}rewards/rejected\" in logs:\n            logs[f\"{prefix}rewards/margins\"] = logs[f\"{prefix}rewards/chosen\"] - logs[f\"{prefix}rewards/rejected\"]\n        # Add averaged stored metrics to logs\n        for key, metrics in self._stored_metrics[train_eval].items():\n            logs[f\"{prefix}{key}\"] = torch.Tensor(metrics).mean().item()\n        del self._stored_metrics[train_eval]\n        return super().log(logs)\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"bco\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom datasets import Dataset, IterableDataset\nfrom transformers import PreTrainedTokenizerBase, TrainerCallback\nfrom transformers.modeling_utils import PreTrainedModel\nfrom transformers.trainer_utils import EvalPrediction\nfrom transformers.training_args import OptimizerNames\nfrom transformers.utils import is_apex_available\n\nfrom ..models.modeling_base import GeometricMixtureWrapper\nfrom ..models.utils import unwrap_model_for_generation\nfrom .nash_md_config import NashMDConfig\nfrom .online_dpo_trainer import OnlineDPOTrainer\nfrom .utils import empty_cache, get_reward, truncate_right\n\n\nif is_apex_available():\n    from apex import amp\n\n\nclass NashMDTrainer(OnlineDPOTrainer):\n    r\"\"\"\n    Initialize NashMDTrainer as a subclass of [`OnlineDPOConfig`].\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForCausalLM`.\n        ref_model (`PreTrainedModelWrapper`):\n            Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no\n            reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.\n        reward_model (`transformers.PreTrainedModel`):\n            The reward model to score completions with, preferably an `AutoModelForSequenceClassification`.\n        judge (`BasePairwiseJudge`):\n            The judge to use for pairwise comparison of model completions.\n        args (`NashMDConfig`):\n            The NashMD config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        peft_config (`Dict`):\n            The peft config to use for training.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"nash-md\"]\n\n    def __init__(\n        self,\n        model: Union[PreTrainedModel, nn.Module] = None,\n        ref_model: Union[PreTrainedModel, nn.Module] = None,\n        reward_model: Optional[nn.Module] = None,\n        args: Optional[NashMDConfig] = None,\n        data_collator: Optional[Callable] = None,\n        train_dataset: Optional[Union[Dataset, IterableDataset]] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n    ) -> None:\n        super().__init__(\n            model=model,\n            ref_model=ref_model,\n            reward_model=reward_model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            peft_config=peft_config,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        self._mixture_coef = self.args.mixture_coef\n\n        # Overwrite the stats dictionary to include NashMD specific statistics\n        self.stats = {\n            # Remove \"non_score_reward\", \"rlhf_reward\", \"scores_margin\"\n            # Add \"mixture_coef\"\n            \"loss/kl\": [],\n            \"objective/entropy\": [],\n            \"loss/score\": [],\n            \"rewards/chosen\": [],\n            \"rewards/rejected\": [],\n            \"rewards/accuracies\": [],\n            \"rewards/margins\": [],\n            \"logps/chosen\": [],\n            \"logps/rejected\": [],\n            \"val/model_contain_eos_token\": [],\n            \"val/ref_contain_eos_token\": [],\n            \"beta\": [],\n            \"mixture_coef\": [],\n        }\n\n    @property\n    def mixture_coef(self):\n        if isinstance(self._mixture_coef, list):\n            epoch = self.state.epoch\n            return self._mixture_coef[epoch] if epoch < len(self._mixture_coef) else self._mixture_coef[-1]\n        else:\n            return self._mixture_coef\n\n    def _generate_completions(self, model, prompts):\n        with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n            model_output = unwrapped_model.generate(\n                input_ids=prompts[\"input_ids\"],\n                attention_mask=prompts[\"attention_mask\"],\n                generation_config=self.generation_config,\n            )\n\n            ref_model = model if self.ref_model is None else self.ref_model\n            with torch.no_grad(), unwrap_model_for_generation(ref_model, self.accelerator) as unwrapped_ref_model:\n                mixture_model = GeometricMixtureWrapper(\n                    model=unwrapped_model,\n                    ref_model=unwrapped_ref_model,\n                    generation_config=self.generation_config,\n                    mixture_coef=self.mixture_coef,\n                    device=self.accelerator.device,\n                )\n\n                mixture_output = mixture_model.generate(\n                    input_ids=prompts[\"input_ids\"],\n                    attention_mask=prompts[\"attention_mask\"],\n                    generation_config=self.generation_config,\n                )\n\n        return model_output, mixture_output\n\n    def _process_completions(self, model_output, mixture_output, prompts):\n        context_length = prompts[\"input_ids\"].shape[1]\n\n        # Process model completions\n        model_completion_ids = model_output[:, context_length:]\n        model_completion_ids, model_completion_mask = truncate_right(\n            model_completion_ids, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id\n        )\n        model_data = {\n            \"input_ids\": torch.cat((prompts[\"input_ids\"], model_completion_ids), dim=1),\n            \"attention_mask\": torch.cat((prompts[\"attention_mask\"], model_completion_mask), dim=1),\n        }\n\n        # Process reference model completions\n        mixture_completion_ids = mixture_output[:, context_length:]\n        mixture_completion_ids, mixture_completion_mask = truncate_right(\n            mixture_completion_ids, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id\n        )\n        mixture_data = {\n            \"input_ids\": torch.cat((prompts[\"input_ids\"], mixture_completion_ids), dim=1),\n            \"attention_mask\": torch.cat((prompts[\"attention_mask\"], mixture_completion_mask), dim=1),\n        }\n\n        return model_data, mixture_data\n\n    def _compute_rewards(self, model_data, mixture_data, context_length):\n        with torch.no_grad():\n            _, model_scores, _ = get_reward(\n                self.reward_model, model_data[\"input_ids\"], self.tokenizer.pad_token_id, context_length\n            )\n            _, mixture_scores, _ = get_reward(\n                self.reward_model, mixture_data[\"input_ids\"], self.tokenizer.pad_token_id, context_length\n            )\n\n        # Apply EOS penalty if needed\n        if self.args.missing_eos_penalty is not None:\n            model_contain_eos = torch.any(model_data[\"input_ids\"] == self.tokenizer.eos_token_id, dim=-1)\n            mixture_contain_eos = torch.any(mixture_data[\"input_ids\"] == self.tokenizer.eos_token_id, dim=-1)\n            model_scores[~model_contain_eos] -= self.args.missing_eos_penalty\n            mixture_scores[~mixture_contain_eos] -= self.args.missing_eos_penalty\n\n        return model_scores, mixture_scores\n\n    def _compute_logprobs(self, model, model_data, context_length):\n        def compute_logprobs_for_data(m, data):\n            output = m(data[\"input_ids\"], attention_mask=data[\"attention_mask\"])\n            logits = output.logits[:, context_length - 1 : -1]\n            logprobs = F.log_softmax(logits, dim=-1)\n            token_logprobs = torch.gather(logprobs, 2, data[\"input_ids\"][:, context_length:].unsqueeze(-1)).squeeze(-1)\n            return token_logprobs\n\n        # Compute logprobs for model completions under the model\n        model_logprobs_model_data = compute_logprobs_for_data(model, model_data)\n\n        # Compute logprobs of model completions under the reference model\n        with torch.no_grad():\n            if self.ref_model is None:\n                with model.disable_adapter():\n                    ref_logprobs_model_data = compute_logprobs_for_data(model, model_data)\n            else:\n                ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data)\n\n        # Mask padding tokens\n        model_padding_mask = model_data[\"attention_mask\"][:, context_length:] == 0\n        model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0)\n        ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0)\n\n        return (model_logprobs_model_data, ref_logprobs_model_data)\n\n    def _compute_losses(\n        self,\n        model_logprobs_model_data,\n        ref_logprobs_model_data,\n        model_data_scores,\n        mixture_data_scores,\n    ):\n        # Compute log probs\n        model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)\n        ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)\n\n        # probability of the model data vs the mixture data\n        probability = F.sigmoid(model_data_scores - mixture_data_scores)\n\n        # reinforce score where 0.5 is a control variate\n        score = (probability - 0.5) * model_logprobs_model_data_sum\n\n        # kl divergence\n        kl_div = model_logprobs_model_data_sum - ref_logprobs_model_data_sum\n\n        # final loss\n        loss = self.beta * kl_div - score\n\n        return loss.mean(), score, kl_div\n\n    def _log_statistics(\n        self,\n        model_data,\n        mixture_data,\n        model_logprobs_model_data,\n        ref_logprobs_model_data,\n        model_scores,\n        mixture_scores,\n        score,\n        kl_div,\n        context_length,\n    ):\n        # Helper function to gather and compute mean\n        def gather_mean(tensor):\n            return self.accelerator.gather(tensor).mean().item()\n\n        # Log score\n        self.stats[\"loss/score\"].append(gather_mean(score))\n        # Log KL divergence\n        self.stats[\"loss/kl\"].append(gather_mean(kl_div))\n\n        # Log logprobs\n        model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)\n        ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)\n\n        self.stats[\"logps/chosen\"].append(gather_mean(model_logprobs_model_data_sum))\n        self.stats[\"logps/rejected\"].append(gather_mean(ref_logprobs_model_data_sum))\n\n        # Log rewards\n        self.stats[\"rewards/chosen\"].append(gather_mean(model_scores))\n        self.stats[\"rewards/rejected\"].append(gather_mean(mixture_scores))\n\n        # Calculate entropy for model data\n        entropy_model_data = -model_logprobs_model_data.sum(1)\n        self.stats[\"objective/entropy\"].append(gather_mean(entropy_model_data))\n\n        # Calculate margins\n        margin = model_scores - mixture_scores\n        self.stats[\"rewards/margins\"].append(gather_mean(margin))\n\n        # Calculate accuracy\n        accuracy = (margin > 0).float()\n        self.stats[\"rewards/accuracies\"].append(gather_mean(accuracy))\n\n        # Log EOS token statistics\n        model_eos = (model_data[\"input_ids\"][:, context_length:] == self.tokenizer.eos_token_id).any(dim=1)\n        mixture_eos = (mixture_data[\"input_ids\"][:, context_length:] == self.tokenizer.eos_token_id).any(dim=1)\n        self.stats[\"val/model_contain_eos_token\"].append(gather_mean(model_eos.float()))\n        self.stats[\"val/ref_contain_eos_token\"].append(gather_mean(mixture_eos.float()))\n\n        # Log beta and mixture coef\n        self.stats[\"beta\"].append(self.beta)\n        self.stats[\"mixture_coef\"].append(self.mixture_coef)\n\n    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:\n        model.train()\n\n        # need the prompt_ only\n        inputs = self._prepare_inputs(inputs)\n        context_length = inputs[\"prompt_input_ids\"].shape[1]\n        prompts = {\n            \"input_ids\": inputs[\"prompt_input_ids\"],\n            \"attention_mask\": inputs[\"prompt_attention_mask\"],\n        }\n        del inputs\n\n        # Sample completions from both the model and the reference model\n        model_output, mixture_output = self._generate_completions(model, prompts)\n\n        # Process model completions\n        model_data, mixture_data = self._process_completions(model_output, mixture_output, prompts)\n\n        # Compute rewards\n        model_data_scores, mixture_data_scores = self._compute_rewards(model_data, mixture_data, context_length)\n\n        # Compute logprobs\n        model_logprobs_model_data, ref_logprobs_model_data = self._compute_logprobs(model, model_data, context_length)\n\n        # Compute loss\n        loss, score, kl_div = self._compute_losses(\n            model_logprobs_model_data, ref_logprobs_model_data, model_data_scores, mixture_data_scores\n        )\n\n        # Log everything\n        self._log_statistics(\n            model_data,\n            mixture_data,\n            model_logprobs_model_data.detach(),\n            ref_logprobs_model_data,\n            model_data_scores,\n            mixture_data_scores,\n            score.detach(),\n            kl_div.detach(),\n            context_length,\n        )\n\n        if (\n            self.args.torch_empty_cache_steps is not None\n            and self.state.global_step % self.args.torch_empty_cache_steps == 0\n        ):\n            empty_cache()\n\n        kwargs = {}\n        # For LOMO optimizers you need to explicitly use the learning rate\n        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:\n            kwargs[\"learning_rate\"] = self._get_learning_rate()\n\n        if self.args.n_gpu > 1:\n            loss = loss.mean()  # mean() to average on multi-gpu parallel training\n\n        if self.use_apex:\n            with amp.scale_loss(loss, self.optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            self.accelerator.backward(loss, **kwargs)\n\n        return loss.detach() / self.args.gradient_accumulation_steps\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport os\nimport sys\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Literal, Optional\n\nimport numpy as np\nimport tyro\nfrom transformers import is_wandb_available\nfrom typing_extensions import Annotated\n\nfrom trl.trainer.utils import exact_div\n\nfrom ..core import flatten_dict\n\n\nJSONDict = Annotated[Optional[dict], tyro.conf.arg(metavar=\"JSON\", constructor=json.loads)]\n\n\n@dataclass\nclass PPOConfig:\n    r\"\"\"\n    Configuration class for the [`PPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[: -len(\".py\")]`):\n            Name of this experiment.\n        seed (`int`, *optional*, defaults to `0`):\n            Random seed.\n        log_with (`Optional[Literal[\"wandb\", \"tensorboard\"]]`, *optional*, defaults to `None`):\n            Log with either `\"wandb\"` or `\"tensorboard\"`. Check\n            [tracking](https://huggingface.co/docs/accelerate/usage_guides/tracking) for more details.\n        task_name (`Optional[str]`, *optional*, defaults to `None`):\n            Name of task to use - used only for tracking purposes.\n        model_name (`Optional[str]`, *optional*, defaults to `\"gpt2\"`):\n            Name of model to use - used only for tracking purposes.\n        query_dataset (`Optional[str]`, *optional*, defaults to `\"stanfordnlp/imdb\"`):\n            Name of dataset to query - used only for tracking purposes.\n        reward_model (`Optional[str]`, *optional*, defaults to `\"sentiment-analysis:lvwerra/distilbert-imdb\"`):\n            Reward model to use - used only for tracking purposes.\n        remove_unused_columns (`bool`, *optional*, defaults to `True`):\n            Remove unused columns from the dataset.\n        tracker_kwargs (`JSONDict`, *optional*, defaults to `{}`):\n            Keyword arguments for the tracker (e.g. `python ppo.py --tracker_kwargs='{\"wandb\": {\"entity\": \"my_wandb_entity\", \"name\": \"my_exp_name\"}}'`.\n        accelerator_kwargs (`JSONDict`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator.\n        project_kwargs (`JSONDict`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator project config (e.g. `logging_dir`).\n        tracker_project_name (`str`, *optional*, defaults to `\"trl\"`):\n            Name of project to use for tracking.\n        push_to_hub_if_best_kwargs (`JSONDict`, *optional*, defaults to `{}`):\n            Keyword arguments for pushing model to the hub during training (e.g. repo_id).\n        steps (`int`, *optional*, defaults to `20000`):\n            Number of training steps.\n        learning_rate (`float`, *optional*, defaults to `1.41e-5`):\n            Learning rate for the optimizer.\n        adap_kl_ctrl (`bool`, *optional*, defaults to `True`):\n            Use adaptive KL control, otherwise linear.\n        init_kl_coef (`Optional[float]`, *optional*, defaults to `0.2`):\n            Initial KL penalty coefficient (used for adaptive and linear control).\n        kl_penalty (`Literal[\"kl\", \"abs\", \"mse\", \"full\"]`, *optional*, defaults to `\"kl\"`):\n            kl penalty options. Possible values are:\n\n                - `\"kl\"`: model_logp - ref_logp\n                - `\"abs\"`: abs(kl)\n                - `\"mse\"`: mean squared error mse(kl)\n                - `\"full\"`: the actual kl for all tokens in the distribution.\n\n        target (`float`, *optional*, defaults to `6.0`):\n            Target KL value for adaptive KL control.\n        horizon (`float`, *optional*, defaults to `10000.0`):\n            Horizon for adaptive KL control.\n        gamma (`float`, *optional*, defaults to `1.0`):\n            Gamma parameter for advantage calculation.\n        lam (`float`, *optional*, defaults to `0.95`):\n            Lambda parameter for advantage calculation.\n        cliprange (`float`, *optional*, defaults to `0.2`):\n            Range for clipping in PPO policy gradient loss.\n        cliprange_value (`float`, *optional*, defaults to `0.2`):\n            Range for clipping values in loss calculation.\n        vf_coef (`float`, *optional*, defaults to `0.1`):\n            Scaling factor for value loss.\n        batch_size (`int`, *optional*, defaults to `128`):\n            Number of samples per optimisation step.\n        forward_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            DEPRECATED: use `mini_batch_size` instead, which does the same thing.\n        mini_batch_size (`int`, *optional*, defaults to `128`):\n            Number of samples optimized in each mini batch.\n        gradient_accumulation_steps (`int`, *optional*, defaults to `1`):\n            Number of gradient accumulation steps.\n        world_size (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for distributed training.\n        ppo_epochs (`int`, *optional*, defaults to `4`):\n            Number of optimisation epochs per batch of samples.\n        optimize_device_cache (`bool`, *optional*, defaults to `False`):\n            Optimize device cache for slightly more memory-efficient training.\n        early_stopping (`bool`, *optional*, defaults to `False`):\n            Whether to stop the PPO optimization loop early is the KL too high.\n        target_kl (`float`, *optional*, defaults to `1.0`):\n            Stop early if we exceed this value by over 50%.\n        compare_steps (`int`, *optional*, defaults to `1`):\n            Compare the current step with the previous `compare_steps` steps.\n        ratio_threshold (`float`, *optional*, defaults to `10.0`):\n            Skip mini-batches with high PPO ratios that can cause loss spikes.\n        use_score_scaling (`bool`, *optional*, defaults to `False`):\n            Use score scaling.\n        use_score_norm (`bool`, *optional*, defaults to `False`):\n            Use score normalization. Only applicable if `use_score_scaling` is True.\n        score_clip (`Optional[float]`, *optional*, defaults to `None`):\n            Score clipping.\n        whiten_rewards (`bool`, *optional*, defaults to `False`):\n            Whiten the rewards before computing advantages.\n        is_encoder_decoder (`Optional[bool]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        is_peft_model (`Optional[bool]`, *optional*, defaults to `None`):\n            Whether the model is a PEFT model.\n        backward_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Number of samples optimized in an `optimizer.step()` call.\n        global_backward_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Effective `backward_batch_size` across all processes.\n        global_batch_size (`Optional[int]`, *optional*, defaults to `None`):\n            Effective `batch_size` across all processes.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n    \"\"\"\n\n    exp_name: str = os.path.basename(sys.argv[0])[: -len(\".py\")]\n    seed: int = 0\n    log_with: Optional[Literal[\"wandb\", \"tensorboard\"]] = None\n    task_name: Optional[str] = None\n    model_name: str = \"gpt2\"\n    query_dataset: str = \"stanfordnlp/imdb\"\n    reward_model: str = \"sentiment-analysis:lvwerra/distilbert-imdb\"\n    remove_unused_columns: bool = True\n    tracker_kwargs: JSONDict = field(default_factory=dict)\n    accelerator_kwargs: JSONDict = field(default_factory=dict)\n    project_kwargs: JSONDict = field(default_factory=dict)\n    tracker_project_name: str = \"trl\"\n    push_to_hub_if_best_kwargs: JSONDict = field(default_factory=dict)\n    steps: int = 20000\n    learning_rate: float = 1.41e-5\n    adap_kl_ctrl: bool = True\n    init_kl_coef: float = 0.2\n    kl_penalty: Literal[\"kl\", \"abs\", \"mse\", \"full\"] = \"kl\"\n    target: float = 6.0\n    horizon: float = 10000.0\n    gamma: float = 1.0\n    lam: float = 0.95\n    cliprange: float = 0.2\n    cliprange_value: float = 0.2\n    vf_coef: float = 0.1\n    batch_size: int = 128\n    forward_batch_size: Optional[int] = None\n    mini_batch_size: int = 128\n    gradient_accumulation_steps: int = 1\n    world_size: tyro.conf.Suppress[int] = None\n    ppo_epochs: int = 4\n    max_grad_norm: Optional[float] = None\n    optimize_cuda_cache: Optional[bool] = None\n    optimize_device_cache: bool = False\n    early_stopping: bool = False\n    target_kl: float = 1.0\n    compare_steps: int = 1\n    ratio_threshold: float = 10.0\n    use_score_scaling: bool = False\n    use_score_norm: bool = False\n    score_clip: Optional[float] = None\n    whiten_rewards: bool = False\n    gradient_checkpointing: bool = False\n    is_encoder_decoder: Optional[tyro.conf.Suppress[bool]] = None\n    is_peft_model: Optional[tyro.conf.Suppress[bool]] = None\n    backward_batch_size: tyro.conf.Suppress[int] = None\n    global_backward_batch_size: Optional[tyro.conf.Suppress[int]] = None\n    global_batch_size: tyro.conf.Suppress[int] = None\n    dataset_num_proc: Optional[int] = None\n\n    if optimize_cuda_cache is not None:\n        warnings.warn(\n            \"The `optimize_cuda_cache` argument will be deprecated soon, please use `optimize_device_cache` instead.\"\n        )\n\n        if optimize_device_cache is True:\n            raise ValueError(\"Both `optimize_device_cache` and `optimize_cuda_cache` were provided\")\n\n        optimize_device_cache = optimize_cuda_cache\n\n    def __post_init__(self):\n        warnings.warn(\n            \"`PPOConfig` is deprecated and will be removed in the future. Please use `PPOv2Config` with `PPOv2Trainer` instead.\",\n            FutureWarning,\n        )\n        if self.forward_batch_size is not None:\n            warnings.warn(\n                \"Note that using `forward_batch_size` is deprecated, use `mini_batch_size` instead. By setting it you overwrite `mini_batch_size` which affects both the batch size during forward passes and also the mini batch size for PPO optimization.\"\n            )\n            self.mini_batch_size = self.forward_batch_size\n\n        self.backward_batch_size = self.mini_batch_size * self.gradient_accumulation_steps\n        exact_div(\n            self.batch_size,\n            self.backward_batch_size,\n            \"`batch_size` must be a multiple of `mini_batch_size * gradient_accumulation_steps`\",\n        )\n\n        # check if wandb is installed\n        if self.log_with == \"wandb\":\n            # raise error if wandb is not installed\n            if not is_wandb_available():\n                raise ImportError(\n                    \"Please install wandb to use wandb logging. You can do this by running `pip install wandb`.\"\n                )\n\n        self.total_ppo_epochs = int(np.ceil(self.steps / self.batch_size))\n        assert self.kl_penalty in [\"kl\", \"abs\", \"mse\", \"full\"]\n\n    def to_dict(self):\n        output_dict = {}\n        for key, value in self.__dict__.items():\n            output_dict[key] = value\n        return flatten_dict(output_dict)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Literal, Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass KTOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`KTOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        learning_rate (`float`, *optional*, defaults to `5e-7`):\n            Initial learning rate for [`AdamW`] optimizer. The default value replaces that of [`~transformers.TrainingArguments`].\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        max_prompt_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the prompt. This argument is required if you want to use the default data collator.\n        max_completion_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the completion. This argument is required if you want to use the default data collator\n            and your model is an encoder-decoder.\n        beta (`float`, *optional*, defaults to `0.1`):\n            Parameter controlling the deviation from the reference model. Higher β means less deviation from the\n            reference model.\n        loss_type (`str`, *optional*, defaults to `\"kto\"`):\n            Type of loss to use. Possible values are:\n\n                - `\"kto\"`: KTO loss from the [KTO](https://huggingface.co/papers/2402.01306) paper.\n                - `\"apo_zero_unpaired\"`: Unpaired variant of APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper.\n\n        desirable_weight (`float`, *optional*, defaults to `1.0`):\n            Desirable losses are weighed by this factor to counter unequal number of desirable and undesirable paris.\n        undesirable_weight (`float`, *optional*, defaults to `1.0`):\n            Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs.\n        label_pad_token_id (`int`, *optional*, defaults to `-100`):\n            Label pad token id. This argument is required if you want to use the default data collator.\n        padding_value (`Optional[int]`, *optional*, defaults to `None`):\n            Padding value to use. If `None`, the padding value of the tokenizer is used.\n        truncation_mode (`str`, *optional*, defaults to `\"keep_end\"`):\n            Truncation mode to use when the prompt is too long. Possible values are `\"keep_end\"` or `\"keep_start\"`.\n            This argument is required if you want to use the default data collator.\n        generate_during_eval (`bool`, *optional*, defaults to `False`):\n            If `True`, generates and logs completions from both the model and the reference model to W&B during\n            evaluation.\n        is_encoder_decoder (`Optional[bool]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):\n            Whether to precompute reference model log probabilities for training and evaluation datasets. This is\n            useful when training without the reference model to reduce the total GPU memory needed.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        ref_model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model\n            from a string.\n        dataset_num_proc: (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n    \"\"\"\n\n    learning_rate: float = 5e-7\n    max_length: Optional[int] = None\n    max_prompt_length: Optional[int] = None\n    max_completion_length: Optional[int] = None\n    beta: float = 0.1\n    loss_type: Literal[\"kto\", \"apo_zero_unpaired\"] = \"kto\"\n    desirable_weight: float = 1.0\n    undesirable_weight: float = 1.0\n    label_pad_token_id: int = -100\n    padding_value: Optional[int] = None\n    truncation_mode: str = \"keep_end\"\n    generate_during_eval: bool = False\n    is_encoder_decoder: Optional[bool] = None\n    precompute_ref_log_probs: bool = False\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    ref_model_init_kwargs: Optional[Dict[str, Any]] = None\n    dataset_num_proc: Optional[int] = None\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport sys\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Literal, Optional\n\nfrom transformers import is_bitsandbytes_available, is_torchvision_available\n\nfrom ..core import flatten_dict\n\n\n@dataclass\nclass DDPOConfig:\n    r\"\"\"\n    Configuration class for the [`DDPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(\".py\")]`):\n            Name of this experiment (by default is the file name without the extension name).\n        run_name (`str`, *optional*, defaults to `\"\"`):\n            Name of this run.\n        seed (`int`, *optional*, defaults to `0`):\n            Random seed.\n        log_with (`Optional[Literal[\"wandb\", \"tensorboard\"]]`, *optional*, defaults to `None`):\n            Log with either 'wandb' or 'tensorboard', check\n            https://huggingface.co/docs/accelerate/usage_guides/tracking for more details.\n        tracker_kwargs (`Dict`, *optional*, defaults to `{}`):\n            Keyword arguments for the tracker (e.g. wandb_project).\n        accelerator_kwargs (`Dict`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator.\n        project_kwargs (`Dict`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator project config (e.g. `logging_dir`).\n        tracker_project_name (`str`, *optional*, defaults to `\"trl\"`):\n            Name of project to use for tracking.\n        logdir (`str`, *optional*, defaults to `\"logs\"`):\n            Top-level logging directory for checkpoint saving.\n        num_epochs (`int`, *optional*, defaults to `100`):\n            Number of epochs to train.\n        save_freq (`int`, *optional*, defaults to `1`):\n            Number of epochs between saving model checkpoints.\n        num_checkpoint_limit (`int`, *optional*, defaults to `5`):\n            Number of checkpoints to keep before overwriting old ones.\n        mixed_precision (`str`, *optional*, defaults to `\"fp16\"`):\n            Mixed precision training.\n        allow_tf32 (`bool`, *optional*, defaults to `True`):\n            Allow `tf32` on Ampere GPUs.\n        resume_from (`str`, *optional*, defaults to `\"\"`):\n            Resume training from a checkpoint.\n        sample_num_steps (`int`, *optional*, defaults to `50`):\n            Number of sampler inference steps.\n        sample_eta (`float`, *optional*, defaults to `1.0`):\n            Eta parameter for the DDIM sampler.\n        sample_guidance_scale (`float`, *optional*, defaults to `5.0`):\n            Classifier-free guidance weight.\n        sample_batch_size (`int`, *optional*, defaults to `1`):\n            Batch size (per GPU) to use for sampling.\n        sample_num_batches_per_epoch (`int`, *optional*, defaults to `2`):\n            Number of batches to sample per epoch.\n        train_batch_size (`int`, *optional*, defaults to `1`):\n            Batch size (per GPU) to use for training.\n        train_use_8bit_adam (`bool`, *optional*, defaults to `False`):\n            Use 8bit Adam optimizer from bitsandbytes.\n        train_learning_rate (`float`, *optional*, defaults to `3e-4`):\n            Learning rate.\n        train_adam_beta1 (`float`, *optional*, defaults to `0.9`):\n            Adam beta1.\n        train_adam_beta2 (`float`, *optional*, defaults to `0.999`):\n            Adam beta2.\n        train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`):\n            Adam weight decay.\n        train_adam_epsilon (`float`, *optional*, defaults to `1e-8`):\n            Adam epsilon.\n        train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`):\n            Number of gradient accumulation steps.\n        train_max_grad_norm (`float`, *optional*, defaults to `1.0`):\n            Maximum gradient norm for gradient clipping.\n        train_num_inner_epochs (`int`, *optional*, defaults to `1`):\n            Number of inner epochs per outer epoch.\n        train_cfg (`bool`, *optional*, defaults to `True`):\n            Whether or not to use classifier-free guidance during training.\n        train_adv_clip_max (`float`, *optional*, defaults to `5.0`):\n            Clip advantages to the range.\n        train_clip_range (`float`, *optional*, defaults to `1e-4`):\n            PPO clip range.\n        train_timestep_fraction (`float`, *optional*, defaults to `1.0`):\n            Fraction of timesteps to train on.\n        per_prompt_stat_tracking (`bool`, *optional*, defaults to `False`):\n            Whether to track statistics for each prompt separately.\n        per_prompt_stat_tracking_buffer_size (`int`, *optional*, defaults to `16`):\n            Number of reward values to store in the buffer for each prompt.\n        per_prompt_stat_tracking_min_count (`int`, *optional*, defaults to `16`):\n            Minimum number of reward values to store in the buffer.\n        async_reward_computation (`bool`, *optional*, defaults to `False`):\n            Whether to compute rewards asynchronously.\n        max_workers (`int`, *optional*, defaults to `2`):\n            Maximum number of workers to use for async reward computation.\n        negative_prompts (`Optional[str]`, *optional*, defaults to `\"\"`):\n            Comma-separated list of prompts to use as negative examples.\n    \"\"\"\n\n    exp_name: str = os.path.basename(sys.argv[0])[: -len(\".py\")]\n    run_name: str = \"\"\n    seed: int = 0\n    log_with: Optional[Literal[\"wandb\", \"tensorboard\"]] = None\n    tracker_kwargs: dict = field(default_factory=dict)\n    accelerator_kwargs: dict = field(default_factory=dict)\n    project_kwargs: dict = field(default_factory=dict)\n    tracker_project_name: str = \"trl\"\n    logdir: str = \"logs\"\n    num_epochs: int = 100\n    save_freq: int = 1\n    num_checkpoint_limit: int = 5\n    mixed_precision: str = \"fp16\"\n    allow_tf32: bool = True\n    resume_from: str = \"\"\n    sample_num_steps: int = 50\n    sample_eta: float = 1.0\n    sample_guidance_scale: float = 5.0\n    sample_batch_size: int = 1\n    sample_num_batches_per_epoch: int = 2\n    train_batch_size: int = 1\n    train_use_8bit_adam: bool = False\n    train_learning_rate: float = 3e-4\n    train_adam_beta1: float = 0.9\n    train_adam_beta2: float = 0.999\n    train_adam_weight_decay: float = 1e-4\n    train_adam_epsilon: float = 1e-8\n    train_gradient_accumulation_steps: int = 1\n    train_max_grad_norm: float = 1.0\n    train_num_inner_epochs: int = 1\n    train_cfg: bool = True\n    train_adv_clip_max: float = 5.0\n    train_clip_range: float = 1e-4\n    train_timestep_fraction: float = 1.0\n    per_prompt_stat_tracking: bool = False\n    per_prompt_stat_tracking_buffer_size: int = 16\n    per_prompt_stat_tracking_min_count: int = 16\n    async_reward_computation: bool = False\n    max_workers: int = 2\n    negative_prompts: str = \"\"\n\n    def to_dict(self):\n        output_dict = {}\n        for key, value in self.__dict__.items():\n            output_dict[key] = value\n        return flatten_dict(output_dict)\n\n    def __post_init__(self):\n        if self.log_with not in [\"wandb\", \"tensorboard\"]:\n            warnings.warn(\n                \"Accelerator tracking only supports image logging if `log_with` is set to 'wandb' or 'tensorboard'.\"\n            )\n\n        if self.log_with == \"wandb\" and not is_torchvision_available():\n            warnings.warn(\"Wandb image logging requires torchvision to be installed\")\n\n        if self.train_use_8bit_adam and not is_bitsandbytes_available():\n            raise ImportError(\n                \"You need to install bitsandbytes to use 8bit Adam. \"\n                \"You can install it with `pip install bitsandbytes`.\"\n            )\n\n\n# KTO Authors: Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager, nullcontext\nfrom copy import deepcopy\nfrom functools import wraps\nfrom operator import itemgetter\nfrom typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.amp as amp\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import PartialState\nfrom accelerate.utils import is_deepspeed_available, tqdm\nfrom datasets import Dataset, concatenate_datasets\nfrom torch.utils.data import DataLoader, SequentialSampler\nfrom transformers import (\n    AutoModelForCausalLM,\n    DataCollator,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainingArguments,\n    is_wandb_available,\n)\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalLoopOutput, has_length\nfrom transformers.utils import is_peft_available\n\nfrom ..models import PreTrainedModelWrapper, create_reference_model\nfrom .kto_config import KTOConfig\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    disable_dropout_in_model,\n    pad_to_length,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\n\nif is_wandb_available():\n    import wandb\n\nif is_deepspeed_available():\n    import deepspeed\n\nif TYPE_CHECKING:\n    from transformers import PreTrainedModel, PreTrainedTokenizer\n\nRUNNING_NAME = \"running.pt\"\n\n\ndef _get_kl_dataset(batch: Dict[str, List[Any]]) -> Dict[str, List[Any]]:\n    \"\"\"Creates mismatched pairs of prompts and completions for the KL dataset by adding a +1 offset to the order of completions.\"\"\"\n    batch[\"answer_input_ids\"] = [batch[\"answer_input_ids\"][-1]] + batch[\"answer_input_ids\"][:-1]\n    batch[\"answer_attention_mask\"] = [batch[\"answer_attention_mask\"][-1]] + batch[\"answer_attention_mask\"][:-1]\n    return batch\n\n\ndef _tokenize(\n    batch: Dict[str, List[Any]],\n    tokenizer: \"PreTrainedTokenizer\",\n) -> Dict[str, List[Any]]:\n    \"\"\"Tokenize a batch from a KTO specific dataset.\"\"\"\n    prompt_tokenized = tokenizer(batch[\"prompt\"], add_special_tokens=False)\n    prompt_input_ids = prompt_tokenized[\"input_ids\"]\n    prompt_attention_mask = prompt_tokenized[\"attention_mask\"]\n    prompt_and_completion = [prompt + completion for prompt, completion in zip(batch[\"prompt\"], batch[\"completion\"])]\n    full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False)\n    full_input_ids = full_tokenized[\"input_ids\"]\n    full_attention_mask = full_tokenized[\"attention_mask\"]\n\n    answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids)]\n    answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask)]\n\n    # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`\n    full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids)]\n    # Prepare input tokens for token by token comparison\n    full_input_ids = [np.array(f) for f in full_input_ids]\n    for full, concat in zip(full_input_ids, full_concat_input_ids):\n        if len(full) != len(concat):\n            raise ValueError(\"Prompt input ids and answer input ids should have the same length.\")\n\n    # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens\n    # can be merged together when tokenizing prompt+answer. This could result\n    # on the last token from the prompt being different when tokenized on its own\n    # vs when done as prompt+answer.\n    response_token_ids_start_idx = [len(p) for p in prompt_input_ids]\n\n    # If tokenized prompt is different than both prompt+answer, then it means the\n    # last token has changed due to merging.\n    for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx)):\n        if not np.array_equal(p, f[:r]):\n            response_token_ids_start_idx[idx] -= 1\n\n    prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx)]\n    prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx)]\n\n    for p, m in zip(prompt_input_ids, prompt_attention_mask):\n        if len(p) != len(m):\n            raise ValueError(\"Prompt input ids and attention mask should have the same length.\")\n\n    answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx)]\n    answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx)]\n\n    output = dict(\n        prompt_input_ids=prompt_input_ids,\n        prompt_attention_mask=prompt_attention_mask,\n        answer_input_ids=answer_input_ids,\n        answer_attention_mask=answer_attention_mask,\n    )\n\n    return output\n\n\ndef _process_tokens(example: Dict[str, Any], model: \"PreTrainedModel\" = None, **kwargs) -> Dict:\n    \"\"\"Process tokens of a KTO specific dataset.\n\n    At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation\n    in case the prompt + completion responses is/are too long. First\n    we truncate the prompt; if we're still too long, we truncate the completion.\n\n    We also create the labels for the completion responses, which are of length equal to\n    the sum of the length of the prompt and the completion response, with\n    label_pad_token_id  for the prompt tokens.\n    \"\"\"\n    prompt = example[\"prompt\"]\n    completion = example[\"completion\"]\n\n    batch = {\n        f\"{kwargs['prefix']}prompt\": prompt,\n        f\"{kwargs['prefix']}completion\": completion,\n        f\"{kwargs['prefix']}label\": example[\"label\"],\n    }\n\n    if not kwargs[\"is_encoder_decoder\"]:\n        # Check issues below for more details\n        #  1. https://github.com/huggingface/trl/issues/907\n        #  2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n        #  3. https://github.com/LianjiaTech/BELLE/issues/337\n\n        if not isinstance(prompt, str):\n            raise ValueError(f\"prompt should be an str but got {type(prompt)}\")\n\n        if not isinstance(completion, str):\n            raise ValueError(f\"completion should be an str but got {type(completion)}\")\n\n        # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer\n        all_tokens = {\n            \"prompt_input_ids\": example[\"prompt_input_ids\"],\n            \"prompt_attention_mask\": example[\"prompt_attention_mask\"],\n            \"answer_input_ids\": example[\"answer_input_ids\"],\n            \"answer_attention_mask\": example[\"answer_attention_mask\"],\n        }\n\n        # calculate max length by checking if BOS/EOS is already there\n        max_length = kwargs[\"max_length\"]\n        bos_token_id = kwargs[\"tokenizer\"].bos_token_id\n        eos_token_id = kwargs[\"tokenizer\"].eos_token_id\n        if bos_token_id != all_tokens[\"prompt_input_ids\"][0]:\n            max_length -= 1\n        if eos_token_id != all_tokens[\"answer_input_ids\"][-1]:\n            max_length -= 1\n\n        # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the prompt\n        if len(all_tokens[\"prompt_input_ids\"]) + len(all_tokens[\"answer_input_ids\"]) > max_length:\n            for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                if kwargs[\"truncation_mode\"] == \"keep_start\":\n                    all_tokens[k] = all_tokens[k][: kwargs[\"max_prompt_length\"]]\n                elif kwargs[\"truncation_mode\"] == \"keep_end\":\n                    all_tokens[k] = all_tokens[k][-kwargs[\"max_prompt_length\"] :]\n                else:\n                    raise ValueError(f\"Unknown truncation mode: {kwargs['truncation_mode']}\")\n\n        # if that's still too long, truncate the response\n        if len(all_tokens[\"prompt_input_ids\"]) + len(all_tokens[\"answer_input_ids\"]) > max_length:\n            for k in [\"answer_input_ids\", \"answer_attention_mask\"]:\n                all_tokens[k] = all_tokens[k][: max_length - kwargs[\"max_prompt_length\"]]\n\n        # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens\n        batch[f\"{kwargs['prefix']}prompt_input_ids\"] = all_tokens[\"prompt_input_ids\"]\n        batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = all_tokens[\"prompt_attention_mask\"]\n        batch[f\"{kwargs['prefix']}completion_input_ids\"] = (\n            all_tokens[\"prompt_input_ids\"] + all_tokens[\"answer_input_ids\"]\n        )\n        batch[f\"{kwargs['prefix']}completion_attention_mask\"] = (\n            all_tokens[\"prompt_attention_mask\"] + all_tokens[\"answer_attention_mask\"]\n        )\n\n        # add BOS, which affects both prompt and the full completion\n        if len(all_tokens[\"prompt_input_ids\"]) == 0 or bos_token_id != all_tokens[\"prompt_input_ids\"][0]:\n            batch[f\"{kwargs['prefix']}prompt_input_ids\"] = [bos_token_id] + batch[\n                f\"{kwargs['prefix']}prompt_input_ids\"\n            ]\n            batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = [1] + batch[f\"{kwargs['prefix']}prompt_attention_mask\"]\n            batch[f\"{kwargs['prefix']}completion_input_ids\"] = [bos_token_id] + batch[\n                f\"{kwargs['prefix']}completion_input_ids\"\n            ]\n            batch[f\"{kwargs['prefix']}completion_attention_mask\"] = [1] + batch[\n                f\"{kwargs['prefix']}completion_attention_mask\"\n            ]\n        # add EOS, which affects only the full completion\n        if len(all_tokens[\"answer_input_ids\"]) == 0 or eos_token_id != all_tokens[\"answer_input_ids\"][-1]:\n            batch[f\"{kwargs['prefix']}completion_input_ids\"] = batch[f\"{kwargs['prefix']}completion_input_ids\"] + [\n                eos_token_id\n            ]\n            batch[f\"{kwargs['prefix']}completion_attention_mask\"] = batch[\n                f\"{kwargs['prefix']}completion_attention_mask\"\n            ] + [1]\n\n        batch[f\"{kwargs['prefix']}completion_labels\"] = batch[f\"{kwargs['prefix']}completion_input_ids\"][:]\n        batch[f\"{kwargs['prefix']}completion_labels\"][: len(batch[f\"{kwargs['prefix']}prompt_input_ids\"])] = [\n            kwargs[\"label_pad_token_id\"]\n        ] * len(batch[f\"{kwargs['prefix']}prompt_input_ids\"])\n    else:\n        completion_tokens = kwargs[\"tokenizer\"](\n            completion, truncation=True, max_length=kwargs[\"max_completion_length\"], add_special_tokens=True\n        )\n        prompt_tokens = kwargs[\"tokenizer\"](\n            prompt, truncation=True, max_length=kwargs[\"max_prompt_length\"], add_special_tokens=True\n        )\n\n        batch[f\"{kwargs['prefix']}prompt_input_ids\"] = prompt_tokens[\"input_ids\"]\n        batch[f\"{kwargs['prefix']}prompt_attention_mask\"] = prompt_tokens[\"attention_mask\"]\n\n        batch[f\"{kwargs['prefix']}completion_labels\"] = completion_tokens[\"input_ids\"]\n        batch[f\"{kwargs['prefix']}completion_attention_mask\"] = completion_tokens[\"attention_mask\"]\n        if model is not None and hasattr(model, \"prepare_decoder_input_ids_from_labels\"):\n            batch[f\"{kwargs['prefix']}completion_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                labels=torch.tensor(batch[\"completion_labels\"])\n            )\n\n    return batch\n\n\nclass KTOTrainer(Trainer):\n    r\"\"\"\n    Initialize KTOTrainer.\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForSequenceClassification`.\n        ref_model (`PreTrainedModelWrapper`):\n            Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no\n            reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.\n        args (`KTOConfig`):\n            The arguments to use for training.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        data_collator (`transformers.DataCollator`, *optional*, defaults to `None`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Dict`, defaults to `None`):\n            The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        disable_dropout (`bool`, defaults to `True`):\n            Whether or not to disable dropouts in `model` and `ref_model`.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n        model_adapter_name (`str`, defaults to `None`):\n            Name of the train target PEFT adapter, when using LoRA with multiple adapters.\n        ref_adapter_name (`str`, defaults to `None`):\n            Name of the reference PEFT adapter, when using LoRA with multiple adapters.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"kto\"]\n\n    def __init__(\n        self,\n        model: Union[PreTrainedModel, nn.Module, str] = None,\n        ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        args: KTOConfig = None,\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        data_collator: Optional[DataCollator] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n        model_adapter_name: Optional[str] = None,\n        ref_adapter_name: Optional[str] = None,\n    ):\n        if type(args) is TrainingArguments:\n            raise ValueError(\"Please use `KTOConfig` instead TrainingArguments.\")\n\n        if not isinstance(model, str) and ref_model is model:\n            raise ValueError(\n                \"`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the \"\n                \"same as `model`, you must mass a copy of it, or `None` if you use peft.\"\n            )\n\n        if args.model_init_kwargs is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\"You passed model_kwargs to the KTOTrainer. But your model is already instantiated.\")\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if args.ref_model_init_kwargs is None:\n            ref_model_init_kwargs = {}\n        elif not isinstance(ref_model, str):\n            raise ValueError(\n                \"You passed ref_model_kwargs to the KTOTrainer. But your ref_model is already instantiated.\"\n            )\n        else:\n            ref_model_init_kwargs = args.ref_model_init_kwargs\n            torch_dtype = ref_model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                ref_model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the KTOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        if isinstance(ref_model, str):\n            warnings.warn(\n                \"You passed a ref model_id to the KTOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM`\"\n            )\n            ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs)\n\n        # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`\n        # has been called in order to properly call autocast if needed.\n        self._peft_has_been_casted_to_bf16 = False\n\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it with `pip install peft` to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            # if model is a peft model and we have a peft_config, we merge and unload it first\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_loaded_in_4bit\", False):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n\n                prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                if _support_gc_kwargs:\n                    prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n            elif getattr(args, \"gradient_checkpointing\", False):\n                # For backward compatibility with older versions of transformers\n                if hasattr(model, \"enable_input_require_grads\"):\n                    model.enable_input_require_grads()\n                else:\n\n                    def make_inputs_require_grad(module, input, output):\n                        output.requires_grad_(True)\n\n                    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            # get peft model with the given config\n            model = get_peft_model(model, peft_config)\n            if args.bf16 and getattr(model, \"is_loaded_in_4bit\", False):\n                peft_module_casting_to_bf16(model)\n                # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager\n                self._peft_has_been_casted_to_bf16 = True\n\n        # For models that use gradient_checkpointing, we need to attach a hook that enables input\n        # to explicitly have `requires_grad=True`, otherwise training will either silently\n        # fail or completely fail.\n        elif getattr(args, \"gradient_checkpointing\", False):\n            # For backward compatibility with older versions of transformers\n            if hasattr(model, \"enable_input_require_grads\"):\n                model.enable_input_require_grads()\n            else:\n\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        if args.generate_during_eval and not is_wandb_available():\n            raise ValueError(\n                \"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install with `pip install wandb` to resolve.\"\n            )\n\n        if model is not None:\n            self.is_encoder_decoder = model.config.is_encoder_decoder\n        elif args.is_encoder_decoder is None:\n            raise ValueError(\"When no model is provided, you need to pass the parameter is_encoder_decoder.\")\n        else:\n            self.is_encoder_decoder = args.is_encoder_decoder\n\n        self.is_peft_model = is_peft_available() and isinstance(model, PeftModel)\n        self.model_adapter_name = model_adapter_name\n        self.ref_adapter_name = ref_adapter_name\n\n        if ref_model:\n            self.ref_model = ref_model\n        elif self.is_peft_model or args.precompute_ref_log_probs:\n            # The `model` with adapters turned off will be used as the reference model\n            self.ref_model = None\n        else:\n            self.ref_model = create_reference_model(model)\n\n        if tokenizer is None:\n            raise ValueError(\n                \"max_length or a tokenizer must be specified when using the default DPODataCollatorWithPadding\"\n            )\n        if args.max_length is None:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding, you should set `max_length` in the KTOTrainer's init\"\n                \" it will be set to `512` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_length = 512\n        if args.max_length is not None:\n            max_length = args.max_length\n\n        if args.max_prompt_length is None:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the KTOTrainer's init\"\n                \" it will be set to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_prompt_length = 128\n        if args.max_prompt_length is not None:\n            max_prompt_length = args.max_prompt_length\n\n        max_completion_length = None\n        if args.max_completion_length is None and self.is_encoder_decoder:\n            warnings.warn(\n                \"When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the KTOTrainer's init\"\n                \" it will be set to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_completion_length = 128\n        if args.max_completion_length is not None and self.is_encoder_decoder:\n            max_completion_length = args.max_completion_length\n\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(\n                pad_token_id=tokenizer.pad_token_id,\n                label_pad_token_id=args.label_pad_token_id,\n                is_encoder_decoder=self.is_encoder_decoder,\n            )\n\n            if args.remove_unused_columns:\n                args.remove_unused_columns = False\n                # warn users\n                warnings.warn(\n                    \"When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your KTOConfig\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_dpo_data_collator = True\n        else:\n            self.use_dpo_data_collator = False\n\n        # disable dropout in the model and reference model\n        disable_dropout_in_model(model)\n        if self.ref_model is not None:\n            disable_dropout_in_model(self.ref_model)\n\n        self.loss_type = args.loss_type\n        self.max_length = max_length\n        self.generate_during_eval = args.generate_during_eval\n        self.label_pad_token_id = args.label_pad_token_id\n        self.padding_value = args.padding_value if args.padding_value is not None else tokenizer.pad_token_id\n        self.max_prompt_length = max_prompt_length\n        self.truncation_mode = args.truncation_mode\n        self.max_completion_length = max_completion_length\n        self.tokenizer = tokenizer\n        self.precompute_ref_log_probs = args.precompute_ref_log_probs\n\n        # Not all losses require a KL calculation\n        self.calculate_KL = True\n        if self.loss_type in [\"apo_zero_unpaired\"]:\n            self.calculate_KL = False\n\n        # Since ref_logs are precomputed on the first call to get_train/eval_dataloader\n        # keep track of first called to avoid computation of future calls\n        self._precomputed_train_ref_log_probs = False\n        self._precomputed_eval_ref_log_probs = False\n\n        # metric\n        self._stored_metrics = defaultdict(lambda: defaultdict(list))\n\n        # KTO parameter\n        self.beta = args.beta\n        self.desirable_weight = args.desirable_weight\n        self.undesirable_weight = args.undesirable_weight\n        self.aux_loss_enabled = getattr(model.config, \"output_router_logits\", False)\n\n        with PartialState().local_main_process_first():\n            # Shuffle the datasets\n            train_dataset = train_dataset.shuffle(seed=args.data_seed)\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.shuffle(seed=args.data_seed)\n\n            # Tokenize and prepare the training datasets\n            train_dataset = train_dataset.map(\n                _tokenize,\n                batched=True,\n                fn_kwargs={\"tokenizer\": self.tokenizer},\n                num_proc=args.dataset_num_proc,\n                desc=\"Tokenizing train dataset\",\n            )\n\n            fn_kwargs = {\n                \"prefix\": \"\",\n                \"is_encoder_decoder\": self.is_encoder_decoder,\n                \"tokenizer\": self.tokenizer,\n                \"max_length\": self.max_length,\n                \"truncation_mode\": self.truncation_mode,\n                \"label_pad_token_id\": self.label_pad_token_id,\n                \"max_prompt_length\": self.max_prompt_length,\n                \"max_completion_length\": self.max_completion_length,\n            }\n\n            train_dataset = train_dataset.map(\n                _process_tokens,\n                fn_kwargs=fn_kwargs,\n                num_proc=args.dataset_num_proc,\n                desc=\"Processing tokenized train dataset\",\n            )\n\n            # Tokenize and prepare the eval datasets\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(\n                    _tokenize,\n                    fn_kwargs={\"tokenizer\": self.tokenizer},\n                    batched=True,\n                    num_proc=args.dataset_num_proc,\n                    desc=\"Tokenizing eval dataset\",\n                )\n\n                eval_dataset = eval_dataset.map(\n                    _process_tokens,\n                    fn_kwargs=fn_kwargs,\n                    num_proc=args.dataset_num_proc,\n                    desc=\"Processing tokenized eval dataset\",\n                )\n\n            # Get KL datasets if needed\n            if self.calculate_KL:\n                total_batch_size = (\n                    max(torch.cuda.device_count(), 1)\n                    * args.per_device_train_batch_size\n                    * args.gradient_accumulation_steps\n                )\n                if total_batch_size <= 1:\n                    raise ValueError(\n                        \"Batch size is 1 (too small). KTO will not work properly because the KL term will be equivalent to the implied reward.\"\n                    )\n\n                # create pairs for estimating the KL term by flipping the matched pairs in each batch of size total_batch_size\n                # i.e., (x_1, y_1), ..., (x_n, y_n) --> (x_1, y_n), ..., (x_n, y_1) = (x'_1, y'_1), ..., (x'_n, y'_n)\n                train_kl_dataset = train_dataset.map(\n                    _get_kl_dataset,\n                    batched=True,\n                    batch_size=total_batch_size,\n                    num_proc=args.dataset_num_proc,\n                    desc=\"Extracting KL train dataset\",\n                )\n\n                fn_kwargs[\"prefix\"] = \"KL_\"\n                train_kl_dataset = train_kl_dataset.map(\n                    _process_tokens,\n                    fn_kwargs=fn_kwargs,\n                    num_proc=args.dataset_num_proc,\n                    remove_columns=[c for c in train_kl_dataset.column_names if c in train_dataset.column_names],\n                    desc=\"Processing tokenized train KL dataset\",\n                )\n\n                # merge the datasets\n                train_dataset = concatenate_datasets([train_dataset, train_kl_dataset], axis=1)\n\n                if eval_dataset is not None:\n                    # Get KL dataset\n                    eval_kl_dataset = eval_dataset.map(\n                        _get_kl_dataset,\n                        batched=True,\n                        batch_size=total_batch_size,\n                        num_proc=args.dataset_num_proc,\n                        desc=\"Extracting eval KL dataset\",\n                    )\n\n                    eval_kl_dataset = eval_kl_dataset.map(\n                        _process_tokens,\n                        fn_kwargs=fn_kwargs,\n                        num_proc=args.dataset_num_proc,\n                        remove_columns=[c for c in eval_kl_dataset.column_names if c in eval_dataset.column_names],\n                        desc=\"Processing tokenized eval KL dataset\",\n                    )\n\n                    # merge the datasets\n                    eval_dataset = concatenate_datasets([eval_dataset, eval_kl_dataset], axis=1)\n\n            # calculate dataset desirability balance\n            num_desirable = max(sum(train_dataset[\"label\"]), 1)\n            num_undesirable = max(len(train_dataset[\"label\"]) - num_desirable, 1)  # \"label\" is binary\n\n            if num_desirable != num_undesirable:\n                # The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306\n                des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2)\n                des_weight_upper_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2)\n                und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2)\n                und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2)\n\n                des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound\n                und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound\n\n                if not (des_weight_in_range or und_weight_in_range):\n                    warnings.warn(\n                        f\"\"\"\n                        You have different amounts of desirable/positive and undesirable/negative examples but the\n                        weights on the desirable and undesirable losses don't seem to be in an ideal range. Based\n                        on your data, we recommend EITHER desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}]\n                        or undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH).\n                        See the documentation on how to optimally set these weights.\"\"\",\n                        UserWarning,\n                    )\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n        # Deepspeed Zero-3 does not support precompute_ref_log_probs\n        if self.is_deepspeed_enabled:\n            if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs:\n                raise ValueError(\n                    \"You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`.\"\n                )\n\n        if self.ref_model is None:\n            if not (self.is_peft_model or self.precompute_ref_log_probs):\n                raise ValueError(\n                    \"No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`\"\n                )\n        else:\n            if self.is_deepspeed_enabled:\n                self.ref_model = self._prepare_deepspeed(self.ref_model)\n            else:\n                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)\n\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n    @contextmanager\n    def null_ref_context(self):\n        \"\"\"Context manager for handling null reference model (that is, peft adapter manipulation).\"\"\"\n        with self.accelerator.unwrap_model(\n            self.model\n        ).disable_adapter() if self.is_peft_model and not self.ref_adapter_name else nullcontext():\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.ref_adapter_name)\n            yield\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.model_adapter_name or \"default\")\n\n    def get_train_dataloader(self) -> DataLoader:\n        \"\"\"\n        Returns the training [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`.\n        \"\"\"\n\n        if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_train_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params))\n            reference_completion_logps = []\n            reference_KL_logps = []\n\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Train dataset reference log probs\"):\n                reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch)\n\n                reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp)\n                reference_completion_logps.append(reference_completion_logp.cpu())\n\n                if self.calculate_KL:\n                    reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp)\n                    reference_KL_logps.append(reference_KL_logp.cpu())\n\n            self.train_dataset = self.train_dataset.add_column(\n                name=\"reference_logps\", column=torch.cat(reference_completion_logps).float().numpy()\n            )\n\n            if self.calculate_KL:\n                self.train_dataset = self.train_dataset.add_column(\n                    name=\"reference_KL_logps\", column=torch.cat(reference_KL_logps).float().numpy()\n                )\n\n            self._precomputed_train_ref_log_probs = True\n\n        return super().get_train_dataloader()\n\n    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:\n        \"\"\"\n        Returns the evaluation [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`.\n\n        Args:\n            eval_dataset (`torch.utils.data.Dataset`, *optional*):\n                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted\n                by the `model.forward()` method are automatically removed. It must implement `__len__`.\n        \"\"\"\n        if eval_dataset is None and self.eval_dataset is None:\n            raise ValueError(\"Trainer: evaluation requires an eval_dataset.\")\n        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset\n\n        if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_eval_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))\n\n            reference_completion_logps = []\n            reference_KL_logps = []\n\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Eval dataset reference log probs\"):\n                reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch)\n\n                reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp)\n                reference_completion_logps.append(reference_completion_logp.cpu())\n\n                if self.calculate_KL:\n                    reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp)\n                    reference_KL_logps.append(reference_KL_logp.cpu())\n\n            eval_dataset = eval_dataset.add_column(\n                name=\"reference_logps\", column=torch.cat(reference_completion_logps).float().numpy()\n            )\n            if self.calculate_KL:\n                eval_dataset = eval_dataset.add_column(\n                    name=\"reference_KL_logps\", column=torch.cat(reference_KL_logps).float().numpy()\n                )\n\n            # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs\n            if self.eval_dataset is not None:\n                self.eval_dataset = eval_dataset\n            self._precomputed_eval_ref_log_probs = True\n\n        return super().get_eval_dataloader(eval_dataset=eval_dataset)\n\n    def compute_reference_log_probs(self, padded_batch: Dict) -> Dict:\n        \"\"\"Computes log probabilities of the reference model for a single padded batch of a KTO specific dataset.\"\"\"\n        with torch.no_grad():\n            if self.ref_model is None:\n                with self.null_ref_context():\n                    if self.is_encoder_decoder:\n                        completion_logits = self.model(\n                            padded_batch[\"prompt_input_ids\"],\n                            attention_mask=padded_batch[\"prompt_attention_mask\"],\n                            decoder_input_ids=padded_batch.get(\"completion_decoder_input_ids\"),\n                            labels=padded_batch[\"completion_labels\"],\n                        ).logits\n\n                        if self.calculate_KL:\n                            KL_logits = self.model(\n                                padded_batch[\"KL_prompt_input_ids\"],\n                                attention_mask=padded_batch[\"KL_prompt_attention_mask\"],\n                                decoder_input_ids=padded_batch.get(\"KL_completion_decoder_input_ids\"),\n                                labels=padded_batch[\"KL_completion_labels\"],\n                            ).logits\n                    else:\n                        completion_logits = self.model(\n                            padded_batch[\"completion_input_ids\"],\n                            attention_mask=padded_batch[\"completion_attention_mask\"],\n                        ).logits\n\n                        if self.calculate_KL:\n                            KL_logits = self.model(\n                                padded_batch[\"KL_completion_input_ids\"],\n                                attention_mask=padded_batch[\"KL_completion_attention_mask\"],\n                            ).logits\n            else:\n                if self.is_encoder_decoder:\n                    completion_logits = self.ref_model(\n                        padded_batch[\"prompt_input_ids\"],\n                        attention_mask=padded_batch[\"prompt_attention_mask\"],\n                        decoder_input_ids=padded_batch.get(\"completion_decoder_input_ids\"),\n                        labels=padded_batch[\"completion_labels\"],\n                    ).logits\n\n                    if self.calculate_KL:\n                        KL_logits = self.ref_model(\n                            padded_batch[\"KL_prompt_input_ids\"],\n                            attention_mask=padded_batch[\"KL_prompt_attention_mask\"],\n                            decoder_input_ids=padded_batch.get(\"KL_completion_decoder_input_ids\"),\n                            labels=padded_batch[\"KL_completion_labels\"],\n                        ).logits\n                else:\n                    completion_logits = self.ref_model(\n                        padded_batch[\"completion_input_ids\"], attention_mask=padded_batch[\"completion_attention_mask\"]\n                    ).logits\n\n                    if self.calculate_KL:\n                        KL_logits = self.ref_model(\n                            padded_batch[\"KL_completion_input_ids\"],\n                            attention_mask=padded_batch[\"KL_completion_attention_mask\"],\n                        ).logits\n\n        completion_logps = self.get_batch_logps(\n            completion_logits,\n            padded_batch[\"completion_labels\"],\n            average_log_prob=False,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        if self.calculate_KL:\n            KL_logps = self.get_batch_logps(\n                KL_logits,\n                padded_batch[\"KL_completion_labels\"],\n                average_log_prob=False,\n                is_encoder_decoder=self.is_encoder_decoder,\n                label_pad_token_id=self.label_pad_token_id,\n            )\n        else:\n            KL_logps = None\n\n        return completion_logps, KL_logps\n\n    @staticmethod\n    def get_batch_logps(\n        logits: torch.FloatTensor,\n        labels: torch.LongTensor,\n        average_log_prob: bool = False,\n        label_pad_token_id: int = -100,\n        is_encoder_decoder: bool = False,\n    ) -> torch.FloatTensor:\n        \"\"\"Compute the log probabilities of the given labels under the given logits.\n\n        Args:\n            logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n            labels: Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are ignored. Shape: (batch_size, sequence_length)\n            average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the (non-masked) tokens.\n\n        Returns:\n            A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the given logits.\n        \"\"\"\n        if logits.shape[:-1] != labels.shape:\n            raise ValueError(\"Logits (batch and sequence length dim) and labels must have the same shape.\")\n\n        if not is_encoder_decoder:\n            labels = labels[:, 1:].clone()\n            logits = logits[:, :-1, :]\n        else:\n            # Fixes end-dec RuntimeError\n            labels = labels.clone()\n\n        loss_mask = labels != label_pad_token_id\n\n        # dummy token; we'll ignore the losses on these tokens later\n        labels[labels == label_pad_token_id] = 0\n\n        per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n\n        if average_log_prob:\n            return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n        else:\n            return (per_token_logps * loss_mask).sum(-1)\n\n    def forward(\n        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        if self.calculate_KL:\n            KL_logps = None\n            KL_model_kwargs = (\n                {\n                    \"input_ids\": batch[\"KL_prompt_input_ids\"],\n                    \"attention_mask\": batch[\"KL_prompt_attention_mask\"],\n                    \"labels\": batch[\"KL_completion_labels\"],\n                    \"decoder_input_ids\": batch.get(\"KL_completion_decoder_input_ids\"),\n                }\n                if self.is_encoder_decoder\n                else {\n                    \"input_ids\": batch[\"KL_completion_input_ids\"],\n                    \"attention_mask\": batch[\"KL_completion_attention_mask\"],\n                }\n            )\n            with torch.no_grad():\n                KL_logits = model(\n                    **KL_model_kwargs,\n                ).logits\n\n            KL_logps = self.get_batch_logps(\n                KL_logits,\n                batch[\"KL_completion_labels\"],\n                average_log_prob=False,\n                is_encoder_decoder=self.is_encoder_decoder,\n                label_pad_token_id=self.label_pad_token_id,\n            )\n        else:\n            KL_logps = None\n\n        model_kwargs = (\n            {\n                \"labels\": batch[\"completion_labels\"],\n                \"decoder_input_ids\": batch.get(\"completion_decoder_input_ids\"),\n            }\n            if self.is_encoder_decoder\n            else {}\n        )\n        if self.aux_loss_enabled:\n            model_kwargs[\"output_router_logits\"] = True\n\n        outputs = model(\n            batch[\"completion_input_ids\"],\n            attention_mask=batch[\"completion_attention_mask\"],\n            **model_kwargs,\n        )\n        completion_logits = outputs.logits\n\n        completion_logps = self.get_batch_logps(\n            completion_logits,\n            batch[\"completion_labels\"],\n            average_log_prob=False,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        if completion_logps.shape[0] != len(batch[\"label\"]):\n            raise ValueError(\n                \"There is a mismatch between the number of examples in this batch and the number of \"\n                \"examples for which an output sequence was predicted.\"\n            )\n\n        chosen_idx = [i for i in range(completion_logps.shape[0]) if batch[\"label\"][i] is True]\n        rejected_idx = [i for i in range(completion_logps.shape[0]) if batch[\"label\"][i] is False]\n\n        chosen_logps = completion_logps[chosen_idx, ...]\n        rejected_logps = completion_logps[rejected_idx, ...]\n\n        chosen_logits = completion_logits[chosen_idx, ...]\n        rejected_logits = completion_logits[rejected_idx, ...]\n\n        if self.aux_loss_enabled:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps, outputs.aux_loss)\n        else:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps)\n\n    def kto_loss(\n        self,\n        policy_chosen_logps: torch.FloatTensor,\n        policy_rejected_logps: torch.FloatTensor,\n        policy_KL_logps: torch.FloatTensor,\n        reference_chosen_logps: torch.FloatTensor,\n        reference_rejected_logps: torch.FloatTensor,\n        reference_KL_logps: torch.FloatTensor,\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Compute the KTO loss for a batch of policy and reference model log probabilities.\n\n        Args:\n            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,)\n            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,)\n            policy_KL_logps: Log probabilities of the policy model for the KL responses. Shape: (batch_size,)\n            reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,)\n            reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in batch_size,)\n            reference_KL_logps: Log probabilities of the reference model for the KL responses. Shape: (batch_size,)\n\n        Returns:\n            A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, KL).\n            The losses tensor contains the KTO loss for each example in the batch.\n            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.\n            The KL tensor contains the detached KL divergence estimate between the policy and reference models.\n        \"\"\"\n        if self.calculate_KL:\n            kl = (policy_KL_logps - reference_KL_logps).mean().detach()\n            kl = self.accelerator.gather(kl).mean().clamp(min=0)\n        else:\n            kl = torch.zeros(1).to(policy_chosen_logps.device)\n\n        # Chosen losses\n        if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0:\n            chosen_logratios = policy_chosen_logps - reference_chosen_logps\n\n            if self.loss_type == \"kto\":\n                # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306)\n                chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl))\n            elif self.loss_type == \"apo_zero_unpaired\":\n                # Unpaired variant of Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266)\n                # Use this loss when you believe the chosen outputs are better than your model's default output\n                chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios)\n\n            chosen_rewards = self.beta * chosen_logratios.detach()\n\n        else:\n            # lists can't be empty -- if they are, then accelerate.gather will hang\n            chosen_losses = torch.Tensor([]).to(self.accelerator.device)\n            chosen_rewards = torch.Tensor([]).to(self.accelerator.device)\n\n        # Rejected losses\n        if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0:\n            rejected_logratios = policy_rejected_logps - reference_rejected_logps\n\n            if self.loss_type == \"kto\":\n                rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios))\n            elif self.loss_type == \"apo_zero_unpaired\":\n                rejected_losses = F.sigmoid(self.beta * rejected_logratios)\n\n            rejected_rewards = self.beta * rejected_logratios.detach()\n        else:\n            # lists can't be empty -- if they are, then accelerate.gather will hang\n            rejected_losses = torch.Tensor([]).to(self.accelerator.device)\n            rejected_rewards = torch.Tensor([]).to(self.accelerator.device)\n\n        losses = torch.cat(\n            (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses),\n            0,\n        )\n\n        return losses, chosen_rewards, rejected_rewards, kl\n\n    def get_batch_loss_metrics(\n        self,\n        model,\n        batch: Dict[str, Union[List, torch.LongTensor]],\n    ):\n        \"\"\"Compute the KTO loss and other metrics for the given batch of inputs for train or test.\"\"\"\n        metrics = {}\n        batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}\n\n        forward_output = self.forward(model, batch)\n        (\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_chosen_logits,\n            policy_rejected_logits,\n            policy_KL_logps,\n        ) = forward_output[:5]\n        if self.aux_loss_enabled:\n            aux_loss = forward_output[5]\n\n        # if reference_logps in batch use them, otherwise use the reference model\n        if \"reference_logps\" in batch:\n            chosen_idx = [i for i in range(batch[\"reference_logps\"].shape[0]) if batch[\"label\"][i] is True]\n            rejected_idx = [i for i in range(batch[\"reference_logps\"].shape[0]) if batch[\"label\"][i] is False]\n\n            reference_chosen_logps = batch[\"reference_logps\"][chosen_idx, ...]\n            reference_rejected_logps = batch[\"reference_logps\"][rejected_idx, ...]\n            if self.calculate_KL:\n                reference_KL_logps = batch[\"reference_KL_logps\"]\n            else:\n                reference_KL_logps = None\n        else:\n            with torch.no_grad():\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        (\n                            reference_chosen_logps,\n                            reference_rejected_logps,\n                            _,\n                            _,\n                            reference_KL_logps,\n                        ) = self.forward(self.model, batch)[:5]\n                else:\n                    (\n                        reference_chosen_logps,\n                        reference_rejected_logps,\n                        _,\n                        _,\n                        reference_KL_logps,\n                    ) = self.forward(self.ref_model, batch)[:5]\n\n        losses, chosen_rewards, rejected_rewards, kl = self.kto_loss(\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_KL_logps,\n            reference_chosen_logps,\n            reference_rejected_logps,\n            reference_KL_logps,\n        )\n        metrics[\"kl\"] = kl.item()\n\n        num_chosen = torch.Tensor([len(chosen_rewards)]).to(self.accelerator.device)\n        num_rejected = torch.Tensor([len(rejected_rewards)]).to(self.accelerator.device)\n\n        all_num_chosen = self.accelerator.gather(num_chosen).sum().item()\n        all_num_rejected = self.accelerator.gather(num_rejected).sum().item()\n\n        if all_num_chosen > 0:\n            metrics[\"rewards/chosen_sum\"] = self.accelerator.gather(chosen_rewards.nansum()).nansum().item()\n            metrics[\"logps/chosen_sum\"] = self.accelerator.gather(policy_chosen_logps.nansum()).nansum().item()\n            metrics[\"logits/chosen\"] = (\n                self.accelerator.gather(policy_chosen_logits.nansum()).nansum().item() / all_num_chosen\n            )\n            metrics[\"count/chosen\"] = all_num_chosen\n\n        if all_num_rejected > 0:\n            metrics[\"rewards/rejected_sum\"] = self.accelerator.gather(rejected_rewards.nansum()).nansum().item()\n            metrics[\"logps/rejected_sum\"] = self.accelerator.gather(policy_rejected_logps.nansum()).nansum().item()\n            metrics[\"logits/rejected\"] = (\n                self.accelerator.gather(policy_rejected_logits.nansum()).nansum().item() / all_num_rejected\n            )\n            metrics[\"count/rejected\"] = all_num_rejected\n\n        loss = losses.nanmean()\n        if self.aux_loss_enabled:\n            loss += getattr(model.config, \"router_aux_loss_coef\", 0.0) * aux_loss\n\n        return loss, metrics\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        compute_loss_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with compute_loss_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs)\n\n        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:\n        loss = loss.to(self.args.device)\n        # force log the metrics\n        if self.accelerator.is_main_process:\n            self.store_metrics(metrics, train_eval=\"train\")\n\n        if return_outputs:\n            return (loss, metrics)\n        return loss\n\n    def store_metrics(self, metrics: Dict[str, float], train_eval: Literal[\"train\", \"eval\"] = \"train\") -> None:\n        for key, value in metrics.items():\n            self._stored_metrics[train_eval][key].append(value)\n\n    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:\n        if self.train_dataset is None or not has_length(self.train_dataset):\n            return None\n        return SequentialSampler(self.train_dataset)\n\n    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:\n        \"\"\"Generate samples from the model and reference model for the given batch of inputs.\"\"\"\n\n        # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with\n        # the torch cuda amp context manager as some hidden states are silently casted to full precision.\n        generate_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with generate_context_manager:\n            policy_output = model.generate(\n                input_ids=batch[\"prompt_input_ids\"],\n                attention_mask=batch[\"prompt_attention_mask\"],\n                max_length=self.max_length,\n                do_sample=True,\n                pad_token_id=self.tokenizer.pad_token_id,\n            )\n\n            # if reference_output in batch use that otherwise use the reference model\n            if \"reference_output\" in batch:\n                reference_output = batch[\"reference_output\"]\n            else:\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        reference_output = self.model.generate(\n                            input_ids=batch[\"prompt_input_ids\"],\n                            attention_mask=batch[\"prompt_attention_mask\"],\n                            max_length=self.max_length,\n                            do_sample=True,\n                            pad_token_id=self.tokenizer.pad_token_id,\n                        )\n                else:\n                    reference_output = self.ref_model.generate(\n                        input_ids=batch[\"prompt_input_ids\"],\n                        attention_mask=batch[\"prompt_attention_mask\"],\n                        max_length=self.max_length,\n                        do_sample=True,\n                        pad_token_id=self.tokenizer.pad_token_id,\n                    )\n\n        policy_output = pad_to_length(policy_output, self.max_length, self.tokenizer.pad_token_id)\n        policy_output_decoded = self.tokenizer.batch_decode(policy_output, skip_special_tokens=True)\n\n        reference_output = pad_to_length(reference_output, self.max_length, self.tokenizer.pad_token_id)\n        reference_output_decoded = self.tokenizer.batch_decode(reference_output, skip_special_tokens=True)\n\n        return policy_output_decoded, reference_output_decoded\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ):\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        if ignore_keys is None:\n            if hasattr(model, \"config\"):\n                ignore_keys = getattr(model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        prediction_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n        with torch.no_grad(), prediction_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs)\n\n        # force log the metrics\n        if self.accelerator.is_main_process:\n            self.store_metrics(metrics, train_eval=\"eval\")\n\n        if prediction_loss_only:\n            return (loss.detach(), None, None)\n\n        # logits for the chosen and rejected samples from model\n        logits_dict = {\n            \"eval_logits/chosen\": metrics[\"logits/chosen\"],\n            \"eval_logits/rejected\": metrics[\"logits/rejected\"],\n        }\n        logits = torch.tensor(\n            [v for k, v in logits_dict.items() if k not in ignore_keys], device=self.accelerator.device\n        )\n        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)\n\n        return (loss.detach(), logits, labels)\n\n    def evaluation_loop(\n        self,\n        dataloader: DataLoader,\n        description: str,\n        prediction_loss_only: Optional[bool] = None,\n        ignore_keys: Optional[List[str]] = None,\n        metric_key_prefix: str = \"eval\",\n    ) -> EvalLoopOutput:\n        \"\"\"\n        Overriding built-in evaluation loop to store metrics for each batch.\n        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.\n\n        Works both with or without labels.\n        \"\"\"\n\n        # Sample and save to game log if requested (for one batch to save time)\n        if self.generate_during_eval:\n            # Generate random indices within the range of the total number of samples\n            num_samples = len(dataloader.dataset)\n            random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)\n\n            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader\n            random_batch_dataset = dataloader.dataset.select(random_indices)\n            random_batch = self.data_collator(random_batch_dataset)\n            random_batch = self._prepare_inputs(random_batch)\n\n            target_indicies = [i for i in range(len(random_batch[\"label\"])) if random_batch[\"label\"][i] is False]\n            target_batch = {\n                \"prompt_input_ids\": random_batch[\"prompt_input_ids\"][target_indicies],\n                \"prompt_attention_mask\": random_batch[\"prompt_attention_mask\"][target_indicies],\n                \"prompt\": itemgetter(*target_indicies)(random_batch[\"prompt\"]),\n            }\n            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, target_batch)\n\n            self.log(\n                {\n                    \"game_log\": wandb.Table(\n                        columns=[\"Prompt\", \"Policy\", \"Ref Model\"],\n                        rows=[\n                            [prompt, pol[len(prompt) :], ref[len(prompt) :]]\n                            for prompt, pol, ref in zip(\n                                target_batch[\"prompt\"], policy_output_decoded, ref_output_decoded\n                            )\n                        ],\n                    )\n                }\n            )\n            self.state.log_history.pop()\n\n        # Base evaluation\n        initial_output = super().evaluation_loop(\n            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix\n        )\n\n        return initial_output\n\n    def log(self, logs: Dict[str, float]) -> None:\n        \"\"\"\n        Log `logs` on the various objects watching training, including stored metrics.\n\n        Args:\n            logs (`Dict[str, float]`):\n                The values to log.\n        \"\"\"\n        # logs either has 'loss' or 'eval_loss'\n        train_eval = \"train\" if \"loss\" in logs else \"eval\"\n        # train metrics should have no prefix, eval should have 'eval_'\n        prefix = \"eval_\" if train_eval == \"eval\" else \"\"\n        # accumulate average metrics from sums and lengths\n        for split in [\"chosen\", \"rejected\"]:\n            if f\"count/{split}\" in self._stored_metrics[train_eval]:\n                count_sum = torch.Tensor(self._stored_metrics[train_eval][f\"count/{split}\"]).sum().item()\n                logs[f\"{prefix}rewards/{split}\"] = (\n                    torch.Tensor(self._stored_metrics[train_eval][f\"rewards/{split}_sum\"]).sum().item() / count_sum\n                )\n                logs[f\"{prefix}logps/{split}\"] = (\n                    torch.Tensor(self._stored_metrics[train_eval][f\"logps/{split}_sum\"]).sum().item() / count_sum\n                )\n                for key in [f\"count/{split}\", f\"rewards/{split}_sum\", f\"logps/{split}_sum\"]:\n                    del self._stored_metrics[train_eval][key]\n        # calculate reward margin\n        if f\"{prefix}rewards/chosen\" in logs and f\"{prefix}rewards/rejected\" in logs:\n            logs[f\"{prefix}rewards/margins\"] = logs[f\"{prefix}rewards/chosen\"] - logs[f\"{prefix}rewards/rejected\"]\n        # Add averaged stored metrics to logs\n        for key, metrics in self._stored_metrics[train_eval].items():\n            logs[f\"{prefix}{key}\"] = torch.Tensor(metrics).mean().item()\n        del self._stored_metrics[train_eval]\n        return super().log(logs)\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"kto\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport sys\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, Literal, Optional, Tuple\n\nfrom transformers import is_bitsandbytes_available, is_torchvision_available\n\nfrom ..core import flatten_dict\n\n\n@dataclass\nclass AlignPropConfig:\n    r\"\"\"\n    Configuration class for the [`AlignPropTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(\".py\")]`):\n            Name of this experiment (defaults to the file name without the extension).\n        run_name (`str`, *optional*, defaults to `\"\"`):\n            Name of this run.\n        log_with (`Optional[Literal[\"wandb\", \"tensorboard\"]]`, *optional*, defaults to `None`):\n            Log with either `\"wandb\"` or `\"tensorboard\"`. Check\n            [tracking](https://huggingface.co/docs/accelerate/usage_guides/tracking) for more details.\n        log_image_freq (`int`, *optional*, defaults to `1`):\n            Frequency for logging images.\n        tracker_kwargs (`Dict[str, Any]`, *optional*, defaults to `{}`):\n            Keyword arguments for the tracker (e.g., `wandb_project`).\n        accelerator_kwargs (`Dict[str, Any]`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator.\n        project_kwargs (`Dict[str, Any]`, *optional*, defaults to `{}`):\n            Keyword arguments for the accelerator project config (e.g., `logging_dir`).\n        tracker_project_name (`str`, *optional*, defaults to `\"trl\"`):\n            Name of project to use for tracking.\n        logdir (`str`, *optional*, defaults to `\"logs\"`):\n            Top-level logging directory for checkpoint saving.\n        num_epochs (`int`, *optional*, defaults to `100`):\n            Number of epochs to train.\n        save_freq (`int`, *optional*, defaults to `1`):\n            Number of epochs between saving model checkpoints.\n        num_checkpoint_limit (`int`, *optional*, defaults to `5`):\n            Number of checkpoints to keep before overwriting old ones.\n        mixed_precision (`str`, *optional*, defaults to `\"fp16\"`):\n            Mixed precision training.\n        allow_tf32 (`bool`, *optional*, defaults to `True`):\n            Allow `tf32` on Ampere GPUs.\n        resume_from (`str`, *optional*, defaults to `\"\"`):\n            Path to resume training from a checkpoint.\n        sample_num_steps (`int`, *optional*, defaults to `50`):\n            Number of sampler inference steps.\n        sample_eta (`float`, *optional*, defaults to `1.0`):\n            Eta parameter for the DDIM sampler.\n        sample_guidance_scale (`float`, *optional*, defaults to `5.0`):\n            Classifier-free guidance weight.\n        train_use_8bit_adam (`bool`, *optional*, defaults to `False`):\n            Whether to use the 8bit Adam optimizer from `bitsandbytes`.\n        train_learning_rate (`float`, *optional*, defaults to `1e-3`):\n            Learning rate.\n        train_adam_beta1 (`float`, *optional*, defaults to `0.9`):\n            Beta1 for Adam optimizer.\n        train_adam_beta2 (`float`, *optional*, defaults to `0.999`):\n            Beta2 for Adam optimizer.\n        train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`):\n            Weight decay for Adam optimizer.\n        train_adam_epsilon (`float`, *optional*, defaults to `1e-8`):\n            Epsilon value for Adam optimizer.\n        train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`):\n            Number of gradient accumulation steps.\n        train_max_grad_norm (`float`, *optional*, defaults to `1.0`):\n            Maximum gradient norm for gradient clipping.\n        negative_prompts (`Optional[str]`, *optional*, defaults to `None`):\n            Comma-separated list of prompts to use as negative examples.\n        truncated_backprop_rand (`bool`, *optional*, defaults to `True`):\n            If `True`, randomized truncation to different diffusion timesteps is used.\n        truncated_backprop_timestep (`int`, *optional*, defaults to `49`):\n            Absolute timestep to which the gradients are backpropagated. Used only if `truncated_backprop_rand=False`.\n        truncated_rand_backprop_minmax (`Tuple[int, int]`, *optional*, defaults to `(0, 50)`):\n            Range of diffusion timesteps for randomized truncated backpropagation.\n    \"\"\"\n\n    exp_name: str = os.path.basename(sys.argv[0])[: -len(\".py\")]\n    run_name: str = \"\"\n    seed: int = 0\n    log_with: Optional[Literal[\"wandb\", \"tensorboard\"]] = None\n    log_image_freq: int = 1\n    tracker_kwargs: Dict[str, Any] = field(default_factory=dict)\n    accelerator_kwargs: Dict[str, Any] = field(default_factory=dict)\n    project_kwargs: Dict[str, Any] = field(default_factory=dict)\n    tracker_project_name: str = \"trl\"\n    logdir: str = \"logs\"\n    num_epochs: int = 100\n    save_freq: int = 1\n    num_checkpoint_limit: int = 5\n    mixed_precision: str = \"fp16\"\n    allow_tf32: bool = True\n    resume_from: str = \"\"\n    sample_num_steps: int = 50\n    sample_eta: float = 1.0\n    sample_guidance_scale: float = 5.0\n    train_batch_size: int = 1\n    train_use_8bit_adam: bool = False\n    train_learning_rate: float = 1e-3\n    train_adam_beta1: float = 0.9\n    train_adam_beta2: float = 0.999\n    train_adam_weight_decay: float = 1e-4\n    train_adam_epsilon: float = 1e-8\n    train_gradient_accumulation_steps: int = 1\n    train_max_grad_norm: float = 1.0\n    negative_prompts: Optional[str] = None\n    truncated_backprop_rand: bool = True\n    truncated_backprop_timestep: int = 49\n    truncated_rand_backprop_minmax: Tuple[int, int] = (0, 50)\n\n    def to_dict(self):\n        output_dict = {}\n        for key, value in self.__dict__.items():\n            output_dict[key] = value\n        return flatten_dict(output_dict)\n\n    def __post_init__(self):\n        if self.log_with not in [\"wandb\", \"tensorboard\"]:\n            warnings.warn(\n                \"Accelerator tracking only supports image logging if `log_with` is set to 'wandb' or 'tensorboard'.\"\n            )\n\n        if self.log_with == \"wandb\" and not is_torchvision_available():\n            warnings.warn(\"Wandb image logging requires torchvision to be installed\")\n\n        if self.train_use_8bit_adam and not is_bitsandbytes_available():\n            raise ImportError(\n                \"You need to install bitsandbytes to use 8bit Adam. \"\n                \"You can install it with `pip install bitsandbytes`.\"\n            )\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport dataclasses\nimport inspect\nimport warnings\nfrom functools import wraps\nfrom typing import Callable, Dict, List, Optional, Tuple, Union\n\nimport datasets\nimport torch\nimport torch.nn as nn\nfrom accelerate.state import PartialState\nfrom datasets import Dataset\nfrom datasets.arrow_writer import SchemaInferenceError\nfrom datasets.builder import DatasetGenerationError\nfrom huggingface_hub.utils._deprecation import _deprecate_arguments\nfrom transformers import (\n    AutoModelForCausalLM,\n    AutoTokenizer,\n    DataCollator,\n    DataCollatorForLanguageModeling,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n)\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalPrediction\nfrom transformers.utils import is_peft_available\n\nfrom ..extras.dataset_formatting import get_formatting_func_from_dataset\nfrom ..import_utils import is_liger_kernel_available\nfrom .sft_config import SFTConfig\nfrom .utils import (\n    ConstantLengthDataset,\n    DataCollatorForCompletionOnlyLM,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training\n\nif is_liger_kernel_available():\n    from liger_kernel.transformers import AutoLigerKernelForCausalLM\n\n\nclass SFTTrainer(Trainer):\n    r\"\"\"\n    Class definition of the Supervised Finetuning Trainer (SFT Trainer).\n    This class is a wrapper around the `transformers.Trainer` class and inherits all of its attributes and methods.\n    The trainer takes care of properly initializing the PeftModel in case a user passes a `PeftConfig` object.\n\n    Args:\n        model (Union[`transformers.PreTrainedModel`, `nn.Module`, `str`]):\n            The model to train, can be a `PreTrainedModel`, a `torch.nn.Module` or a string with the model name to\n            load from cache or download. The model can be also converted to a `PeftModel` if a `PeftConfig` object is\n            passed to the `peft_config` argument.\n        args (`Optional[SFTConfig]`):\n            The arguments to tweak for training. Will default to a basic instance of [`SFTConfig`] with the `output_dir`\n            set to a directory named *tmp_trainer* in the current directory if not provided.\n        data_collator (`Optional[transformers.DataCollator]`):\n            The data collator to use for training.\n        train_dataset (`Optional[datasets.Dataset]`):\n            The dataset to use for training. We recommend users to use `trl.trainer.ConstantLengthDataset` to create their dataset.\n        eval_dataset (Optional[Union[`datasets.Dataset`, Dict[`str`, `datasets.Dataset`]]]):\n            The dataset to use for evaluation. We recommend users to use `trl.trainer.ConstantLengthDataset` to create their dataset.\n        tokenizer (`Optional[transformers.PreTrainedTokenizer]`):\n            The tokenizer to use for training. If not specified, the tokenizer associated to the model will be used.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        compute_metrics (`Callable[[transformers.EvalPrediction], Dict]`, *optional* defaults to None):\n            The function used to compute metrics during evaluation. It should return a dictionary mapping metric names to metric values.\n            If not specified, only the loss will be computed during evaluation.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Optional[PeftConfig]`):\n            The PeftConfig object to use to initialize the PeftModel.\n        formatting_func (`Optional[Callable]`):\n            The formatting function to be used for creating the `ConstantLengthDataset`.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"sft\"]\n\n    @_deprecate_arguments(\n        version=\"1.0.0\",\n        deprecated_args=[\n            \"dataset_text_field\",\n            \"packing\",\n            \"max_seq_length\",\n            \"dataset_num_proc\",\n            \"dataset_batch_size\",\n            \"neftune_noise_alpha\",\n            \"model_init_kwargs\",\n            \"dataset_kwargs\",\n            \"eval_packing\",\n            \"num_of_sequences\",\n            \"chars_per_token\",\n        ],\n        custom_message=\"Deprecated positional argument(s) used in SFTTrainer, please use the SFTConfig to set these arguments instead.\",\n    )\n    def __init__(\n        self,\n        model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        args: Optional[SFTConfig] = None,\n        data_collator: Optional[DataCollator] = None,  # type: ignore\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        peft_config: Optional[\"PeftConfig\"] = None,\n        dataset_text_field: Optional[str] = None,\n        packing: Optional[bool] = False,\n        formatting_func: Optional[Callable] = None,\n        max_seq_length: Optional[int] = None,\n        infinite: Optional[bool] = None,\n        num_of_sequences: Optional[int] = None,\n        chars_per_token: Optional[float] = None,\n        dataset_num_proc: Optional[int] = None,\n        dataset_batch_size: Optional[int] = None,\n        neftune_noise_alpha: Optional[float] = None,\n        model_init_kwargs: Optional[Dict] = None,\n        dataset_kwargs: Optional[Dict] = None,\n        eval_packing: Optional[bool] = None,\n    ):\n        if args is None:\n            output_dir = \"tmp_trainer\"\n            warnings.warn(f\"No `SFTConfig` passed, using `output_dir={output_dir}`.\")\n            args = SFTConfig(output_dir=output_dir)\n        elif args is not None and args.__class__.__name__ == \"TrainingArguments\":\n            args_as_dict = args.to_dict()\n            # Manually copy token values as TrainingArguments.to_dict() redacts them\n            args_as_dict.update({k: getattr(args, k) for k in args_as_dict.keys() if k.endswith(\"_token\")})\n            args = SFTConfig(**args_as_dict)\n\n        if neftune_noise_alpha is not None:\n            warnings.warn(\n                \"You passed a `neftune_noise_alpha` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.neftune_noise_alpha = neftune_noise_alpha\n\n        if model_init_kwargs is not None:\n            warnings.warn(\n                \"You passed `model_init_kwargs` to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.model_init_kwargs = model_init_kwargs\n        if getattr(args, \"model_init_kwargs\", None) is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\"You passed model_init_kwargs to the SFTConfig, but your model is already instantiated.\")\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the SFTConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if infinite is not None:\n            warnings.warn(\n                \"The `infinite` argument is deprecated and will be removed in a future version of TRL. Use `TrainingArguments.max_steps` or `TrainingArguments.num_train_epochs` instead to control training length.\"\n            )\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the SFTTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            if args.use_liger:\n                model = AutoLigerKernelForCausalLM.from_pretrained(model, **model_init_kwargs)\n            else:\n                model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        if packing:\n            warnings.warn(\n                \"You passed a `packing` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.packing = packing\n        if eval_packing is not None:\n            warnings.warn(\n                \"You passed a `eval_packing` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.eval_packing = eval_packing\n\n        if args.packing and data_collator is not None and isinstance(data_collator, DataCollatorForCompletionOnlyLM):\n            raise ValueError(\n                \"You passed a `DataCollatorForCompletionOnlyLM` to the SFTTrainer. This is not compatible with the `packing` argument.\"\n            )\n\n        if is_peft_available() and peft_config is not None:\n            if not isinstance(peft_config, PeftConfig):\n                raise ValueError(\n                    \"If you want to use the PeftModel, you need to pass a PeftConfig object to the SFTTrainer.\"\n                    f\" and you passed a {type(peft_config)}.\"\n                )\n\n            if not isinstance(model, PeftModel):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n                gradient_checkpointing_kwargs = getattr(args, \"gradient_checkpointing_kwargs\", None) or {}\n                is_sharded_qlora = False\n                # Below is to support QLoRA + FSDP / DS-Zero3 - one should never call\n                # peft_module_casting_to_bf16 or prepare_model_for_kbit_training when doing\n                # QLoRA + FSDP / DS-Zero3\n                if getattr(model, \"is_loaded_in_4bit\", False):\n                    for _, param in model.named_parameters():\n                        if param.__class__.__name__ == \"Params4bit\":\n                            is_sharded_qlora = param.data.device.type == \"cpu\"\n                            break\n                if getattr(model, \"is_loaded_in_8bit\", False) or (\n                    getattr(model, \"is_loaded_in_4bit\", False) and not is_sharded_qlora\n                ):\n                    prepare_model_kwargs = {\n                        \"use_gradient_checkpointing\": getattr(args, \"gradient_checkpointing\", False)\n                    }\n\n                    if _support_gc_kwargs:\n                        prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = gradient_checkpointing_kwargs\n\n                    model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n\n                    if args is not None:\n                        args = dataclasses.replace(args, gradient_checkpointing=False)\n                elif getattr(args, \"gradient_checkpointing\", False) and (\n                    \"use_reentrant\" not in gradient_checkpointing_kwargs\n                    or gradient_checkpointing_kwargs[\"use_reentrant\"]\n                ):\n                    # For backward compatibility with older versions of transformers\n                    if hasattr(model, \"enable_input_require_grads\"):\n                        model.enable_input_require_grads()\n                    else:\n\n                        def make_inputs_require_grad(module, input, output):\n                            output.requires_grad_(True)\n\n                        model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n                if (\n                    \"autocast_adapter_dtype\" in list(inspect.signature(get_peft_model).parameters)\n                    and getattr(model, \"is_loaded_in_4bit\", False)\n                    and is_sharded_qlora\n                ):\n                    model = get_peft_model(model, peft_config, autocast_adapter_dtype=False)\n                else:\n                    model = get_peft_model(model, peft_config)\n                if (\n                    args is not None\n                    and args.bf16\n                    and getattr(model, \"is_loaded_in_4bit\", False)\n                    and not is_sharded_qlora\n                ):\n                    peft_module_casting_to_bf16(model)\n\n        if tokenizer is None:\n            tokenizer = AutoTokenizer.from_pretrained(model.config._name_or_path)\n            if getattr(tokenizer, \"pad_token\", None) is None:\n                tokenizer.pad_token = tokenizer.eos_token\n\n        if max_seq_length is not None:\n            warnings.warn(\n                \"You passed a `max_seq_length` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.max_seq_length = max_seq_length\n\n        if args.max_seq_length is None:\n            # to overcome some issues with broken tokenizers\n            args.max_seq_length = min(tokenizer.model_max_length, 1024)\n\n            warnings.warn(\n                f\"You didn't pass a `max_seq_length` argument to the SFTTrainer, this will default to {args.max_seq_length}\"\n            )\n\n        if dataset_num_proc is not None:\n            warnings.warn(\n                \"You passed a `dataset_num_proc` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.dataset_num_proc = dataset_num_proc\n        self.dataset_num_proc = args.dataset_num_proc\n\n        if dataset_batch_size is not None:\n            warnings.warn(\n                \"You passed a `dataset_batch_size` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.dataset_batch_size = dataset_batch_size\n        self.dataset_batch_size = args.dataset_batch_size\n\n        if dataset_text_field is not None:\n            warnings.warn(\n                \"You passed a `dataset_text_field` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.dataset_text_field = dataset_text_field\n\n        if dataset_kwargs is not None:\n            warnings.warn(\n                \"You passed a `dataset_kwargs` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.dataset_kwargs = dataset_kwargs\n        if args.dataset_kwargs is None:\n            args.dataset_kwargs = {}\n\n        if formatting_func is None and args.dataset_text_field is None:\n            # check if dataset has ChatML format or instruction format and is supported\n            # if not stays #None\n            formatting_func = get_formatting_func_from_dataset(train_dataset, tokenizer)\n            # if a template is detected, we don't need to add special tokens again\n            if formatting_func is not None:\n                args.dataset_kwargs[\"add_special_tokens\"] = False\n\n        if not args.packing:\n            if data_collator is None:\n                data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n\n        if num_of_sequences is not None:\n            warnings.warn(\n                \"You passed a `num_of_sequences` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.num_of_sequences = num_of_sequences\n\n        if chars_per_token is not None:\n            warnings.warn(\n                \"You passed a `chars_per_token` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\"\n            )\n            args.chars_per_token = chars_per_token\n\n        # Pre-process the datasets only once per node. The remaining processes will use the cache.\n        with PartialState().local_main_process_first():\n            if train_dataset is not None:\n                train_dataset = self._prepare_dataset(\n                    train_dataset,\n                    tokenizer,\n                    args.packing,\n                    args.dataset_text_field,\n                    args.max_seq_length,\n                    formatting_func,\n                    args.num_of_sequences,\n                    args.chars_per_token,\n                    remove_unused_columns=args.remove_unused_columns if args is not None else True,\n                    **args.dataset_kwargs,\n                )\n            if eval_dataset is not None:\n                _multiple = isinstance(eval_dataset, dict)\n                _eval_datasets = eval_dataset if _multiple else {\"singleton\": eval_dataset}\n\n                eval_packing = args.packing if args.eval_packing is None else args.eval_packing\n\n                for _eval_dataset_name, _eval_dataset in _eval_datasets.items():\n                    _eval_datasets[_eval_dataset_name] = self._prepare_dataset(\n                        _eval_dataset,\n                        tokenizer,\n                        eval_packing,\n                        args.dataset_text_field,\n                        args.max_seq_length,\n                        formatting_func,\n                        args.num_of_sequences,\n                        args.chars_per_token,\n                        remove_unused_columns=args.remove_unused_columns if args is not None else True,\n                        **args.dataset_kwargs,\n                    )\n                if not _multiple:\n                    eval_dataset = _eval_datasets[\"singleton\"]\n\n        if tokenizer.padding_side is not None and tokenizer.padding_side != \"right\":\n            warnings.warn(\n                \"You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to \"\n                \"overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code.\"\n            )\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if self.train_dataset is not None:\n            if self.args.max_steps > 0 and args.packing:\n                warnings.warn(\n                    \"You passed `packing=True` to the SFTTrainer/SFTConfig, and you are training your model with `max_steps` strategy. The dataset will be iterated until the `max_steps` are reached.\"\n                )\n                self.train_dataset.infinite = True\n            elif self.args.max_steps == -1 and args.packing:\n                self.train_dataset.infinite = False\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"sft\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n    def _prepare_dataset(\n        self,\n        dataset,\n        tokenizer,\n        packing,\n        dataset_text_field,\n        max_seq_length,\n        formatting_func,\n        num_of_sequences,\n        chars_per_token,\n        remove_unused_columns=True,\n        append_concat_token=True,\n        add_special_tokens=True,\n        skip_prepare_dataset=False,\n    ):\n        if dataset is None:\n            raise ValueError(\"The dataset should not be None\")\n\n        if skip_prepare_dataset:\n            return dataset\n\n        # If the dataset is already preprocessed (tokenized), return as-is. Only works if dataset is\n        # a datasets.Dataset or datasets.IterableDataset -- not for torch Dataset\n        column_names = (\n            dataset.column_names if isinstance(dataset, (datasets.Dataset, datasets.IterableDataset)) else None\n        )\n        if column_names and \"input_ids\" in column_names:\n            if formatting_func is not None:\n                warnings.warn(\n                    \"You passed a dataset that is already processed (contains an `input_ids` field) together with a valid formatting function. Therefore `formatting_func` will be ignored.\"\n                )\n\n            return dataset\n\n        # check if torch dataset / dataloader and do nothing\n        # see https://github.com/huggingface/trl/pull/1468 for why datasets.IterableDataset needs a separate check\n        if isinstance(\n            dataset, (torch.utils.data.IterableDataset, torch.utils.data.Dataset, ConstantLengthDataset)\n        ) and not isinstance(dataset, datasets.IterableDataset):\n            return dataset\n\n        # If we aren't skipping data preparation, then a dataset_text_field or formatting_func must be provided.\n        if dataset_text_field is None and formatting_func is None:\n            raise ValueError(\n                \"You need to provide either `dataset_text_field` or `formatting_func` argument. Alternatively, you \"\n                \"can skip the dataset preparation by using `SFTConfig(dataset_kwargs={'skip_prepare_dataset': True})`.\"\n            )\n\n        if not packing:\n            return self._prepare_non_packed_dataloader(\n                tokenizer,\n                dataset,\n                dataset_text_field,\n                max_seq_length,\n                formatting_func,\n                add_special_tokens,\n                remove_unused_columns,\n            )\n\n        else:\n            return self._prepare_packed_dataloader(\n                tokenizer,\n                dataset,\n                dataset_text_field,\n                max_seq_length,\n                num_of_sequences,\n                chars_per_token,\n                formatting_func,\n                append_concat_token,\n                add_special_tokens,\n            )\n\n    def _prepare_non_packed_dataloader(\n        self,\n        tokenizer,\n        dataset,\n        dataset_text_field,\n        max_seq_length,\n        formatting_func=None,\n        add_special_tokens=True,\n        remove_unused_columns=True,\n    ):\n        use_formatting_func = formatting_func is not None and dataset_text_field is None\n\n        # Inspired from: https://huggingface.co/learn/nlp-course/chapter7/6?fw=pt\n        def tokenize(element):\n            outputs = tokenizer(\n                element[dataset_text_field] if not use_formatting_func else formatting_func(element),\n                add_special_tokens=add_special_tokens,\n                truncation=True,\n                padding=False,\n                max_length=max_seq_length,\n                return_overflowing_tokens=False,\n                return_length=False,\n            )\n\n            if use_formatting_func and not isinstance(formatting_func(element), list):\n                raise ValueError(\n                    \"The `formatting_func` should return a list of processed strings since it can lead to silent bugs.\"\n                )\n\n            return {\"input_ids\": outputs[\"input_ids\"], \"attention_mask\": outputs[\"attention_mask\"]}\n\n        signature_columns = [\"input_ids\", \"labels\", \"attention_mask\"]\n\n        if dataset.column_names is not None:  # None for IterableDataset\n            extra_columns = list(set(dataset.column_names) - set(signature_columns))\n        else:\n            extra_columns = []\n\n        if not remove_unused_columns and len(extra_columns) > 0:\n            warnings.warn(\n                \"You passed `remove_unused_columns=False` on a non-packed dataset. This might create some issues with the default collator and yield to errors. If you want to \"\n                f\"inspect dataset other columns (in this case {extra_columns}), you can subclass `DataCollatorForLanguageModeling` in case you used the default collator and create your own data collator in order to inspect the unused dataset columns.\"\n            )\n\n        map_kwargs = {\n            \"batched\": True,\n            \"remove_columns\": dataset.column_names if remove_unused_columns else None,\n            \"batch_size\": self.dataset_batch_size,\n        }\n        if isinstance(dataset, datasets.Dataset):\n            map_kwargs[\"num_proc\"] = self.dataset_num_proc  # this arg is not available for IterableDataset\n        tokenized_dataset = dataset.map(tokenize, **map_kwargs)\n\n        return tokenized_dataset\n\n    def _prepare_packed_dataloader(\n        self,\n        tokenizer,\n        dataset,\n        dataset_text_field,\n        max_seq_length,\n        num_of_sequences,\n        chars_per_token,\n        formatting_func=None,\n        append_concat_token=True,\n        add_special_tokens=True,\n    ):\n        if dataset_text_field is not None or formatting_func is not None:\n            if tokenizer is None:\n                raise ValueError(\"You need to pass a tokenizer when using `dataset_text_field` with `SFTTrainer`.\")\n\n            constant_length_iterator = ConstantLengthDataset(\n                tokenizer,\n                dataset,\n                dataset_text_field=dataset_text_field,\n                formatting_func=formatting_func,\n                seq_length=max_seq_length,\n                infinite=False,\n                num_of_sequences=num_of_sequences,\n                chars_per_token=chars_per_token,\n                eos_token_id=tokenizer.eos_token_id,\n                append_concat_token=append_concat_token,\n                add_special_tokens=add_special_tokens,\n            )\n\n            if isinstance(dataset, datasets.IterableDataset):\n                return constant_length_iterator\n\n            def data_generator(constant_length_iterator):\n                yield from constant_length_iterator\n\n            try:\n                packed_dataset = Dataset.from_generator(\n                    data_generator, gen_kwargs={\"constant_length_iterator\": constant_length_iterator}\n                )\n            except (DatasetGenerationError, SchemaInferenceError) as exc:\n                raise ValueError(\n                    \"Error occurred while packing the dataset. \"\n                    \"Make sure that your dataset has enough samples to at least yield one packed sequence.\"\n                ) from exc\n            return packed_dataset\n        else:\n            raise ValueError(\n                \"You need to pass a `dataset_text_field` or `formatting_func` argument to the SFTTrainer if you want to use the `ConstantLengthDataset`.\"\n            )\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom dataclasses import dataclass\n\nfrom ..trainer.utils import OnPolicyConfig\n\n\n@dataclass\nclass PPOv2Config(OnPolicyConfig):\n    r\"\"\"\n    Configuration class for the [`PPOv2Trainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[:-3]`):\n            Name of this experiment.\n        reward_model_path (`str`, *optional*, defaults to `\"EleutherAI/pythia-160m\"`):\n            Path to the reward model.\n        num_ppo_epochs (`int`, *optional*, defaults to `4`):\n            Number of epochs to train.\n        whiten_rewards (`bool`, *optional*, defaults to `False`):\n            Whether to whiten the rewards.\n        kl_coef (`float`, *optional*, defaults to `0.05`):\n            KL coefficient.\n        cliprange (`float`, *optional*, defaults to `0.2`):\n            Clip range.\n        vf_coef (`float`, *optional*, defaults to `0.1`):\n            Value function coefficient.\n        cliprange_value (`float`, *optional*, defaults to `0.2`):\n            Clip range for the value function.\n        gamma (`float`, *optional*, defaults to `1.0`):\n            Discount factor.\n        lam (`float`, *optional*, defaults to `0.95`):\n            Lambda value for GAE.\n    \"\"\"\n\n    exp_name: str = os.path.basename(__file__)[: -len(\".py\")]\n    reward_model_path: str = \"EleutherAI/pythia-160m\"\n    num_ppo_epochs: int = 4\n    whiten_rewards: bool = False\n    kl_coef: float = 0.05\n    cliprange: float = 0.2\n    vf_coef: float = 0.1\n    cliprange_value: float = 0.2\n    gamma: float = 1.0\n    lam: float = 0.95\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass SFTConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`SFTTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        dataset_text_field (`Optional[str]`, *optional*, defaults to `None`):\n            Name of the text field of the dataset. If provided, the trainer will automatically create a\n            [`ConstantLengthDataset`] based on `dataset_text_field`.\n        packing (`bool`, *optional*, defaults to `False`):\n            Controls whether the [`ConstantLengthDataset`] packs the sequences of the dataset.\n        max_seq_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum sequence length for the [`ConstantLengthDataset`] and for automatically creating the dataset. If\n            `None`, it uses the smaller value between `tokenizer.model_max_length` and `1024`.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset. Only used when `packing=False`.\n        dataset_batch_size (`Union[int, None]`, *optional*, defaults to `1000`):\n            Number of examples to tokenize per batch. If `dataset_batch_size <= 0` or `dataset_batch_size is None`,\n            tokenizes the full dataset as a single batch.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        dataset_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Dictionary of optional keyword arguments to pass when creating packed or non-packed datasets.\n        eval_packing (`Optional[bool]`, *optional*, defaults to `None`):\n            Whether to pack the eval dataset. If `None`, uses the same value as `packing`.\n        num_of_sequences (`int`, *optional*, defaults to `1024`):\n            Number of sequences to use for the [`ConstantLengthDataset`].\n        chars_per_token (`float`, *optional*, defaults to `3.6`):\n            Number of characters per token to use for the [`ConstantLengthDataset`]. See\n            [chars_token_ratio](https://github.com/huggingface/trl/blob/08f550674c553c36c51d1027613c29f14f3676a5/examples/stack_llama/scripts/supervised_finetuning.py#L53) for more details.\n        use_liger (`bool`, *optional*, defaults to `False`):\n            Monkey patch the model with Liger kernels to increase throughput and reduce memory usage.\n    \"\"\"\n\n    dataset_text_field: Optional[str] = None\n    packing: bool = False\n    max_seq_length: Optional[int] = None\n    dataset_num_proc: Optional[int] = None\n    dataset_batch_size: int = 1000\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    dataset_kwargs: Optional[Dict[str, Any]] = None\n    eval_packing: Optional[bool] = None\n    num_of_sequences: int = 1024\n    chars_per_token: float = 3.6\n    use_liger: bool = False\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Optional\n\nfrom .sft_config import SFTConfig\n\n\n@dataclass\nclass GKDConfig(SFTConfig):\n    \"\"\"\n    Configuration class for GKDTrainer.\n\n    Args:\n        temperature (`float`, *optional*, defaults to `0.9`):\n            Temperature for sampling. The higher the temperature, the more random the completions.\n        lmbda (`float`, *optional*, defaults to `0.5`):\n            Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy\n            student-generated outputs).\n        beta (`float`, *optional*, defaults to `0.5`):\n            Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When\n            beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence.\n        max_new_tokens (`int`, *optional*, defaults to `128`):\n            Maximum number of tokens to generate per completion.\n        teacher_model_name_or_path (`Optional[str]`, *optional*, defaults to `None`):\n            Model name or path of the teacher model. If `None`, the teacher model will be the same as the model\n            being trained.\n        teacher_model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model\n            from a string.\n        disable_dropout (`bool`, *optional*, defaults to `True`):\n            Whether or not to disable dropouts in `model`.\n    \"\"\"\n\n    temperature: float = 0.9\n    lmbda: float = 0.5\n    beta: float = 0.5\n    max_new_tokens: int = 128\n    teacher_model_name_or_path: Optional[str] = None\n    teacher_model_init_kwargs: Optional[Dict[str, Any]] = None\n    disable_dropout: bool = True\n\n    def __post_init__(self):\n        super().__post_init__()\n        # check lmbda and beta are in the range [0, 1]\n        if self.lmbda < 0.0 or self.lmbda > 1.0:\n            raise ValueError(\"lmbda must be in the range [0.0, 1.0].\")\n        if self.beta < 0.0 or self.beta > 1.0:\n            raise ValueError(\"beta must be in the range [0.0, 1.0].\")\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport concurrent.futures\nimport logging\nimport random\nfrom abc import ABC, abstractmethod\nfrom typing import List, Optional, Union\n\nimport numpy as np\nfrom accelerate import Accelerator\nfrom huggingface_hub import InferenceClient\nfrom transformers.utils import is_openai_available\n\nfrom ..import_utils import is_llmblender_available\n\n\nif is_llmblender_available():\n    import llm_blender\n\nif is_openai_available():\n    from openai import OpenAI\n\n\nDEFAULT_PAIRWISE_SYSTEM_PROMPT = '''I require a leaderboard for various large language models. I'll provide you with prompts given to these models and their corresponding outputs. Your task is to assess these responses, and select the model that produces the best output from a human perspective.\n\n## Instruction\n\n{{\n    \"instruction\": \"\"\"{prompt}\"\"\",\n}}\n\n## Model Outputs\n\nHere are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier.\n\n{{\n    {{\n        \"model_identifier\": \"0\",\n        \"output\": \"\"\"{response0}\"\"\"\n    }},\n    {{\n        \"model_identifier\": \"1\",\n        \"output\": \"\"\"{response1}\"\"\"\n    }}\n}}\n\n## Task\n\nEvaluate the models on the basis of the quality and relevance of their results, and select the model that generated the best result. Reply with the identifier of the best model. Our evaluation will only take into account the first character of your answer, so make sure it contains only one of the identifiers and nothing else (no quotation marks, no spaces, no new lines, ...).\n'''\n\n\nclass BaseJudge(ABC):\n    \"\"\"\n    Base class for judges. The subclasses of this class should implement the `judge` method.\n    \"\"\"\n\n    @abstractmethod\n    def judge(self, prompts: List[str], completions: List[str], shuffle_order: bool = True) -> List:\n        raise NotImplementedError(\"Judge subclasses must implement the `judge` method.\")\n\n\nclass BaseRankJudge(ABC):\n    \"\"\"\n    Base class for LLM ranking judges.\n\n    Example:\n    ```python\n    class MyRankJudge(BaseRankJudge):\n        def judge(self, prompts, completions, shuffle_order=True):\n            return ...  # Your ranking logic here\n\n    judge = MyRankJudge()\n    judge.judge(\n        prompts=[\"The capital of France is\", \"The capital of Germany is\"],\n        completions=[[\" Paris\", \" Marseille\", \"Lyon\"], [\" Munich\", \" Berlin\"]]\n    )  # [[0, 1, 2], [1, 0]]\n    ```\n    \"\"\"\n\n    @abstractmethod\n    def judge(self, prompts: List[str], completions: List[List[str]], shuffle_order: bool = True) -> List[List[int]]:\n        \"\"\"\n        Judge the completion for the given prompts and return the ranks of each completion.\n\n        Args:\n            prompts (`List[str]`): List of prompts.\n            completions (`List[List[str]]`): List of completions list, where each element is a list of completions for the corresponding prompt.\n            shuffle_order (`bool`): Whether to shuffle the order of the completions to avoid positional bias.\n\n        Returns:\n            List of lists of idxs, where each list contains the ranks of the completions for the corresponding prompt.\n            E.g., [1, 2, 0] means that the second completion (idx=1) is the best, followed by the third, and then the first.\n        \"\"\"\n        raise NotImplementedError(\"Judge subclasses must implement the `judge` method.\")\n\n\nclass BasePairwiseJudge(BaseJudge):\n    \"\"\"\n    Base class for pairwise judges.\n    \"\"\"\n\n    @abstractmethod\n    def judge(self, prompts: List[str], completions: List[List[str]], shuffle_order: bool = True) -> List[int]:\n        \"\"\"\n        Judge the completion pairs for the given prompts.\n\n        Args:\n            prompts (`List[str]`): List of prompts.\n            completions (`List[List[str]]`): List of completions pairs, where each element is a pair of completions for the corresponding prompt.\n            shuffle_order (`bool`): Whether to shuffle the order of the completions to avoid positional bias.\n\n        Returns:\n            List of idxs, where each idx is the rank of the best completion for the corresponding prompt.\n            E.g., 1 means that the second completion (idx=1) is the best.\n\n        Note:\n            If the judge returns -1 for any prompt, it indicates that the inner process used to compute the preference has failed.\n            For instance, this could occur if the underlying language model returned an invalid answer.\n            In such cases, the caller should handle these invalid indices appropriately, possibly by implementing fallback logic or error handling.\n        \"\"\"\n        raise NotImplementedError(\"Judge subclasses must implement the `judge` method.\")\n\n\nclass RandomRankJudge(BaseRankJudge):\n    \"\"\"\n    Random rank, for testing purposes.\n    \"\"\"\n\n    def judge(self, prompts, completions, shuffle_order=True):\n        num_completions = [len(completions[i]) for i in range(len(prompts))]\n        return [random.sample(range(n), n) for n in num_completions]\n\n\nclass RandomPairwiseJudge(BasePairwiseJudge):\n    \"\"\"\n    Random pairwise judge, for testing purposes.\n    \"\"\"\n\n    def judge(self, prompts, completions, shuffle_order=True):\n        return [random.randint(0, len(completion) - 1) for completion in completions]\n\n\nclass PairRMJudge(BasePairwiseJudge):\n    \"\"\"\n    LLM judge based on the PairRM model from AllenAI.\n\n    See: https://huggingface.co/llm-blender/PairRM\n    \"\"\"\n\n    def __init__(self):\n        if not is_llmblender_available():\n            raise ValueError(\"llm-blender is not installed. Please install it with 'pip install llm-blender'.\")\n        self.blender = llm_blender.Blender()\n        self.blender.loadranker(\"llm-blender/PairRM\", device=Accelerator().device)\n\n    def judge(self, prompts: List[str], completions: List[List[str]], shuffle_order: bool = True) -> List[int]:\n        # Shuffle the order of the completions to avoid positional bias\n        if shuffle_order:\n            flip_mask = np.random.choice([True, False], size=len(prompts))\n            completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)]\n\n        # Rank the completions\n        ranks = self.blender.rank(prompts, completions)\n        ranks -= 1  # PairRM is 1-indexed, so we subtract 1 to make it 0-indexed\n\n        # Flip back the ranks to the original order if needed\n        if shuffle_order:\n            ranks[flip_mask] = ranks[flip_mask][:, ::-1]\n\n        # Return the ranks\n        return ranks[:, 0].tolist()\n\n\nclass HfPairwiseJudge(BasePairwiseJudge):\n    \"\"\"\n    Pairwise judge based on the Hugging Face API with chat completion.\n\n    This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt.\n\n    Args:\n        model (`str`, *optional*): The model to use for the judge. Defaults to \"meta-llama/Meta-Llama-3-70B-Instruct\".\n        token (`str`, *optional*): The Hugging Face API token to use for the InferenceClient.\n        system_prompt (`str`, *optional*): The system prompt to be used for the judge. If not provided, a default prompt is used.\n            Note that the system prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`.\n            Also, the inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token response.\n    \"\"\"\n\n    def __init__(\n        self,\n        model=\"meta-llama/Meta-Llama-3-70B-Instruct\",\n        token: Optional[str] = None,\n        system_prompt: Optional[str] = None,\n    ):\n        self.client = InferenceClient(model=model, token=token)\n        self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT\n\n    def judge(self, prompts: List[str], completions: List[List[str]], shuffle_order: bool = True) -> List[int]:\n        # Shuffle the order of the completions to avoid positional bias\n        if shuffle_order:\n            flip_mask = np.random.choice([True, False], size=len(prompts))\n            completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)]\n\n        # Define a function to get the rank for a single prompt, will be called concurrently\n        def get_rank(prompt, candidates):\n            content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1])\n            completion = self.client.chat_completion(messages=[{\"role\": \"user\", \"content\": content}], max_tokens=1)\n            response = completion.choices[0].message.content\n            if response in [\"0\", \"1\"]:\n                return int(response)\n            else:\n                logging.debug(f\"Invalid response from the judge model: '{response}'. Returning -1.\")\n                return -1\n\n        # Call the completions concurrently\n        with concurrent.futures.ThreadPoolExecutor() as executor:\n            ranks = list(executor.map(get_rank, prompts, completions))\n\n        # Flip back the ranks to the original order if needed\n        if shuffle_order:\n            ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)]\n\n        # Return the ranks\n        return ranks\n\n\nclass OpenAIPairwiseJudge(BasePairwiseJudge):\n    \"\"\"\n    Judge based on the OpenAI API.\n\n    This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt.\n\n    Args:\n        model (`str`, *optional*): The model to use for the judge. Defaults to `\"gpt-4-turbo-preview\"`.\n        system_prompt (`str`, *optional*): The system prompt to be used for the judge. If not provided, a default prompt is used.\n            Note that the system prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`.\n            Also, the inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token response.\n        max_requests (`int`, *optional*): The maximum number of requests to make to the OpenAI API. Defaults to 1000. If set to `None`, there is no limit.\n    \"\"\"\n\n    def __init__(\n        self, model=\"gpt-4-turbo-preview\", system_prompt: Optional[str] = None, max_requests: Union[int, None] = 1_000\n    ):\n        if not is_openai_available():\n            raise ValueError(\"OpenAI client is not installed. Please install it with 'pip install openai'.\")\n        self.client = OpenAI()\n        self.model = model\n        self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT\n        self.max_requests = max_requests\n        self.num_requests = 0\n        self._warned = False\n\n    def judge(self, prompts: List[str], completions: List[List[str]], shuffle_order: bool = True) -> List[int]:\n        # Check if the limit of requests is reached, if so, use random choice instead\n        if self.max_requests is not None and self.num_requests >= self.max_requests:\n            if not self._warned:  # Print the warning only once\n                logging.warning(\n                    f\"Reached the maximum number of requests ({self.max_requests}). From now on, returning -1 instead. \"\n                    \" To increase the limit, set `max_requests` to a higher value, or to `None` for no limit.\"\n                )\n                self._warned = True\n            return [-1] * len(prompts)\n\n        # Shuffle the order of the completions to avoid positional bias\n        if shuffle_order:\n            flip_mask = np.random.choice([True, False], size=len(prompts))\n            completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)]\n\n        # Define a function to get the rank for a single prompt, will be called concurrently\n        def get_rank(prompt, candidates):\n            content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1])\n            messages = [{\"role\": \"user\", \"content\": content}]\n            completion = self.client.chat.completions.create(model=self.model, messages=messages, max_tokens=1)\n            response = completion.choices[0].message.content\n            if response in [\"0\", \"1\"]:\n                return int(response)\n            else:\n                logging.debug(f\"Invalid response from the judge model: '{response}'. Returning -1.\")\n                return -1\n\n        # Call the completions concurrently\n        with concurrent.futures.ThreadPoolExecutor() as executor:\n            ranks = list(executor.map(get_rank, prompts, completions))\n\n        # Flip back the ranks to the original order if needed\n        if shuffle_order:\n            ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)]\n\n        # Update the number of requests\n        self.num_requests += len(prompts)\n\n        # Return the ranks\n        return ranks\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import List, Union\n\nfrom trl.trainer.online_dpo_config import OnlineDPOConfig\n\n\n@dataclass\nclass NashMDConfig(OnlineDPOConfig):\n    r\"\"\"\n    Configuration class for the [`NashMDTrainer`].\n\n    Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following:\n\n    Parameters:\n        mixture_coef (`float` or `list[float]`, *optional*, defaults to `0.5`):\n            Logit mixture coefficient for the model and reference model. If a list of floats is provided then the\n            mixture coefficient is selected for each new epoch and the last coefficient is used for the rest of the\n            epochs.\n    \"\"\"\n\n    mixture_coef: Union[float, List[float]] = 0.5\n\n\n# DPO Authors: Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn 2023\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import contextmanager, nullcontext\nfrom copy import deepcopy\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n\nimport torch\nimport torch.amp as amp\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import PartialState\nfrom accelerate.utils import is_deepspeed_available, tqdm\nfrom datasets import Dataset\nfrom huggingface_hub.utils._deprecation import _deprecate_arguments\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    AutoModelForCausalLM,\n    DataCollator,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    is_wandb_available,\n)\nfrom transformers.models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalLoopOutput\nfrom transformers.utils import is_peft_available\n\nfrom ..models import PreTrainedModelWrapper, create_reference_model\nfrom .callbacks import SyncRefModelCallback\nfrom .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    RunningMoments,\n    add_bos_token_if_needed,\n    add_eos_token_if_needed,\n    cap_exp,\n    disable_dropout_in_model,\n    pad_to_length,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\n\nif is_wandb_available():\n    import wandb\n\nif is_deepspeed_available():\n    import deepspeed\n\n\ndef _tokenize(\n    features: Dict[str, List],\n    tokenizer: PreTrainedTokenizerBase,\n    args: DPOConfig,\n    processor: Optional[Callable] = None,\n    model: Optional[PreTrainedModel] = None,\n) -> Dict[str, List]:\n    \"\"\"\n    Tokenizes and processes a batch of input features using the provided tokenizer and processor.\n    \"\"\"\n    batch = defaultdict(list)\n\n    if model is None:\n        prompt = features[\"prompt\"]\n        images = features.get(\"images\", [None] * len(features[\"prompt\"]))\n\n        prompt_tokens = _process_prompt(prompt, processor, tokenizer, images)\n        chosen_tokens = _process_answer(prompt, features[\"chosen\"], processor, tokenizer, images)\n        rejected_tokens = _process_answer(prompt, features[\"rejected\"], processor, tokenizer, images)\n\n        prompt_len_input_ids = _adjust_prompt_length(prompt_tokens, chosen_tokens, rejected_tokens)\n\n        prompt_tokens, chosen_tokens, rejected_tokens = _add_special_tokens(\n            tokenizer, prompt_len_input_ids, prompt_tokens, chosen_tokens, rejected_tokens\n        )\n\n        _truncate_tokens(chosen_tokens, rejected_tokens, prompt_tokens, args)\n\n        _build_sequence_tokens(batch, chosen_tokens, args, \"chosen\")\n        _build_sequence_tokens(batch, rejected_tokens, args, \"rejected\")\n\n        _append_prompt_tokens_to_batch(batch, prompt_tokens)\n\n    else:\n        _tokenize_encoder_decoder(\n            batch, tokenizer, features[\"prompt\"], features[\"chosen\"], features[\"rejected\"], args, model\n        )\n\n    return dict(batch)\n\n\ndef _process_prompt(\n    prompts: List[str], processor: Optional[Callable], tokenizer: PreTrainedTokenizerBase, images: List[Optional[Any]]\n) -> List[Dict[str, List[int]]]:\n    \"\"\"\n    Processes a list of prompts by tokenizing them, optionally using a processor for additional processing.\n    \"\"\"\n    if processor:\n        processor_kwargs = (\n            {\"add_special_tokens\": False} if \"add_special_tokens\" in inspect.signature(processor).parameters else {}\n        )\n        prompt_tokens = []\n        for prompt, image in zip(prompts, images):\n            tokens = processor(images=image, text=prompt, **processor_kwargs)\n            tokens = {k: v[0] for k, v in tokens.items()}\n            if not isinstance(tokens[\"input_ids\"], list):\n                tokens[\"input_ids\"] = tokens[\"input_ids\"].tolist()\n                tokens[\"attention_mask\"] = tokens[\"attention_mask\"].tolist()\n            prompt_tokens.append(tokens)\n    else:\n        prompt_tokens = [tokenizer(prompt, add_special_tokens=False) for prompt in prompts]\n    return [{f\"prompt_{k}\": v for k, v in tokens.items()} for tokens in prompt_tokens]\n\n\ndef _process_answer(\n    prompts: List[str],\n    answers: List[str],\n    processor: Optional[Callable],\n    tokenizer: PreTrainedTokenizerBase,\n    images: List[Optional[Any]],\n) -> List[Dict[str, Any]]:\n    return [\n        _build_tokenized_answer(prompt, answer, image, processor=processor, tokenizer=tokenizer)\n        for prompt, answer, image in zip(prompts, answers, images)\n    ]\n\n\ndef _adjust_prompt_length(\n    prompt_tokens: List[Dict[str, List[int]]],\n    chosen_tokens: List[Dict[str, List[int]]],\n    rejected_tokens: List[Dict[str, List[int]]],\n) -> List[int]:\n    prompt_len_input_ids = []\n    for p_tokens, c_tokens, r_tokens in zip(prompt_tokens, chosen_tokens, rejected_tokens):\n        c_len = len(c_tokens[\"prompt_input_ids\"])\n        r_len = len(r_tokens[\"prompt_input_ids\"])\n        min_len = min(c_len, r_len)\n\n        for k, v in p_tokens.items():\n            p_tokens[k] = v[:min_len]\n\n        num_diff_tokens = sum([a != b for a, b in zip(c_tokens[\"prompt_input_ids\"], r_tokens[\"prompt_input_ids\"])])\n        num_diff_len = abs(c_len - r_len)\n        if num_diff_tokens > 1 or num_diff_len > 1:\n            raise ValueError(\n                \"Chosen and rejected prompt_input_ids might only differ on the last token due to tokenizer merge ops.\"\n            )\n        prompt_len_input_ids.append(min_len)\n    return prompt_len_input_ids\n\n\ndef _add_special_tokens(\n    tokenizer: PreTrainedTokenizerBase,\n    prompt_len_input_ids: List[int],\n    prompt_tokens: List[Dict[str, List[int]]],\n    chosen_tokens: List[Dict[str, List[int]]],\n    rejected_tokens: List[Dict[str, List[int]]],\n) -> Tuple[List[Dict[str, List[int]]], List[Dict[str, List[int]]], List[Dict[str, List[int]]]]:\n    for i in range(len(prompt_tokens)):\n        prompt_tokens[i], chosen_tokens[i], rejected_tokens[i] = add_bos_token_if_needed(\n            tokenizer.bos_token_id,\n            prompt_len_input_ids[i],\n            prompt_tokens[i],\n            len(chosen_tokens[i][\"prompt_input_ids\"]),\n            chosen_tokens[i],\n            len(rejected_tokens[i][\"prompt_input_ids\"]),\n            rejected_tokens[i],\n        )\n\n        chosen_tokens[i], rejected_tokens[i] = add_eos_token_if_needed(\n            tokenizer.eos_token_id, chosen_tokens[i], rejected_tokens[i]\n        )\n    return prompt_tokens, chosen_tokens, rejected_tokens\n\n\ndef _truncate_tokens(\n    chosen_tokens: List[Dict[str, List[int]]],\n    rejected_tokens: List[Dict[str, List[int]]],\n    prompt_tokens: List[Dict[str, List[int]]],\n    args: DPOConfig,\n) -> None:\n    \"\"\"\n    Truncates the tokens in chosen, rejected, and prompt sequences to ensure they fit within the maximum length constraints.\n    \"\"\"\n    if args.truncation_mode not in [\"keep_start\", \"keep_end\"]:\n        raise ValueError(f\"Invalid truncation mode: {args.truncation_mode}\")\n\n    for c_tokens, r_tokens, p_tokens in zip(chosen_tokens, rejected_tokens, prompt_tokens):\n        longer_response_length = max(len(c_tokens[\"input_ids\"]), len(r_tokens[\"input_ids\"]))\n\n        # if combined sequence is too long, truncate the prompt\n        for answer_tokens in [c_tokens, r_tokens, p_tokens]:\n            if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > args.max_length:\n                if args.truncation_mode == \"keep_start\":\n                    for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                        answer_tokens[k] = answer_tokens[k][: args.max_prompt_length]\n                elif args.truncation_mode == \"keep_end\":\n                    for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                        answer_tokens[k] = answer_tokens[k][-args.max_prompt_length :]\n\n        # if that's still too long, truncate the response from the end\n        for answer_tokens in [c_tokens, r_tokens]:\n            if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > args.max_length:\n                for k in [\"input_ids\", \"attention_mask\"]:\n                    answer_tokens[k] = answer_tokens[k][: args.max_length - args.max_prompt_length]\n\n\ndef _build_sequence_tokens(\n    batch: Dict[str, List[int]], tokens: List[Dict[str, List[int]]], args: DPOConfig, prefix: str\n) -> None:\n    for token in tokens:\n        sequence_tokens = {f\"{prefix}_{k}\": token[f\"prompt_{k}\"] + token[k] for k in [\"input_ids\", \"attention_mask\"]}\n        sequence_tokens[f\"{prefix}_labels\"] = sequence_tokens[f\"{prefix}_input_ids\"][:]\n        sequence_tokens[f\"{prefix}_labels\"][: len(token[\"prompt_input_ids\"])] = [args.label_pad_token_id] * len(\n            token[\"prompt_input_ids\"]\n        )\n        for k, v in sequence_tokens.items():\n            batch[k].append(v)\n\n\ndef _append_prompt_tokens_to_batch(batch: Dict[str, List[int]], prompt_tokens: List[Dict[str, List[int]]]) -> None:\n    for p_tokens in prompt_tokens:\n        for k, v in p_tokens.items():\n            batch[k].append(v)\n\n\ndef _tokenize_encoder_decoder(\n    batch: Dict[str, List[int]],\n    tokenizer: PreTrainedTokenizerBase,\n    prompt: List[str],\n    chosen: List[str],\n    rejected: List[str],\n    args: DPOConfig,\n    model: Optional[PreTrainedModel],\n) -> None:\n    chosen_tokens = tokenizer(chosen, truncation=True, max_length=args.max_completion_length, add_special_tokens=True)\n    rejected_tokens = tokenizer(\n        rejected, truncation=True, max_length=args.max_completion_length, add_special_tokens=True\n    )\n    prompt_tokens = tokenizer(prompt, truncation=True, max_length=args.max_prompt_length, add_special_tokens=True)\n\n    batch[\"chosen_labels\"] = chosen_tokens[\"input_ids\"]\n    batch[\"rejected_labels\"] = rejected_tokens[\"input_ids\"]\n    batch[\"prompt_input_ids\"] = prompt_tokens[\"input_ids\"]\n    batch[\"prompt_attention_mask\"] = prompt_tokens[\"attention_mask\"]\n\n    if model is not None and hasattr(model, \"prepare_decoder_input_ids_from_labels\"):\n        # Ensure the sequences are of the same length\n        max_length = max(len(seq) for seq in batch[\"chosen_labels\"] + batch[\"rejected_labels\"])\n        batch[\"chosen_labels\"] = [\n            seq + [tokenizer.pad_token_id] * (max_length - len(seq)) for seq in batch[\"chosen_labels\"]\n        ]\n        batch[\"rejected_labels\"] = [\n            seq + [tokenizer.pad_token_id] * (max_length - len(seq)) for seq in batch[\"rejected_labels\"]\n        ]\n\n        batch[\"rejected_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n            labels=torch.tensor(batch[\"rejected_labels\"])\n        )\n        batch[\"chosen_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n            labels=torch.tensor(batch[\"chosen_labels\"])\n        )\n\n\ndef _build_tokenized_answer(\n    prompt: str,\n    answer: str,\n    images: Optional[List[Any]] = None,\n    processor: Optional[Callable] = None,\n    tokenizer: Optional[PreTrainedTokenizerBase] = None,\n) -> Dict[str, Any]:\n    \"\"\"\n    Build tokenized response, handling vision models and different tokenizers.\n    \"\"\"\n\n    def tokenize(text, images=None):\n        if processor:\n            processor_kwargs = (\n                {\"add_special_tokens\": False}\n                if \"add_special_tokens\" in inspect.signature(processor).parameters\n                else {}\n            )\n            tokenized = processor(images=images, text=text, **processor_kwargs)\n            tokenized = {k: v[0] for k, v in tokenized.items()}\n            if not isinstance(tokenized[\"input_ids\"], list):\n                tokenized[\"input_ids\"] = tokenized[\"input_ids\"].tolist()\n                tokenized[\"attention_mask\"] = tokenized[\"attention_mask\"].tolist()\n        else:\n            tokenized = tokenizer(text, add_special_tokens=False)\n        return tokenized\n\n    full_tokenized = tokenize(prompt + answer, images)\n    prompt_tokenized = tokenize(prompt, images)\n\n    prompt_input_ids = prompt_tokenized[\"input_ids\"]\n    answer_input_ids = full_tokenized[\"input_ids\"][len(prompt_input_ids) :]\n    answer_attention_mask = full_tokenized[\"attention_mask\"][len(prompt_input_ids) :]\n\n    if len(full_tokenized[\"input_ids\"]) != len(prompt_input_ids + answer_input_ids):\n        raise ValueError(\"Prompt input ids and answer input ids should have the same length.\")\n\n    # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens\n    # can be merged together when tokenizing prompt+answer. This could result\n    # on the last token from the prompt being different when tokenized on its own\n    # vs when done as prompt+answer.\n    response_token_ids_start_idx = len(prompt_input_ids)\n\n    # If tokenized prompt is different than both prompt+answer, then it means the\n    # last token has changed due to merging.\n    if prompt_input_ids != full_tokenized[\"input_ids\"][:response_token_ids_start_idx]:\n        response_token_ids_start_idx -= 1\n\n    prompt_input_ids = full_tokenized[\"input_ids\"][:response_token_ids_start_idx]\n    prompt_attention_mask = full_tokenized[\"attention_mask\"][:response_token_ids_start_idx]\n\n    if len(prompt_input_ids) != len(prompt_attention_mask):\n        raise ValueError(\"Prompt input ids and attention mask should have the same length.\")\n\n    return_dict = {\n        \"prompt_input_ids\": prompt_input_ids,\n        \"prompt_attention_mask\": prompt_attention_mask,\n        \"input_ids\": answer_input_ids,\n        \"attention_mask\": answer_attention_mask,\n    }\n    if \"pixel_values\" in full_tokenized:\n        return_dict[\"prompt_pixel_values\"] = full_tokenized[\"pixel_values\"]\n    if \"pixel_attention_mask\" in full_tokenized:\n        return_dict[\"prompt_pixel_attention_mask\"] = full_tokenized[\"pixel_attention_mask\"]\n\n    return return_dict\n\n\nclass DPOTrainer(Trainer):\n    r\"\"\"\n    Initialize DPOTrainer.\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForSequenceClassification`.\n        ref_model (`PreTrainedModelWrapper`):\n            Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no\n            reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.\n        args (`DPOConfig`):\n            The DPO config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Dict`, defaults to `None`):\n            The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"dpo\"]\n\n    @_deprecate_arguments(\n        version=\"1.0.0\",\n        deprecated_args=[\n            \"beta\",\n            \"label_smoothing\",\n            \"loss_type\",\n            \"label_pad_token_id\",\n            \"padding_value\",\n            \"truncation_mode\",\n            \"max_length\",\n            \"max_prompt_length\",\n            \"max_target_length\",\n            \"is_encoder_decoder\",\n            \"disable_dropout\",\n            \"generate_during_eval\",\n            \"precompute_ref_log_probs\",\n            \"dataset_num_proc\",\n            \"model_init_kwargs\",\n            \"ref_model_init_kwargs\",\n            \"model_adapter_name\",\n            \"ref_adapter_name\",\n            \"reference_free\",\n            \"force_use_ref_model\",\n        ],\n        custom_message=\"Deprecated positional argument(s) used in DPOTrainer, please use the DPOConfig to set these arguments instead.\",\n    )\n    def __init__(\n        self,\n        model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        beta: float = 0.1,\n        label_smoothing: float = 0,\n        loss_type: Optional[str] = None,\n        args: Optional[DPOConfig] = None,\n        data_collator: Optional[DataCollator] = None,\n        label_pad_token_id: int = -100,\n        padding_value: Optional[int] = None,\n        truncation_mode: str = \"keep_end\",\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        max_length: Optional[int] = None,\n        max_prompt_length: Optional[int] = None,\n        max_target_length: Optional[int] = None,\n        peft_config: Optional[Dict] = None,\n        is_encoder_decoder: Optional[bool] = None,\n        disable_dropout: bool = True,\n        generate_during_eval: bool = False,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n        precompute_ref_log_probs: bool = False,\n        dataset_num_proc: Optional[int] = None,\n        model_init_kwargs: Optional[Dict] = None,\n        ref_model_init_kwargs: Optional[Dict] = None,\n        model_adapter_name: Optional[str] = None,\n        ref_adapter_name: Optional[str] = None,\n        reference_free: bool = False,\n        force_use_ref_model: bool = False,\n    ):\n        if not isinstance(model, str) and ref_model is model:\n            raise ValueError(\n                \"`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the \"\n                \"same as `model`, you must mass a copy of it, or `None` if you use peft.\"\n            )\n\n        if model_init_kwargs is not None:\n            warnings.warn(\n                \"You passed `model_init_kwargs` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.model_init_kwargs = model_init_kwargs\n\n        if args.model_init_kwargs is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\n                \"You passed model_init_kwargs to the DPOTrainer/DPOConfig, but your model is already instantiated.\"\n            )\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if ref_model_init_kwargs is not None:\n            warnings.warn(\n                \"You passed `ref_model_init_kwargs` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.ref_model_init_kwargs = ref_model_init_kwargs\n\n        if args.ref_model_init_kwargs is None:\n            ref_model_init_kwargs = {}\n        elif not isinstance(ref_model, str):\n            raise ValueError(\n                \"You passed ref_model_init_kwargs to the DPOTrainer/DPOConfig, but your ref_model is already instantiated.\"\n            )\n        else:\n            ref_model_init_kwargs = args.ref_model_init_kwargs\n            torch_dtype = ref_model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the DPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                ref_model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the DPOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        if isinstance(ref_model, str):\n            warnings.warn(\n                \"You passed a ref model_id to the DPOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM`\"\n            )\n            ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs)\n\n        # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`\n        # has been called in order to properly call autocast if needed.\n        self._peft_has_been_casted_to_bf16 = False\n\n        if force_use_ref_model:\n            warnings.warn(\n                \"You passed `force_use_ref_model` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.force_use_ref_model = force_use_ref_model\n\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            # if model is a peft model and we have a peft_config, we merge and unload it first\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            if ref_model is not None and not args.force_use_ref_model:\n                raise ValueError(\n                    \"You passed both a ref_model and a peft_config. For training PEFT adapters with DPO there is no need to pass a reference\"\n                    \" model. Please pass `ref_model=None` in case you want to train PEFT adapters, or pass a ref_model with `force_use_ref_model=True` in DPOTrainer's init.\"\n                    \" if you want to use a different ref_model.\"\n                )\n\n            if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_loaded_in_4bit\", False):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n\n                prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                if _support_gc_kwargs:\n                    prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n            elif getattr(args, \"gradient_checkpointing\", False):\n                # For backward compatibility with older versions of transformers\n                if hasattr(model, \"enable_input_require_grads\"):\n                    model.enable_input_require_grads()\n                else:\n\n                    def make_inputs_require_grad(module, input, output):\n                        output.requires_grad_(True)\n\n                    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            # get peft model with the given config\n            model = get_peft_model(model, peft_config)\n            if args.bf16 and getattr(model, \"is_loaded_in_4bit\", False):\n                peft_module_casting_to_bf16(model)\n                # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager\n                self._peft_has_been_casted_to_bf16 = True\n\n        # For models that use gradient_checkpointing, we need to attach a hook that enables input\n        # to explicitly have `requires_grad=True`, otherwise training will either silently\n        # fail or completely fail.\n        elif getattr(args, \"gradient_checkpointing\", False):\n            # For backward compatibility with older versions of transformers\n            if hasattr(model, \"enable_input_require_grads\"):\n                model.enable_input_require_grads()\n            else:\n\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        if generate_during_eval:\n            warnings.warn(\n                \"You passed `generate_during_eval` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.generate_during_eval = generate_during_eval\n        if args.generate_during_eval and not is_wandb_available():\n            raise ValueError(\n                \"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install `wandb` to resolve.\"\n            )\n\n        if is_encoder_decoder is not None:\n            warnings.warn(\n                \"You passed `is_encoder_decoder` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.is_encoder_decoder = is_encoder_decoder\n        if model is not None:\n            self.is_encoder_decoder = model.config.is_encoder_decoder\n        elif args.is_encoder_decoder is None:\n            raise ValueError(\n                \"When no model is provided, you need to pass the parameter is_encoder_decoder to the DPOTrainer/DPOConfig.\"\n            )\n        else:\n            self.is_encoder_decoder = args.is_encoder_decoder\n\n        if model is not None:\n            self.is_vision_model = model.config.model_type in MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES.keys()\n        else:\n            warnings.warn(\n                \"No model provided, cannot determine if it is a vision model. Setting is_vision_model to False.\"\n            )\n            self.is_vision_model = False\n\n        if self.is_vision_model:\n            self.processor = tokenizer\n            self.tokenizer = tokenizer.tokenizer  # tokenizer is actually a processor at this point\n        else:\n            self.tokenizer = tokenizer\n\n        self.is_peft_model = is_peft_available() and isinstance(model, PeftModel)\n        if model_adapter_name is not None:\n            warnings.warn(\n                \"You passed `model_adapter_name` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.model_adapter_name = model_adapter_name\n        self.model_adapter_name = args.model_adapter_name\n\n        if ref_adapter_name is not None:\n            warnings.warn(\n                \"You passed `ref_adapter_name` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.ref_adapter_name = ref_adapter_name\n        self.ref_adapter_name = args.ref_adapter_name\n\n        if reference_free:\n            warnings.warn(\n                \"You passed `reference_free` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.reference_free = reference_free\n        self.reference_free = args.reference_free\n\n        if precompute_ref_log_probs:\n            warnings.warn(\n                \"You passed `precompute_ref_log_probs` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.precompute_ref_log_probs = precompute_ref_log_probs\n\n        if ref_model:\n            self.ref_model = ref_model\n        elif self.is_peft_model or args.precompute_ref_log_probs:\n            # The `model` with adapters turned off will be used as the reference model\n            self.ref_model = None\n        else:\n            self.ref_model = create_reference_model(model)\n\n        if tokenizer is None:\n            raise ValueError(\"tokenizer must be specified to tokenize a DPO dataset.\")\n\n        if max_length is not None:\n            warnings.warn(\n                \"You passed `max_length` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.max_length = max_length\n        if args.max_length is None:\n            warnings.warn(\n                \"`max_length` is not set in the DPOConfig's init\"\n                \" it will default to `512` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            args.max_length = 512\n\n        if max_prompt_length is not None:\n            warnings.warn(\n                \"You passed `max_prompt_length` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.max_prompt_length = max_prompt_length\n        if args.max_prompt_length is None:\n            warnings.warn(\n                \"`max_prompt_length` is not set in the DPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            args.max_prompt_length = 128\n\n        if max_target_length is not None:\n            warnings.warn(\n                \"You passed `max_target_length` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.max_completion_length = max_target_length\n        if args.max_completion_length is None and self.is_encoder_decoder:\n            warnings.warn(\n                \"When using an encoder decoder architecture, you should set `max_completion_length` in the DPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            args.max_completion_length = 128\n\n        if label_pad_token_id != -100:\n            warnings.warn(\n                \"You passed `label_pad_token_id` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.label_pad_token_id = label_pad_token_id\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(\n                pad_token_id=self.tokenizer.pad_token_id,\n                label_pad_token_id=args.label_pad_token_id,\n                is_encoder_decoder=self.is_encoder_decoder,\n            )\n\n            if args.remove_unused_columns:\n                args.remove_unused_columns = False\n                # warn users\n                warnings.warn(\n                    \"When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_dpo_data_collator = True\n        else:\n            self.use_dpo_data_collator = False\n\n        if not disable_dropout:\n            warnings.warn(\n                \"You passed `disable_dropout` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.disable_dropout = disable_dropout\n        if args.disable_dropout:\n            disable_dropout_in_model(model)\n            if self.ref_model is not None:\n                disable_dropout_in_model(self.ref_model)\n\n        self.max_length = args.max_length\n        self.generate_during_eval = args.generate_during_eval\n        self.label_pad_token_id = args.label_pad_token_id\n        if padding_value is not None:\n            warnings.warn(\n                \"You passed `padding_value` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.padding_value = padding_value\n        self.padding_value = args.padding_value if padding_value is not None else self.tokenizer.pad_token_id\n        self.max_prompt_length = args.max_prompt_length\n        if truncation_mode != \"keep_end\":\n            warnings.warn(\n                \"You passed `truncation_mode` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.truncation_mode = truncation_mode\n        self.truncation_mode = args.truncation_mode\n        self.max_completion_length = args.max_completion_length\n        self.precompute_ref_log_probs = args.precompute_ref_log_probs\n\n        # Since ref_logs are precomputed on the first call to get_train/eval_dataloader\n        # keep track of first called to avoid computation of future calls\n        self._precomputed_train_ref_log_probs = False\n        self._precomputed_eval_ref_log_probs = False\n\n        if loss_type is not None:\n            warnings.warn(\n                \"You passed `loss_type` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.loss_type = loss_type\n        if label_smoothing != 0:\n            warnings.warn(\n                \"You passed `label_smoothing` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.label_smoothing = label_smoothing\n        if (\n            args.loss_type in [\"hinge\", \"ipo\", \"bco_pair\", \"sppo_hard\", \"nca_pair\", \"apo_zero\", \"apo_down\"]\n            and args.label_smoothing > 0\n        ):\n            warnings.warn(\n                \"You are using a loss type that does not support label smoothing. Ignoring label_smoothing parameter.\"\n            )\n        if args.loss_type == \"kto_pair\":\n            raise ValueError(\"Support for kto_pair has been removed in DPOTrainer. Please use KTOTrainer.\")\n\n        if beta != 0.1:\n            warnings.warn(\n                \"You passed `beta` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.beta = beta\n        self.beta = args.beta\n        self.label_smoothing = args.label_smoothing\n        self.loss_type = args.loss_type\n        self.aux_loss_enabled = getattr(model.config, \"output_router_logits\", False)\n\n        self._stored_metrics = defaultdict(lambda: defaultdict(list))\n\n        self.f_divergence_type = args.f_divergence_type\n        self.f_divergence_params = {FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY: args.f_alpha_divergence_coef}\n\n        if dataset_num_proc is not None:\n            warnings.warn(\n                \"You passed `dataset_num_proc` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`.\"\n            )\n            args.dataset_num_proc = dataset_num_proc\n        self.dataset_num_proc = args.dataset_num_proc\n\n        # Compute that only on the main process for faster data processing.\n        # see: https://github.com/huggingface/trl/pull/1255\n        with PartialState().local_main_process_first():\n            # tokenize the dataset, lower writer batch size to avoid OOM (frequent in vision models)\n            fn_kwargs = {\n                \"tokenizer\": self.tokenizer,\n                \"args\": args,\n                \"processor\": self.processor if self.is_vision_model else None,\n                \"model\": model if self.is_encoder_decoder else None,\n            }\n            train_dataset = train_dataset.map(\n                _tokenize,\n                fn_kwargs=fn_kwargs,\n                batched=True,\n                num_proc=self.dataset_num_proc,\n                writer_batch_size=10,\n                desc=\"Tokenizing train dataset\",\n            )\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(\n                    _tokenize,\n                    fn_kwargs=fn_kwargs,\n                    batched=True,\n                    num_proc=self.dataset_num_proc,\n                    writer_batch_size=10,\n                    desc=\"Tokenizing eval dataset\",\n                )\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n        # Deepspeed Zero-3 does not support precompute_ref_log_probs\n        if self.is_deepspeed_enabled:\n            if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs:\n                raise ValueError(\n                    \"You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`.\"\n                )\n\n        if self.ref_model is None:\n            if not (self.is_peft_model or self.precompute_ref_log_probs):\n                raise ValueError(\n                    \"No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`\"\n                )\n            if args.sync_ref_model:\n                raise ValueError(\n                    \"You currently cannot use `ref_model=None` with TR-DPO method. Please provide `ref_model`.\"\n                )\n        else:\n            if self.is_deepspeed_enabled:\n                self.ref_model = self._prepare_deepspeed(self.ref_model)\n            else:\n                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)\n\n        if args.sync_ref_model:\n            if precompute_ref_log_probs:\n                raise ValueError(\n                    \"You cannot use `precompute_ref_log_probs=True` with TR-DPO method. Please set `precompute_ref_log_probs=False`.\"\n                )\n\n            self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator))\n        if self.loss_type == \"bco_pair\":\n            self.running = RunningMoments(self.accelerator)\n\n    def _prepare_deepspeed(self, model: PreTrainedModelWrapper):\n        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473\n        deepspeed_plugin = self.accelerator.state.deepspeed_plugin\n        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)\n\n        if model is not None:\n            if hasattr(model, \"config\"):\n                hidden_size = (\n                    max(model.config.hidden_sizes)\n                    if getattr(model.config, \"hidden_sizes\", None)\n                    else getattr(model.config, \"hidden_size\", None)\n                )\n                if hidden_size is not None and config_kwargs[\"zero_optimization\"][\"stage\"] == 3:\n                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`\n                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081\n                    config_kwargs.update(\n                        {\n                            \"zero_optimization.reduce_bucket_size\": hidden_size * hidden_size,\n                            \"zero_optimization.stage3_param_persistence_threshold\": 10 * hidden_size,\n                            \"zero_optimization.stage3_prefetch_bucket_size\": 0.9 * hidden_size * hidden_size,\n                        }\n                    )\n\n        # If ZeRO-3 is used, we shard both the active and reference model.\n        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)\n        if config_kwargs[\"zero_optimization\"][\"stage\"] != 3:\n            config_kwargs[\"zero_optimization\"][\"stage\"] = 0\n        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)\n        model.eval()\n        return model\n\n    def get_train_dataloader(self) -> DataLoader:\n        \"\"\"\n        Returns the training [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`.\n        \"\"\"\n\n        if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_train_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params))\n\n            reference_chosen_logps = []\n            reference_rejected_logps = []\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Train dataset reference log probs\"):\n                reference_chosen_logp, reference_rejected_logp = self.compute_reference_log_probs(padded_batch)\n                reference_chosen_logp, reference_rejected_logp = self.accelerator.gather_for_metrics(\n                    (reference_chosen_logp, reference_rejected_logp)\n                )\n                reference_chosen_logps.append(reference_chosen_logp.cpu())\n                reference_rejected_logps.append(reference_rejected_logp.cpu())\n\n                # Unnecessary cache clearing to avoid OOM\n                torch.cuda.empty_cache()\n                self.accelerator.free_memory()\n\n            all_reference_chosen_logps = torch.cat(reference_chosen_logps).float().numpy()\n            all_reference_rejected_logps = torch.cat(reference_rejected_logps).float().numpy()\n\n            self.train_dataset = self.train_dataset.add_column(\n                name=\"reference_chosen_logps\", column=all_reference_chosen_logps\n            )\n            self.train_dataset = self.train_dataset.add_column(\n                name=\"reference_rejected_logps\", column=all_reference_rejected_logps\n            )\n\n            self._precomputed_train_ref_log_probs = True\n\n        return super().get_train_dataloader()\n\n    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:\n        \"\"\"\n        Returns the evaluation [`~torch.utils.data.DataLoader`].\n\n        Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`.\n\n        Args:\n            eval_dataset (`torch.utils.data.Dataset`, *optional*):\n                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted\n                by the `model.forward()` method are automatically removed. It must implement `__len__`.\n        \"\"\"\n        if eval_dataset is None and self.eval_dataset is None:\n            raise ValueError(\"Trainer: evaluation requires an eval_dataset.\")\n        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset\n\n        if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs:\n            dataloader_params = {\n                \"batch_size\": self.args.per_device_eval_batch_size,\n                \"collate_fn\": self.data_collator,\n                \"num_workers\": self.args.dataloader_num_workers,\n                \"pin_memory\": self.args.dataloader_pin_memory,\n                \"shuffle\": False,\n            }\n\n            # prepare dataloader\n            data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))\n\n            reference_chosen_logps = []\n            reference_rejected_logps = []\n            for padded_batch in tqdm(iterable=data_loader, desc=\"Eval dataset reference log probs\"):\n                reference_chosen_logp, reference_rejected_logp = self.compute_reference_log_probs(padded_batch)\n                reference_chosen_logp, reference_rejected_logp = self.accelerator.gather_for_metrics(\n                    (reference_chosen_logp, reference_rejected_logp)\n                )\n                reference_chosen_logps.append(reference_chosen_logp.cpu())\n                reference_rejected_logps.append(reference_rejected_logp.cpu())\n\n            all_reference_chosen_logps = torch.cat(reference_chosen_logps).float().numpy()\n            all_reference_rejected_logps = torch.cat(reference_rejected_logps).float().numpy()\n\n            eval_dataset = eval_dataset.add_column(name=\"reference_chosen_logps\", column=all_reference_chosen_logps)\n            eval_dataset = eval_dataset.add_column(\n                name=\"reference_rejected_logps\", column=all_reference_rejected_logps\n            )\n\n            # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs\n            if self.eval_dataset is not None:\n                self.eval_dataset = eval_dataset\n            self._precomputed_eval_ref_log_probs = True\n\n        return super().get_eval_dataloader(eval_dataset=eval_dataset)\n\n    @contextmanager\n    def null_ref_context(self):\n        \"\"\"Context manager for handling null reference model (that is, peft adapter manipulation).\"\"\"\n        with self.accelerator.unwrap_model(\n            self.model\n        ).disable_adapter() if self.is_peft_model and not self.ref_adapter_name else nullcontext():\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.ref_adapter_name)\n            yield\n            if self.ref_adapter_name:\n                self.model.set_adapter(self.model_adapter_name or \"default\")\n\n    def compute_reference_log_probs(self, padded_batch: Dict) -> Dict:\n        \"\"\"Computes log probabilities of the reference model for a single padded batch of a DPO specific dataset.\"\"\"\n        compte_ref_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        # compute reference logps\n        with torch.no_grad(), compte_ref_context_manager:\n            if self.ref_model is None:\n                with self.null_ref_context():\n                    reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(\n                        self.model, padded_batch\n                    )[:2]\n            else:\n                reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(\n                    self.ref_model, padded_batch\n                )[:2]\n\n        return reference_chosen_logps, reference_rejected_logps\n\n    @staticmethod\n    def concatenated_inputs(\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        is_encoder_decoder: bool = False,\n        is_vision_model: bool = False,\n        label_pad_token_id: int = -100,\n        padding_value: int = 0,\n        device: Optional[torch.device] = None,\n    ) -> Dict[str, torch.LongTensor]:\n        \"\"\"Concatenate the chosen and rejected inputs into a single tensor.\n\n        Args:\n            batch: A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors of shape (batch_size, sequence_length).\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n            label_pad_token_id: The label pad token id.\n            padding_value: The padding value to use for the concatenated inputs_ids.\n            device: The device for the concatenated inputs.\n\n        Returns:\n            A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'.\n        \"\"\"\n        concatenated_batch = {}\n\n        if is_encoder_decoder:\n            max_length = max(batch[\"chosen_labels\"].shape[1], batch[\"rejected_labels\"].shape[1])\n        else:\n            max_length = max(batch[\"chosen_input_ids\"].shape[1], batch[\"rejected_input_ids\"].shape[1])\n\n        for k in batch:\n            if k.startswith(\"chosen\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"chosen\", \"concatenated\")\n                concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value)\n        for k in batch:\n            if k.startswith(\"rejected\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"rejected\", \"concatenated\")\n                concatenated_batch[concatenated_key] = torch.cat(\n                    (\n                        concatenated_batch[concatenated_key],\n                        pad_to_length(batch[k], max_length, pad_value=pad_value),\n                    ),\n                    dim=0,\n                ).to(device=device)\n\n        if is_encoder_decoder:\n            concatenated_batch[\"concatenated_input_ids\"] = batch[\"prompt_input_ids\"].repeat(2, 1).to(device=device)\n            concatenated_batch[\"concatenated_attention_mask\"] = (\n                batch[\"prompt_attention_mask\"].repeat(2, 1).to(device=device)\n            )\n            concatenated_batch[\"concatenated_decoder_input_ids\"] = torch.cat(\n                [batch[\"chosen_decoder_input_ids\"], batch[\"rejected_decoder_input_ids\"]], dim=0\n            ).to(device=device)\n\n        if is_vision_model:\n            concatenated_batch[\"pixel_values\"] = torch.cat(\n                [batch[\"prompt_pixel_values\"], batch[\"prompt_pixel_values\"]], dim=0\n            )\n            if \"prompt_pixel_attention_mask\" in batch:\n                concatenated_batch[\"pixel_attention_mask\"] = torch.cat(\n                    [batch[\"prompt_pixel_attention_mask\"], batch[\"prompt_pixel_attention_mask\"]], dim=0\n                )\n        return concatenated_batch\n\n    def dpo_loss(\n        self,\n        policy_chosen_logps: torch.FloatTensor,\n        policy_rejected_logps: torch.FloatTensor,\n        reference_chosen_logps: torch.FloatTensor,\n        reference_rejected_logps: torch.FloatTensor,\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Compute the DPO loss for a batch of policy and reference model log probabilities.\n\n        Args:\n            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)\n            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)\n            reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)\n            reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)\n\n        Returns:\n            A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).\n            The losses tensor contains the DPO loss for each example in the batch.\n            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.\n        \"\"\"\n        chosen_logratios = policy_chosen_logps.to(self.accelerator.device) - (\n            not self.reference_free\n        ) * reference_chosen_logps.to(self.accelerator.device)\n        rejected_logratios = policy_rejected_logps.to(self.accelerator.device) - (\n            not self.reference_free\n        ) * reference_rejected_logps.to(self.accelerator.device)\n\n        if self.f_divergence_type == FDivergenceType.ALPHA_DIVERGENCE.value:\n            # The alpha-divergence formula: (1 - u^-alpha) / alpha\n            # The divergence difference between the chosen and rejected sample is:\n            #     (1 - u[w]^-alpha) / alpha - (1 - u[l]^-alpha) / alpha\n            #        = (u[l]^-alpha - u[w]^-alpha) / alpha\n            # where u[w] and u[l] are the policy/reference probability ratios\n            # for the chosen and rejected samples, respectively.\n            alpha_coef = FDivergenceConstants.ALPHA_DIVERGENCE_COEF_DEFAULT\n            if self.f_divergence_params and FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY in self.f_divergence_params:\n                alpha_coef = float(self.f_divergence_params[FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY])\n            logits = (cap_exp(rejected_logratios * -alpha_coef) - cap_exp(chosen_logratios * -alpha_coef)) / alpha_coef\n        else:\n            pi_logratios = policy_chosen_logps - policy_rejected_logps\n            if self.reference_free:\n                ref_logratios = torch.tensor([0], dtype=pi_logratios.dtype, device=pi_logratios.device)\n            else:\n                ref_logratios = reference_chosen_logps - reference_rejected_logps\n\n            pi_logratios = pi_logratios.to(self.accelerator.device)\n            ref_logratios = ref_logratios.to(self.accelerator.device)\n            logits = pi_logratios - ref_logratios\n\n            if self.f_divergence_type == FDivergenceType.JS_DIVERGENCE.value:\n                # The js-divergence formula: log(2 * u / (1 + u))\n                # The divergence difference between the chosen and rejected sample is:\n                #     log(2 * u[w] / (1 + u[w])) - log(2 * u[l] / (1 + u[l]))\n                #       = log(u[w]) - log(u[l]) - (log(1 + u[w]) - log(1 + u[l]))\n                # where u[w] and u[l] are the policy/reference probability ratios\n                # for the chosen and rejected samples, respectively.\n                logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios)\n\n        # The beta is a temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5.\n        # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and\n        # calculates a conservative DPO loss.\n        if self.loss_type == \"sigmoid\":\n            losses = (\n                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)\n                - F.logsigmoid(-self.beta * logits) * self.label_smoothing\n            )\n        elif self.loss_type == \"robust\":\n            losses = (\n                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)\n                + F.logsigmoid(-self.beta * logits) * self.label_smoothing\n            ) / (1 - 2 * self.label_smoothing)\n        elif self.loss_type == \"exo_pair\":\n            # eqn (16) of the EXO paper: https://huggingface.co/papers/2402.00856\n            import math\n\n            if self.label_smoothing == 0:\n                self.label_smoothing = 1e-3\n            losses = (self.beta * logits).sigmoid() * (\n                F.logsigmoid(self.beta * logits) - math.log(1 - self.label_smoothing)\n            ) + (-self.beta * logits).sigmoid() * (F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing))\n        elif self.loss_type == \"hinge\":\n            losses = torch.relu(1 - self.beta * logits)\n        elif self.loss_type == \"ipo\":\n            # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper.\n            losses = (logits - 1 / (2 * self.beta)) ** 2\n        elif self.loss_type == \"bco_pair\":\n            chosen_logratios = policy_chosen_logps - reference_chosen_logps\n            rejected_logratios = policy_rejected_logps - reference_rejected_logps\n\n            chosen_rewards = self.beta * chosen_logratios\n            rejected_rewards = self.beta * rejected_logratios\n            rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach()\n            self.running.update(rewards)\n            delta = self.running.mean\n\n            losses = -F.logsigmoid((self.beta * chosen_logratios) - delta) - F.logsigmoid(\n                -(self.beta * rejected_logratios - delta)\n            )\n        elif self.loss_type == \"sppo_hard\":\n            # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, estimated using the PairRM score. The probability calculation is conducted outside of the trainer class. The version described here is the hard probability version, where P in Equation (4.7) of Algorithm 1 is set to 1 for the winner and 0 for the loser.\n            a = policy_chosen_logps - reference_chosen_logps\n            b = policy_rejected_logps - reference_rejected_logps\n\n            losses = (a - 0.5 / self.beta) ** 2 + (b + 0.5 / self.beta) ** 2\n        elif self.loss_type == \"nca_pair\":\n            chosen_rewards = (policy_chosen_logps - reference_chosen_logps) * self.beta\n            rejected_rewards = (policy_rejected_logps - reference_rejected_logps) * self.beta\n            losses = (\n                -F.logsigmoid(chosen_rewards)\n                - 0.5 * F.logsigmoid(-chosen_rewards)\n                - 0.5 * F.logsigmoid(-rejected_rewards)\n            )\n        elif self.loss_type == \"aot_pair\":\n            chosen_logratios = policy_chosen_logps - reference_chosen_logps\n            rejected_logratios = policy_rejected_logps - reference_rejected_logps\n\n            chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0)\n            rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0)\n\n            delta = chosen_logratios_sorted - rejected_logratios_sorted\n\n            losses = (\n                -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)\n                - F.logsigmoid(-self.beta * delta) * self.label_smoothing\n            )\n\n        elif self.loss_type == \"aot\":\n            pi_logratios = policy_chosen_logps - policy_rejected_logps\n            ref_logratios = reference_chosen_logps - reference_rejected_logps\n\n            pi_logratios_sorted, _ = torch.sort(pi_logratios, dim=0)\n            ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0)\n\n            delta = pi_logratios_sorted - ref_logratios_sorted\n\n            losses = (\n                -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)\n                - F.logsigmoid(-self.beta * delta) * self.label_smoothing\n            )\n\n        elif self.loss_type == \"apo_zero\":\n            # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266)\n            # Use this loss when you believe the chosen outputs are better than your model's default output\n\n            losses_chosen = 1 - F.sigmoid(self.beta * chosen_logratios)  # Increase chosen likelihood\n            losses_rejected = F.sigmoid(self.beta * rejected_logratios)  # Decrease rejected likelihood\n\n            losses = losses_chosen + losses_rejected\n\n        elif self.loss_type == \"apo_down\":\n            # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266)\n            # Use this loss when you believe the chosen outputs are worse than your model's default output\n\n            losses_chosen = F.sigmoid(self.beta * chosen_logratios)  # Decrease chosen likelihood\n            losses_rejected = 1 - F.sigmoid(\n                self.beta * (chosen_logratios - rejected_logratios)\n            )  # Decrease rejected likelihood more\n\n            losses = losses_chosen + losses_rejected\n\n        else:\n            raise ValueError(\n                f\"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', 'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_pair', 'apo_zero', 'apo_down']\"\n            )\n\n        chosen_rewards = (\n            self.beta\n            * (\n                policy_chosen_logps.to(self.accelerator.device) - reference_chosen_logps.to(self.accelerator.device)\n            ).detach()\n        )\n        rejected_rewards = (\n            self.beta\n            * (\n                policy_rejected_logps.to(self.accelerator.device)\n                - reference_rejected_logps.to(self.accelerator.device)\n            ).detach()\n        )\n\n        return losses, chosen_rewards, rejected_rewards\n\n    @staticmethod\n    def get_batch_logps(\n        logits: torch.FloatTensor,\n        labels: torch.LongTensor,\n        label_pad_token_id: int = -100,\n        is_encoder_decoder: bool = False,\n    ) -> Tuple[torch.FloatTensor, torch.LongTensor]:\n        \"\"\"Compute the log probabilities of the given labels under the given logits.\n\n        Args:\n            logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n            labels: Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are ignored. Shape: (batch_size, sequence_length)\n            label_pad_token_id: The label pad token id.\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n\n        Returns:\n            A Tuple of two tensor of shape ((batch_size,), (batch_size,)) containing the sum of log probabilities of the given labels under the given logits in the first tensor and the number of non-masked tokens in the second tensor.\n        \"\"\"\n        if logits.shape[:-1] != labels.shape:\n            raise ValueError(\n                f\"Logits (batch and sequence length dim) {logits.shape[:-1]} and labels must have the same shape {labels.shape}.\"\n            )\n\n        if not is_encoder_decoder:\n            labels = labels[:, 1:].clone()\n            logits = logits[:, :-1, :]\n        loss_mask = labels != label_pad_token_id\n\n        # dummy token; we'll ignore the losses on these tokens later\n        labels[labels == label_pad_token_id] = 0\n\n        per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n\n        return (per_token_logps * loss_mask).sum(-1), loss_mask.sum(-1)\n\n    def concatenated_forward(\n        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.\n\n        We do this to avoid doing two forward passes, because it's faster for FSDP.\n        \"\"\"\n        concatenated_batch = self.concatenated_inputs(\n            batch,\n            is_encoder_decoder=self.is_encoder_decoder,\n            is_vision_model=self.is_vision_model,\n            label_pad_token_id=self.label_pad_token_id,\n            padding_value=self.padding_value,\n            device=self.accelerator.device,\n        )\n        len_chosen = batch[\"chosen_labels\"].shape[0]\n\n        model_kwargs = {}\n\n        if self.is_encoder_decoder:\n            model_kwargs[\"labels\"] = concatenated_batch[\"concatenated_labels\"]\n            model_kwargs[\"decoder_input_ids\"] = concatenated_batch.get(\"concatenated_decoder_input_ids\")\n\n        if self.is_vision_model:\n            model_kwargs[\"pixel_values\"] = concatenated_batch[\"pixel_values\"]\n            if \"pixel_attention_mask\" in concatenated_batch:\n                model_kwargs[\"pixel_attention_mask\"] = concatenated_batch[\"pixel_attention_mask\"]\n\n        if self.aux_loss_enabled:\n            model_kwargs[\"output_router_logits\"] = True\n\n        outputs = model(\n            concatenated_batch[\"concatenated_input_ids\"],\n            attention_mask=concatenated_batch[\"concatenated_attention_mask\"],\n            use_cache=False,\n            **model_kwargs,\n        )\n        all_logits = outputs.logits\n\n        if all_logits.shape[:2] != concatenated_batch[\"concatenated_labels\"].shape[:2]:\n            # for llava, the model returns logits for the entire sequence, including the image tokens (placed before the text tokens)\n            seq_len = concatenated_batch[\"concatenated_labels\"].shape[1]\n            all_logits = all_logits[:, -seq_len:]\n\n        all_logps, size_completion = self.get_batch_logps(\n            all_logits,\n            concatenated_batch[\"concatenated_labels\"],\n            # average_log_prob=self.loss_type == \"ipo\",\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        def cross_entropy_loss(logits, labels):\n            if not self.is_encoder_decoder:\n                # Shift so that tokens < n predict n\n                logits = logits[..., :-1, :].contiguous()\n                labels = labels[..., 1:].contiguous()\n            # Flatten the tokens\n            loss_fct = nn.CrossEntropyLoss(ignore_index=self.label_pad_token_id)\n            logits = logits.view(-1, logits.shape[-1])\n            labels = labels.view(-1)\n            # Enable model parallelism\n            labels = labels.to(logits.device)\n            loss = loss_fct(logits, labels)\n            return loss\n\n        labels = concatenated_batch[\"concatenated_labels\"].clone()\n        nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen])\n\n        if self.loss_type == \"ipo\":\n            all_logps = all_logps / size_completion\n\n        chosen_logps = all_logps[:len_chosen]\n        rejected_logps = all_logps[len_chosen:]\n\n        chosen_logits = all_logits[:len_chosen]\n        rejected_logits = all_logits[len_chosen:]\n\n        if self.aux_loss_enabled:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss)\n\n        return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss)\n\n    def get_batch_loss_metrics(\n        self,\n        model,\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        train_eval: Literal[\"train\", \"eval\"] = \"train\",\n    ):\n        \"\"\"Compute the DPO loss and other metrics for the given batch of inputs for train or test.\"\"\"\n        metrics = {}\n\n        forward_output = self.concatenated_forward(model, batch)\n        (\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_chosen_logits,\n            policy_rejected_logits,\n            policy_nll_loss,\n        ) = forward_output[:5]\n        if self.aux_loss_enabled:\n            aux_loss = forward_output[5]\n\n        # if reference_chosen_logps and reference_rejected_logps in batch use them, otherwise use the reference model\n        if (\n            \"reference_chosen_logps\" in batch\n            and \"reference_rejected_logps\" in batch\n            and (self.precompute_ref_log_probs or self.args.rpo_alpha is not None)\n        ):\n            reference_chosen_logps = batch[\"reference_chosen_logps\"]\n            reference_rejected_logps = batch[\"reference_rejected_logps\"]\n        else:\n            with torch.no_grad():\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(\n                            self.model, batch\n                        )[:2]\n                else:\n                    reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(\n                        self.ref_model, batch\n                    )[:2]\n\n        losses, chosen_rewards, rejected_rewards = self.dpo_loss(\n            policy_chosen_logps,\n            policy_rejected_logps,\n            reference_chosen_logps,\n            reference_rejected_logps,\n        )\n        reward_accuracies = (chosen_rewards > rejected_rewards).float()\n\n        if self.args.rpo_alpha is not None:\n            # RPO loss from V3 of the paper:\n            losses = losses + policy_nll_loss * self.args.rpo_alpha\n\n        prefix = \"eval_\" if train_eval == \"eval\" else \"\"\n        metrics[f\"{prefix}rewards/chosen\"] = chosen_rewards.mean().cpu()\n        metrics[f\"{prefix}rewards/rejected\"] = rejected_rewards.mean().cpu()\n        metrics[f\"{prefix}rewards/accuracies\"] = reward_accuracies.mean().cpu()\n        metrics[f\"{prefix}rewards/margins\"] = (chosen_rewards - rejected_rewards).mean().cpu()\n        metrics[f\"{prefix}logps/rejected\"] = policy_rejected_logps.detach().mean().cpu()\n        metrics[f\"{prefix}logps/chosen\"] = policy_chosen_logps.detach().mean().cpu()\n        metrics[f\"{prefix}logits/rejected\"] = policy_rejected_logits.detach().mean().cpu()\n        metrics[f\"{prefix}logits/chosen\"] = policy_chosen_logits.detach().mean().cpu()\n        if self.args.rpo_alpha is not None:\n            metrics[f\"{prefix}nll_loss\"] = policy_nll_loss.detach().mean().cpu()\n\n        if self.aux_loss_enabled:\n            return losses.mean() + getattr(model.config, \"router_aux_loss_coef\", 0.0) * aux_loss, metrics\n\n        return losses.mean(), metrics\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n\n        compute_loss_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n        with compute_loss_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"train\")\n\n        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:\n        loss = loss.to(self.args.device)\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"train\")\n\n        if return_outputs:\n            return (loss, metrics)\n        return loss\n\n    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:\n        \"\"\"Generate samples from the model and reference model for the given batch of inputs.\"\"\"\n\n        # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with\n        # the torch cuda amp context manager as some hidden states are silently casted to full precision.\n        generate_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with generate_context_manager:\n            policy_output = model.generate(\n                input_ids=batch[\"prompt_input_ids\"],\n                attention_mask=batch[\"prompt_attention_mask\"],\n                max_length=self.max_length,\n                do_sample=True,\n                pad_token_id=self.tokenizer.pad_token_id,\n            )\n\n            # if reference_output in batch use that otherwise use the reference model\n            if \"reference_output\" in batch:\n                reference_output = batch[\"reference_output\"]\n            else:\n                if self.ref_model is None:\n                    with self.null_ref_context():\n                        reference_output = self.model.generate(\n                            input_ids=batch[\"prompt_input_ids\"],\n                            attention_mask=batch[\"prompt_attention_mask\"],\n                            max_length=self.max_length,\n                            do_sample=True,\n                            pad_token_id=self.tokenizer.pad_token_id,\n                        )\n                else:\n                    reference_output = self.ref_model.generate(\n                        input_ids=batch[\"prompt_input_ids\"],\n                        attention_mask=batch[\"prompt_attention_mask\"],\n                        max_length=self.max_length,\n                        do_sample=True,\n                        pad_token_id=self.tokenizer.pad_token_id,\n                    )\n\n        policy_output = pad_to_length(policy_output, self.max_length, self.tokenizer.pad_token_id)\n        policy_output_decoded = self.tokenizer.batch_decode(policy_output, skip_special_tokens=True)\n\n        reference_output = pad_to_length(reference_output, self.max_length, self.tokenizer.pad_token_id)\n        reference_output_decoded = self.tokenizer.batch_decode(reference_output, skip_special_tokens=True)\n\n        return policy_output_decoded, reference_output_decoded\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ):\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        if ignore_keys is None:\n            if hasattr(model, \"config\"):\n                ignore_keys = getattr(model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        prediction_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with torch.no_grad(), prediction_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"eval\")\n\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"eval\")\n\n        if prediction_loss_only:\n            return (loss.detach(), None, None)\n\n        # logits for the chosen and rejected samples from model\n        logits_dict = {\n            \"eval_logits/chosen\": metrics[\"eval_logits/chosen\"],\n            \"eval_logits/rejected\": metrics[\"eval_logits/rejected\"],\n        }\n        logits = tuple(v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys)\n        logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)\n        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)\n\n        return (loss.detach(), logits, labels)\n\n    def store_metrics(self, metrics: Dict[str, float], train_eval: Literal[\"train\", \"eval\"] = \"train\") -> None:\n        for key, value in metrics.items():\n            self._stored_metrics[train_eval][key].append(value)\n\n    def evaluation_loop(\n        self,\n        dataloader: DataLoader,\n        description: str,\n        prediction_loss_only: Optional[bool] = None,\n        ignore_keys: Optional[List[str]] = None,\n        metric_key_prefix: str = \"eval\",\n    ) -> EvalLoopOutput:\n        \"\"\"\n        Overriding built-in evaluation loop to store metrics for each batch.\n        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.\n\n        Works both with or without labels.\n        \"\"\"\n\n        # Sample and save to game log if requested (for one batch to save time)\n        if self.generate_during_eval:\n            # Generate random indices within the range of the total number of samples\n            num_samples = len(dataloader.dataset)\n            random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)\n\n            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader\n            random_batch_dataset = dataloader.dataset.select(random_indices)\n            random_batch = self.data_collator(random_batch_dataset)\n            random_batch = self._prepare_inputs(random_batch)\n\n            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, random_batch)\n\n            self.log(\n                {\n                    \"game_log\": wandb.Table(\n                        columns=[\"Prompt\", \"Policy\", \"Ref Model\"],\n                        rows=[\n                            [prompt, pol[len(prompt) :], ref[len(prompt) :]]\n                            for prompt, pol, ref in zip(\n                                random_batch[\"prompt\"], policy_output_decoded, ref_output_decoded\n                            )\n                        ],\n                    )\n                }\n            )\n            self.state.log_history.pop()\n\n        # Base evaluation\n        initial_output = super().evaluation_loop(\n            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix\n        )\n\n        return initial_output\n\n    def log(self, logs: Dict[str, float]) -> None:\n        \"\"\"\n        Log `logs` on the various objects watching training, including stored metrics.\n\n        Args:\n            logs (`Dict[str, float]`):\n                The values to log.\n        \"\"\"\n        # logs either has 'loss' or 'eval_loss'\n        train_eval = \"train\" if \"loss\" in logs else \"eval\"\n        # Add averaged stored metrics to logs\n        for key, metrics in self._stored_metrics[train_eval].items():\n            logs[key] = torch.tensor(metrics).mean().item()\n        del self._stored_metrics[train_eval]\n        return super().log(logs)\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"dpo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# flake8: noqa\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# There is a circular import in the PPOTrainer if we let isort sort these\nfrom typing import TYPE_CHECKING\nfrom ..import_utils import _LazyModule, is_diffusers_available, OptionalDependencyNotAvailable\n\n\n_import_structure = {\n    \"callbacks\": [\"RichProgressCallback\", \"SyncRefModelCallback\"],\n    \"utils\": [\n        \"AdaptiveKLController\",\n        \"FixedKLController\",\n        \"ConstantLengthDataset\",\n        \"DataCollatorForCompletionOnlyLM\",\n        \"RunningMoments\",\n        \"disable_dropout_in_model\",\n        \"peft_module_casting_to_bf16\",\n    ],\n    \"dpo_config\": [\"DPOConfig\", \"FDivergenceConstants\", \"FDivergenceType\"],\n    \"dpo_trainer\": [\"DPOTrainer\"],\n    \"cpo_config\": [\"CPOConfig\"],\n    \"cpo_trainer\": [\"CPOTrainer\"],\n    \"alignprop_config\": [\"AlignPropConfig\"],\n    \"alignprop_trainer\": [\"AlignPropTrainer\"],\n    \"iterative_sft_trainer\": [\"IterativeSFTTrainer\"],\n    \"kto_config\": [\"KTOConfig\"],\n    \"kto_trainer\": [\"KTOTrainer\"],\n    \"bco_config\": [\"BCOConfig\"],\n    \"bco_trainer\": [\"BCOTrainer\"],\n    \"model_config\": [\"ModelConfig\"],\n    \"nash_md_config\": [\"NashMDConfig\"],\n    \"nash_md_trainer\": [\"NashMDTrainer\"],\n    \"online_dpo_config\": [\"OnlineDPOConfig\"],\n    \"online_dpo_trainer\": [\"OnlineDPOTrainer\"],\n    \"xpo_config\": [\"XPOConfig\"],\n    \"xpo_trainer\": [\"XPOTrainer\"],\n    \"orpo_config\": [\"ORPOConfig\"],\n    \"orpo_trainer\": [\"ORPOTrainer\"],\n    \"ppo_config\": [\"PPOConfig\"],\n    \"ppo_trainer\": [\"PPOTrainer\"],\n    \"ppov2_config\": [\"PPOv2Config\"],\n    \"ppov2_trainer\": [\"PPOv2Trainer\"],\n    \"reward_config\": [\"RewardConfig\"],\n    \"reward_trainer\": [\"RewardTrainer\", \"compute_accuracy\"],\n    \"rloo_config\": [\"RLOOConfig\"],\n    \"rloo_trainer\": [\"RLOOTrainer\"],\n    \"sft_config\": [\"SFTConfig\"],\n    \"sft_trainer\": [\"SFTTrainer\"],\n    \"base\": [\"BaseTrainer\"],\n    \"ddpo_config\": [\"DDPOConfig\"],\n    \"gkd_trainer\": [\"GKDTrainer\"],\n    \"gkd_config\": [\"GKDConfig\"],\n    \"callbacks\": [\"RichProgressCallback\", \"SyncRefModelCallback\", \"WinRateCallback\", \"LogCompletionsCallback\"],\n    \"judges\": [\n        \"BaseJudge\",\n        \"BaseRankJudge\",\n        \"BasePairwiseJudge\",\n        \"RandomRankJudge\",\n        \"RandomPairwiseJudge\",\n        \"PairRMJudge\",\n        \"HfPairwiseJudge\",\n        \"OpenAIPairwiseJudge\",\n    ],\n}\n\ntry:\n    if not is_diffusers_available():\n        raise OptionalDependencyNotAvailable()\nexcept OptionalDependencyNotAvailable:\n    pass\nelse:\n    _import_structure[\"ddpo_trainer\"] = [\"DDPOTrainer\"]\n\nif TYPE_CHECKING:\n    # isort: off\n    from .callbacks import RichProgressCallback, SyncRefModelCallback\n    from .utils import (\n        AdaptiveKLController,\n        FixedKLController,\n        ConstantLengthDataset,\n        DataCollatorForCompletionOnlyLM,\n        RunningMoments,\n        disable_dropout_in_model,\n        peft_module_casting_to_bf16,\n        empty_cache,\n    )\n\n    # isort: on\n\n    from .base import BaseTrainer\n    from .ddpo_config import DDPOConfig\n\n    from .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType\n    from .dpo_trainer import DPOTrainer\n    from .iterative_sft_trainer import IterativeSFTTrainer\n    from .cpo_config import CPOConfig\n    from .cpo_trainer import CPOTrainer\n    from .alignprop_config import AlignPropConfig\n    from .alignprop_trainer import AlignPropTrainer\n    from .kto_config import KTOConfig\n    from .kto_trainer import KTOTrainer\n    from .bco_config import BCOConfig\n    from .bco_trainer import BCOTrainer\n    from .model_config import ModelConfig\n    from .nash_md_config import NashMDConfig\n    from .nash_md_trainer import NashMDTrainer\n    from .online_dpo_config import OnlineDPOConfig\n    from .online_dpo_trainer import OnlineDPOTrainer\n    from .xpo_config import XPOConfig\n    from .xpo_trainer import XPOTrainer\n    from .orpo_config import ORPOConfig\n    from .orpo_trainer import ORPOTrainer\n    from .ppo_config import PPOConfig\n    from .ppo_trainer import PPOTrainer\n    from .ppov2_config import PPOv2Config\n    from .ppov2_trainer import PPOv2Trainer\n    from .reward_config import RewardConfig\n    from .reward_trainer import RewardTrainer, compute_accuracy\n    from .rloo_config import RLOOConfig\n    from .rloo_trainer import RLOOTrainer\n    from .sft_config import SFTConfig\n    from .sft_trainer import SFTTrainer\n    from .gkd_trainer import GKDTrainer\n    from .gkd_config import GKDConfig\n    from .callbacks import RichProgressCallback, SyncRefModelCallback, WinRateCallback, LogCompletionsCallback\n    from .judges import (\n        BaseJudge,\n        BaseRankJudge,\n        BasePairwiseJudge,\n        RandomRankJudge,\n        RandomPairwiseJudge,\n        PairRMJudge,\n        HfPairwiseJudge,\n        OpenAIPairwiseJudge,\n    )\n\n    try:\n        if not is_diffusers_available():\n            raise OptionalDependencyNotAvailable()\n    except OptionalDependencyNotAvailable:\n        pass\n    else:\n        from .ddpo_trainer import DDPOTrainer\nelse:\n    import sys\n\n    sys.modules[__name__] = _LazyModule(__name__, globals()[\"__file__\"], _import_structure, module_spec=__spec__)\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom huggingface_hub import PyTorchModelHubMixin\n\n\nclass BaseTrainer(PyTorchModelHubMixin):\n    r\"\"\"\n    Base class for all trainers - this base class implements the basic functions that we\n    need for a trainer.\n\n    The trainer needs to have the following functions:\n        - step: takes in a batch of data and performs a step of training\n        - loss: takes in a batch of data and returns the loss\n        - compute_rewards: takes in a batch of data and returns the rewards\n        - _build_models_and_tokenizer: builds the models and tokenizer\n        - _build_dataset: builds the dataset\n    Each user is expected to implement their own trainer class that inherits from this base\n    if they want to use a new training algorithm.\n    \"\"\"\n\n    def __init__(self, config):\n        self.config = config\n\n    def step(self, *args):\n        raise NotImplementedError(\"Not implemented\")\n\n    def loss(self, *args):\n        raise NotImplementedError(\"Not implemented\")\n\n    def compute_rewards(self, *args):\n        raise NotImplementedError(\"Not implemented\")\n\n    def _save_pretrained(self, save_directory):\n        raise NotImplementedError(\"Not implemented\")\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom datasets import Dataset, IterableDataset\nfrom transformers import PreTrainedTokenizerBase, TrainerCallback, is_apex_available\nfrom transformers.modeling_utils import PreTrainedModel\nfrom transformers.trainer_utils import EvalPrediction\nfrom transformers.training_args import OptimizerNames\n\nfrom ..models.utils import unwrap_model_for_generation\nfrom .online_dpo_trainer import OnlineDPOTrainer\nfrom .utils import empty_cache, get_reward, truncate_right\nfrom .xpo_config import XPOConfig\n\n\nif is_apex_available():\n    from apex import amp\n\n\nclass XPOTrainer(OnlineDPOTrainer):\n    r\"\"\"\n    Initialize XPOTrainer as a subclass of [`OnlineDPOConfig`].\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForCausalLM`.\n        ref_model (`PreTrainedModelWrapper`):\n            Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no\n            reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.\n        reward_model (`transformers.PreTrainedModel`):\n            The reward model to score completions with, preferably an `AutoModelForSequenceClassification`.\n        judge (`BasePairwiseJudge`):\n            The judge to use for pairwise comparison of model completions.\n        args (`XPOConfig`):\n            The XPO config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        peft_config (`Dict`):\n            The peft config to use for training.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"xpo\"]\n\n    def __init__(\n        self,\n        model: Union[PreTrainedModel, nn.Module] = None,\n        ref_model: Union[PreTrainedModel, nn.Module] = None,\n        reward_model: Optional[nn.Module] = None,\n        args: Optional[XPOConfig] = None,\n        data_collator: Optional[Callable] = None,\n        train_dataset: Optional[Union[Dataset, IterableDataset]] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n    ) -> None:\n        super().__init__(\n            model=model,\n            ref_model=ref_model,\n            reward_model=reward_model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            peft_config=peft_config,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        self._alpha = self.args.alpha\n\n        # Overwrite the stats dictionary to include XPO specific statistics\n        self.stats = {\n            # Remove \"non_score_reward\", \"rlhf_reward\", \"scores\"\n            # Add \"loss/dpo\", \"loss/xpo\"\n            \"loss/dpo\": [],\n            \"loss/xpo\": [],\n            \"objective/kl\": [],\n            \"objective/entropy\": [],\n            # Replace \"scores\" by \"model_scores\" and \"ref_scores\"\n            \"objective/model_scores\": [],\n            \"objective/ref_scores\": [],\n            \"objective/scores_margin\": [],\n            \"rewards/chosen\": [],\n            \"rewards/rejected\": [],\n            \"rewards/accuracies\": [],\n            \"rewards/margins\": [],\n            \"logps/chosen\": [],\n            \"logps/rejected\": [],\n            # Replace \"contain_eos_token\" by \"model_contain_eos_token\" and \"ref_contain_eos_token\"\n            \"val/model_contain_eos_token\": [],\n            \"val/ref_contain_eos_token\": [],\n            \"alpha\": [],\n            \"beta\": [],\n        }\n\n    @property\n    def alpha(self):\n        if isinstance(self._alpha, list):\n            epoch = self.state.epoch\n            return self._alpha[epoch] if epoch < len(self._alpha) else self._alpha[-1]\n        else:\n            return self._alpha\n\n    def _generate_completions(self, prompts, model):\n        with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:\n            model_output = unwrapped_model.generate(\n                input_ids=prompts[\"input_ids\"],\n                attention_mask=prompts[\"attention_mask\"],\n                generation_config=self.generation_config,\n            )\n\n        ref_model = model if self.ref_model is None else self.ref_model\n        with torch.no_grad(), unwrap_model_for_generation(ref_model, self.accelerator) as unwrapped_ref_model:\n            ref_output = unwrapped_ref_model.generate(\n                input_ids=prompts[\"input_ids\"],\n                attention_mask=prompts[\"attention_mask\"],\n                generation_config=self.generation_config,\n            )\n\n        return model_output, ref_output\n\n    def _process_completions(self, model_output, ref_output, prompts):\n        context_length = prompts[\"input_ids\"].shape[1]\n\n        # Process model completions\n        model_completion_ids = model_output[:, context_length:]\n        model_completion_ids, model_completion_mask = truncate_right(\n            model_completion_ids, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id\n        )\n        model_data = {\n            \"input_ids\": torch.cat((prompts[\"input_ids\"], model_completion_ids), dim=1),\n            \"attention_mask\": torch.cat((prompts[\"attention_mask\"], model_completion_mask), dim=1),\n        }\n\n        # Process reference model completions\n        ref_completion_ids = ref_output[:, context_length:]\n        ref_completion_ids, ref_completion_mask = truncate_right(\n            ref_completion_ids, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id\n        )\n        ref_data = {\n            \"input_ids\": torch.cat((prompts[\"input_ids\"], ref_completion_ids), dim=1),\n            \"attention_mask\": torch.cat((prompts[\"attention_mask\"], ref_completion_mask), dim=1),\n        }\n\n        return model_data, ref_data\n\n    def _compute_rewards(self, model_data, ref_data, context_length):\n        with torch.no_grad():\n            _, model_scores, _ = get_reward(\n                self.reward_model, model_data[\"input_ids\"], self.tokenizer.pad_token_id, context_length\n            )\n            _, ref_scores, _ = get_reward(\n                self.reward_model, ref_data[\"input_ids\"], self.tokenizer.pad_token_id, context_length\n            )\n\n        # Apply EOS penalty if needed\n        if self.args.missing_eos_penalty is not None:\n            model_contain_eos = torch.any(model_data[\"input_ids\"] == self.tokenizer.eos_token_id, dim=-1)\n            ref_contain_eos = torch.any(ref_data[\"input_ids\"] == self.tokenizer.eos_token_id, dim=-1)\n            model_scores[~model_contain_eos] -= self.args.missing_eos_penalty\n            ref_scores[~ref_contain_eos] -= self.args.missing_eos_penalty\n\n        return model_scores, ref_scores\n\n    def _compute_logprobs(self, model, model_data, ref_data, context_length):\n        def compute_logprobs_for_data(m, data):\n            output = m(data[\"input_ids\"], attention_mask=data[\"attention_mask\"])\n            logits = output.logits[:, context_length - 1 : -1]\n            logprobs = F.log_softmax(logits, dim=-1)\n            token_logprobs = torch.gather(logprobs, 2, data[\"input_ids\"][:, context_length:].unsqueeze(-1)).squeeze(-1)\n            return token_logprobs\n\n        # Compute logprobs for model completions\n        model_logprobs_model_data = compute_logprobs_for_data(model, model_data)\n        # Compute logprobs for model on reference completions (for XPO loss)\n        model_logprobs_ref_data = compute_logprobs_for_data(model, ref_data)\n\n        # Compute logprobs for reference model completions\n        with torch.no_grad():\n            if self.ref_model is None:\n                with model.disable_adapter():\n                    ref_logprobs_model_data = compute_logprobs_for_data(model, model_data)\n                    ref_logprobs_ref_data = compute_logprobs_for_data(model, ref_data)\n            else:\n                ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data)\n                ref_logprobs_ref_data = compute_logprobs_for_data(self.ref_model, ref_data)\n\n        # Mask padding tokens\n        model_padding_mask = model_data[\"attention_mask\"][:, context_length:] == 0\n        ref_padding_mask = ref_data[\"attention_mask\"][:, context_length:] == 0\n        model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0)\n        model_logprobs_ref_data = model_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0)\n        ref_logprobs_ref_data = ref_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0)\n        ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0)\n\n        return model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data\n\n    def _compute_losses(\n        self,\n        model_logprobs_model_data,\n        model_logprobs_ref_data,\n        ref_logprobs_ref_data,\n        ref_logprobs_model_data,\n        chosen_mask,\n    ):\n        # Compute log probs\n        model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)\n        model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1)\n        ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1)\n        ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)\n\n        chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)\n        chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)\n        chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs\n\n        rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)\n        rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)\n        rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs\n\n        # Compute logits as the difference between chosen and rejected log ratios\n        logits = chosen_log_ratios - rejected_log_ratios\n\n        if self.args.loss_type == \"sigmoid\":\n            dpo_losses = -F.logsigmoid(self.beta * logits)\n        elif self.args.loss_type == \"ipo\":\n            dpo_losses = (logits - 1 / (2 * self.beta)) ** 2\n        else:\n            raise NotImplementedError(f\"invalid loss type {self.args.loss_type}\")\n\n        # Compute XPO specific loss\n        xpo_losses = self.alpha * model_logprobs_ref_data_sum\n\n        # Total loss\n        loss = (dpo_losses + xpo_losses).mean()\n\n        return loss, dpo_losses, xpo_losses\n\n    def _log_statistics(\n        self,\n        model_data,\n        ref_data,\n        model_logprobs_model_data,\n        model_logprobs_ref_data,\n        ref_logprobs_ref_data,\n        ref_logprobs_model_data,\n        model_scores,\n        ref_scores,\n        dpo_losses,\n        xpo_losses,\n        context_length,\n    ):\n        # Helper function to gather and compute mean\n        def gather_mean(tensor):\n            return self.accelerator.gather(tensor).mean().item()\n\n        # Log losses\n        self.stats[\"loss/dpo\"].append(gather_mean(dpo_losses))\n        self.stats[\"loss/xpo\"].append(gather_mean(xpo_losses))\n\n        # Log scores\n        self.stats[\"objective/model_scores\"].append(gather_mean(model_scores))\n        self.stats[\"objective/ref_scores\"].append(gather_mean(ref_scores))\n        self.stats[\"objective/scores_margin\"].append(gather_mean(model_scores - ref_scores))\n\n        # Determine which model outputs are \"chosen\" vs \"rejected\"\n        chosen_mask = model_scores >= ref_scores\n\n        # Log logprobs\n        model_logprobs_model_data_sum = model_logprobs_model_data.sum(1)\n        model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1)\n        ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1)\n        ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1)\n\n        chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)\n        chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)\n        chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs\n\n        rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum)\n        rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)\n        rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs\n\n        self.stats[\"logps/chosen\"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean()))\n        self.stats[\"logps/rejected\"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean()))\n\n        # Log rewards\n        # Compute various statistics\n        chosen_rewards = chosen_log_ratios * self.beta\n        rejected_rewards = rejected_log_ratios * self.beta\n        self.stats[\"rewards/chosen\"].append(gather_mean(chosen_rewards.mean()))\n        self.stats[\"rewards/rejected\"].append(gather_mean(rejected_rewards.mean()))\n\n        # Calculate KL divergence for model and ref data\n        kl_model_data = model_logprobs_model_data - ref_logprobs_model_data\n        kl_ref_data = model_logprobs_ref_data - ref_logprobs_ref_data\n        mean_kl = (kl_model_data.sum(1) + kl_ref_data.sum(1)).mean() / 2\n        self.stats[\"objective/kl\"].append(gather_mean(mean_kl))\n\n        # Calculate entropy for model and ref data\n        entropy_model_data = -model_logprobs_model_data.sum(1)\n        entropy_ref_data = -model_logprobs_ref_data.sum(1)\n        mean_entropy = (entropy_model_data.mean() + entropy_ref_data.mean()) / 2\n        self.stats[\"objective/entropy\"].append(gather_mean(mean_entropy))\n\n        # Calculate margins\n        margin = chosen_rewards - rejected_rewards\n        self.stats[\"rewards/margins\"].append(gather_mean(margin.mean()))\n\n        # Calculate accuracy\n        accuracy = (margin > 0).float()\n        self.stats[\"rewards/accuracies\"].append(gather_mean(accuracy.mean()))\n\n        # Log EOS token statistics\n        model_eos = (model_data[\"input_ids\"][:, context_length:] == self.tokenizer.eos_token_id).any(dim=1)\n        ref_eos = (ref_data[\"input_ids\"][:, context_length:] == self.tokenizer.eos_token_id).any(dim=1)\n        self.stats[\"val/model_contain_eos_token\"].append(gather_mean(model_eos.float()))\n        self.stats[\"val/ref_contain_eos_token\"].append(gather_mean(ref_eos.float()))\n\n        # Log alpha and beta\n        self.stats[\"alpha\"].append(self.alpha)\n        self.stats[\"beta\"].append(self.beta)\n\n    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:\n        model.train()\n\n        # need the prompt_ only\n        inputs = self._prepare_inputs(inputs)\n        context_length = inputs[\"prompt_input_ids\"].shape[1]\n        prompts = {\n            \"input_ids\": inputs[\"prompt_input_ids\"],\n            \"attention_mask\": inputs[\"prompt_attention_mask\"],\n        }\n        del inputs\n\n        # Sample completions from both the model and the reference model\n        model_output, ref_output = self._generate_completions(prompts, model)\n\n        # Process model completions\n        model_data, ref_data = self._process_completions(model_output, ref_output, prompts)\n\n        # Compute rewards\n        model_data_scores, ref_data_scores = self._compute_rewards(model_data, ref_data, context_length)\n\n        # Compute logprobs\n        model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data = (\n            self._compute_logprobs(model, model_data, ref_data, context_length)\n        )\n\n        # Compute loss\n        loss, dpo_losses, xpo_losses = self._compute_losses(\n            model_logprobs_model_data,\n            model_logprobs_ref_data,\n            ref_logprobs_ref_data,\n            ref_logprobs_model_data,\n            model_data_scores >= ref_data_scores,\n        )\n\n        # Log everything\n        self._log_statistics(\n            model_data,\n            ref_data,\n            model_logprobs_model_data.detach(),\n            model_logprobs_ref_data.detach(),\n            ref_logprobs_ref_data,\n            ref_logprobs_model_data,\n            model_data_scores,\n            ref_data_scores,\n            dpo_losses.detach(),\n            xpo_losses.detach(),\n            context_length,\n        )\n\n        if (\n            self.args.torch_empty_cache_steps is not None\n            and self.state.global_step % self.args.torch_empty_cache_steps == 0\n        ):\n            empty_cache()\n\n        kwargs = {}\n        # For LOMO optimizers you need to explicitly use the learning rate\n        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:\n            kwargs[\"learning_rate\"] = self._get_learning_rate()\n\n        if self.args.n_gpu > 1:\n            loss = loss.mean()  # mean() to average on multi-gpu parallel training\n\n        if self.use_apex:\n            with amp.scale_loss(loss, self.optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            self.accelerator.backward(loss, **kwargs)\n\n        return loss.detach() / self.args.gradient_accumulation_steps\n\n\n# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import List, Literal, Optional\n\n\n@dataclass\nclass ModelConfig:\n    \"\"\"\n    Configuration class for the models.\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        model_name_or_path (`Optional[str]`, *optional*, defaults to `None`):\n            Model checkpoint for weights initialization.\n        model_revision (`str`, *optional*, defaults to `\"main\"`):\n            Specific model version to use. It can be a branch name, a tag name, or a commit id.\n        torch_dtype (`Optional[Literal[\"auto\", \"bfloat16\", \"float16\", \"float32\"]]`, *optional*, defaults to `None`):\n            Override the default `torch.dtype` and load the model under this dtype. Possible values are\n\n                - `\"bfloat16\"`: `torch.bfloat16`\n                - `\"float16\"`: `torch.float16`\n                - `\"float32\"`: `torch.float32`\n                - `\"auto\"`: Automatically derive the dtype from the model's weights.\n\n        trust_remote_code (`bool`, *optional*, defaults to `False`):\n            Whether to allow for custom models defined on the Hub in their own modeling files. This option should only\n            be set to `True` for repositories you trust and in which you have read the code, as it will execute code\n            present on the Hub on your local machine.\n        attn_implementation (`Optional[str]`, *optional*, defaults to `None`):\n            Which attention implementation to use. You can run `--attn_implementation=flash_attention_2`, in which case\n            you must install this manually by running `pip install flash-attn --no-build-isolation`.\n        use_peft (`bool`, *optional*, defaults to `False`):\n            Whether to use PEFT for training.\n        lora_r (`int`, *optional*, defaults to `16`):\n            LoRA R value.\n        lora_alpha (`int`, *optional*, defaults to `32`):\n            LoRA alpha.\n        lora_dropout (`float`, *optional*, defaults to `0.05`):\n            LoRA dropout.\n        lora_target_modules (`Optional[Union[str, List[str]]]`, *optional*, defaults to `None`):\n            LoRA target modules.\n        lora_modules_to_save (`Optional[List[str]]`, *optional*, defaults to `None`):\n            Model layers to unfreeze & train.\n        lora_task_type (`str`, *optional*, defaults to `\"CAUSAL_LM\"`):\n            Task type to pass for LoRA (use `\"SEQ_CLS\"` for reward modeling).\n        use_rslora (`bool`, *optional*, defaults to `False`):\n            Whether to use Rank-Stabilized LoRA, which sets the adapter scaling factor to `lora_alpha/√r`, instead of\n            the original default value of `lora_alpha/r`.\n        load_in_8bit (`bool`, *optional*, defaults to `False`):\n            Whether to use 8 bit precision for the base model. Works only with LoRA.\n        load_in_4bit (`bool`, *optional*, defaults to `False`):\n            Whether to use 4 bit precision for the base model. Works only with LoRA.\n        bnb_4bit_quant_type (`str`, *optional*, defaults to `\"nf4\"`):\n            Quantization type (`\"fp4\"` or `\"nf4\"`).\n        use_bnb_nested_quant (`bool`, *optional*, defaults to `False`):\n            Whether to use nested quantization.\n    \"\"\"\n\n    model_name_or_path: Optional[str] = None\n    model_revision: str = \"main\"\n    torch_dtype: Optional[Literal[\"auto\", \"bfloat16\", \"float16\", \"float32\"]] = None\n    trust_remote_code: bool = False\n    attn_implementation: Optional[str] = None\n    use_peft: bool = False\n    lora_r: int = 16\n    lora_alpha: int = 32\n    lora_dropout: float = 0.05\n    lora_target_modules: Optional[List[str]] = None\n    lora_modules_to_save: Optional[List[str]] = None\n    lora_task_type: str = \"CAUSAL_LM\"\n    use_rslora: bool = False\n    load_in_8bit: bool = False\n    load_in_4bit: bool = False\n    bnb_4bit_quant_type: Literal[\"fp4\", \"nf4\"] = \"nf4\"\n    use_bnb_nested_quant: bool = False\n\n    def __post_init__(self):\n        if self.load_in_8bit and self.load_in_4bit:\n            raise ValueError(\"You can't use 8 bit and 4 bit precision at the same time\")\n\n        if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1:\n            self.lora_target_modules = self.lora_target_modules[0]\n\n\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import List, Optional, Union\n\nimport torch\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import gather_object, is_deepspeed_available\nfrom rich.console import Console, Group\nfrom rich.live import Live\nfrom rich.panel import Panel\nfrom rich.progress import Progress\nfrom transformers import (\n    GenerationConfig,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    TrainerCallback,\n    TrainerControl,\n    TrainerState,\n    TrainingArguments,\n)\nfrom transformers.integrations import WandbCallback\nfrom transformers.trainer_utils import has_length\n\nfrom ..models.utils import unwrap_model_for_generation\nfrom .judges import BasePairwiseJudge\n\n\nif is_deepspeed_available():\n    import deepspeed\n\n\ndef _generate_completions(\n    prompts: List[str],\n    model: PreTrainedModel,\n    tokenizer: PreTrainedTokenizerBase,\n    accelerator: Accelerator,\n    generation_config: Optional[GenerationConfig],\n    batch_size: int = 1,\n) -> List[str]:\n    \"\"\"\n    Generates completions for a list of pre-formatted prompts.\n\n    Args:\n        prompts (List[str]): A list of input prompts for which completions are to be generated.\n        model (PreTrainedModel): The pre-trained model to be used for generation.\n        tokenizer (PreTrainedTokenizerBase): The tokenizer to be used for encoding and decoding.\n        accelerator (Accelerator): The accelerator to be used for model execution.\n        generation_config (GenerationConfig): Configuration for text generation.\n        batch_size (int, optional): The number of prompts to process in each batch. Default is 1.\n\n    Returns:\n        List[str]: A list of generated text completions corresponding to the input prompts.\n    \"\"\"\n    completions = []\n    with unwrap_model_for_generation(model, accelerator) as unwrapped_model:\n        unwrapped_model.eval()\n        for idx in range(0, len(prompts), batch_size):\n            batch = prompts[idx : idx + batch_size]\n            tokenized_batch = tokenizer(batch, return_tensors=\"pt\", padding=True, truncation=True).to(model.device)\n            generations = unwrapped_model.generate(\n                **tokenized_batch,\n                generation_config=generation_config,\n            )\n            for prompt, generation in zip(tokenized_batch.input_ids, generations):\n                # Remove prompt from generation\n                generation = generation[len(prompt) :]\n                completion = tokenizer.decode(generation, skip_special_tokens=True)\n                completions.append(completion)\n        unwrapped_model.train()\n    return completions\n\n\nclass SyncRefModelCallback(TrainerCallback):\n    def __init__(\n        self,\n        ref_model: Union[PreTrainedModel, torch.nn.Module],\n        accelerator: Optional[Accelerator],\n    ):\n        self.accelerator = accelerator\n        self.ref_model = ref_model\n\n    @staticmethod\n    def _sync_target_model(model, target_model, alpha):\n        for target_param, copy_param in zip(target_model.parameters(), model.parameters()):\n            target_param.data.mul_(1.0 - alpha).add_(copy_param.data, alpha=alpha)\n\n    @staticmethod\n    def sync_target_model(model, target_model, alpha):\n        deepspeed_plugin = AcceleratorState().deepspeed_plugin\n        if deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3:\n            with deepspeed.zero.GatheredParameters(\n                list(model.parameters()) + list(target_model.parameters()), modifier_rank=0\n            ):\n                if deepspeed.comm.get_rank() == 0:\n                    SyncRefModelCallback._sync_target_model(model, target_model, alpha)\n        else:\n            SyncRefModelCallback._sync_target_model(model, target_model, alpha)\n\n    def on_step_end(self, args, state, control, **kwargs):\n        model: PreTrainedModel = kwargs[\"model\"]\n\n        if self.ref_model is not None and state.global_step % args.ref_model_sync_steps == 0:\n            if self.accelerator:\n                model = self.accelerator.unwrap_model(model)\n            self.sync_target_model(model, self.ref_model, args.ref_model_mixup_alpha)\n\n\nclass RichProgressCallback(TrainerCallback):\n    \"\"\"\n    A [`TrainerCallback`] that displays the progress of training or evaluation using Rich.\n    \"\"\"\n\n    def __init__(self):\n        self.training_bar = None\n        self.prediction_bar = None\n\n        self.training_task_id = None\n        self.prediction_task_id = None\n\n        self.rich_group = None\n        self.rich_console = None\n\n        self.training_status = None\n        self.current_step = None\n\n    def on_train_begin(self, args, state, control, **kwargs):\n        if state.is_world_process_zero:\n            self.training_bar = Progress()\n            self.prediction_bar = Progress()\n\n            self.rich_console = Console()\n\n            self.training_status = self.rich_console.status(\"Nothing to log yet ...\")\n\n            self.rich_group = Live(Panel(Group(self.training_bar, self.prediction_bar, self.training_status)))\n            self.rich_group.start()\n\n            self.training_task_id = self.training_bar.add_task(\"[blue]Training the model\", total=state.max_steps)\n            self.current_step = 0\n\n    def on_step_end(self, args, state, control, **kwargs):\n        if state.is_world_process_zero:\n            self.training_bar.update(self.training_task_id, advance=state.global_step - self.current_step, update=True)\n            self.current_step = state.global_step\n\n    def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):\n        if state.is_world_process_zero and has_length(eval_dataloader):\n            if self.prediction_task_id is None:\n                self.prediction_task_id = self.prediction_bar.add_task(\n                    \"[blue]Predicting on the evaluation dataset\", total=len(eval_dataloader)\n                )\n            self.prediction_bar.update(self.prediction_task_id, advance=1, update=True)\n\n    def on_evaluate(self, args, state, control, **kwargs):\n        if state.is_world_process_zero:\n            if self.prediction_task_id is not None:\n                self.prediction_bar.remove_task(self.prediction_task_id)\n                self.prediction_task_id = None\n\n    def on_predict(self, args, state, control, **kwargs):\n        if state.is_world_process_zero:\n            if self.prediction_task_id is not None:\n                self.prediction_bar.remove_task(self.prediction_task_id)\n                self.prediction_task_id = None\n\n    def on_log(self, args, state, control, logs=None, **kwargs):\n        if state.is_world_process_zero and self.training_bar is not None:\n            _ = logs.pop(\"total_flos\", None)\n            self.training_status.update(f\"[bold green]Status = {str(logs)}\")\n\n    def on_train_end(self, args, state, control, **kwargs):\n        if state.is_world_process_zero:\n            self.rich_group.stop()\n\n            self.training_bar = None\n            self.prediction_bar = None\n            self.training_task_id = None\n            self.prediction_task_id = None\n            self.rich_group = None\n            self.rich_console = None\n            self.training_status = None\n            self.current_step = None\n\n\nclass WinRateCallback(TrainerCallback):\n    \"\"\"\n    A [`~transformers.TrainerCallback`] that computes the win rate of a model based on a reference.\n\n    It generates completions using prompts from the evaluation dataset and compares the trained model's outputs against\n    a reference. The reference is either the initial version of the model (before training) or the reference model, if\n    available in the trainer. During each evaluation step, a judge determines how often the trained model's completions\n    win against the reference using a judge. The win rate is then logged in the trainer's logs under the key\n    `\"eval_win_rate\"`.\n\n    Usage:\n    ```python\n    trainer = DPOTrainer(...)\n    judge = PairRMJudge()\n    win_rate_callback = WinRateCallback(judge=judge, trainer=trainer)\n    trainer.add_callback(win_rate_callback)\n    ```\n\n    Args:\n        judge (`BasePairwiseJudge`):\n            The judge to use for comparing completions.\n        trainer (`Trainer`):\n            Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `\"prompt\"`\n            column containing the prompts for generating completions. If the `Trainer` has a reference model (via the\n            `ref_model` attribute), it will use this reference model for generating the reference completions;\n            otherwise, it defaults to using the initial model.\n        generation_config (`GenerationConfig`, *optional*):\n            The generation config to use for generating completions.\n        num_prompts (`int`, *optional*):\n            The number of prompts to generate completions for. If not provided, defaults to the number of examples\n            in the evaluation dataset.\n    \"\"\"\n\n    def __init__(\n        self,\n        judge: BasePairwiseJudge,\n        trainer: Trainer,\n        generation_config: Optional[GenerationConfig] = None,\n        num_prompts: int = None,\n    ):\n        self.judge = judge\n        self.trainer = trainer\n        self.generation_config = generation_config\n        self.ref_completions = []\n\n        if self.trainer.eval_dataset is None:\n            raise ValueError(\"Trainer must have an evaluation dataset to use the WinRateCallback.\")\n        else:\n            self.eval_dataset = self.trainer.eval_dataset\n\n        if num_prompts is not None:\n            self.eval_dataset = self.eval_dataset.select(range(num_prompts))\n\n    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):\n        # When the trainer is initialized, we generate completions for the reference model.\n        tokenizer = kwargs[\"tokenizer\"]\n        tokenizer.padding_side = \"left\"\n        accelerator = self.trainer.accelerator\n        # Use the reference model if available, otherwise use the initial model\n        model = getattr(self.trainer, \"ref_model\", None)\n        # At this point, there are two cases where `ref_model` is None:\n        # 1. The method doesn't require a reference model.\n        # 2. The method uses a reference model, but `ref_model` is set to None.\n        #    This occurs when using PEFT, where the reference model can be obtained by simply disabling the model's adapter.\n        #    In theory, we should disable the adapter here, but since it's zero-initialized at the start of training,\n        #    the model behaves identically with or without the adapter.\n        #    Therefore, there's no need to explicitly disable it at this point.\n        if model is None:\n            model = self.trainer.model_wrapped\n        with accelerator.split_between_processes(self.eval_dataset[\"prompt\"]) as prompts:\n            self.ref_completions = _generate_completions(\n                prompts,\n                model=model,\n                tokenizer=tokenizer,\n                accelerator=accelerator,\n                generation_config=self.generation_config,\n                batch_size=args.per_device_eval_batch_size,\n            )\n\n    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):\n        # At every evaluation step, we generate completions for the model and compare them with the reference\n        # completions that have been generated at the beginning of training. We then compute the win rate and log it to\n        # the trainer.\n        tokenizer = kwargs[\"tokenizer\"]\n        tokenizer.padding_side = \"left\"\n        accelerator = self.trainer.accelerator\n        model = self.trainer.model_wrapped\n        with accelerator.split_between_processes(self.eval_dataset[\"prompt\"]) as prompts:\n            completions = _generate_completions(\n                prompts,\n                model=model,\n                tokenizer=tokenizer,\n                accelerator=accelerator,\n                generation_config=self.generation_config,\n                batch_size=args.per_device_eval_batch_size,\n            )\n\n            completions = list(zip(self.ref_completions, completions))\n            winner_indices = self.judge.judge(prompts, completions)\n            winner_indices = gather_object(winner_indices)\n\n        # Logging\n        if self.trainer.accelerator.is_main_process:\n            win_rate = sum(winner_idx == 1 for winner_idx in winner_indices) / len(winner_indices)\n            self.trainer.log({\"eval_win_rate\": win_rate})\n\n\nclass LogCompletionsCallback(WandbCallback):\n    r\"\"\"\n    A [`~transformers.TrainerCallback`] that logs completions to Weights & Biases.\n\n    Usage:\n    ```python\n    trainer = DPOTrainer(...)\n    completions_callback = LogCompletionsCallback(trainer=trainer)\n    trainer.add_callback(completions_callback)\n    ```\n\n    Args:\n        trainer (`Trainer`):\n            Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `\"prompt\"`\n            column containing the prompts for generating completions.\n        generation_config (`GenerationConfig`, *optional*):\n            The generation config to use for generating completions.\n        num_prompts (`int`, *optional*):\n            The number of prompts to generate completions for. If not provided, defaults to the number of examples in the evaluation dataset.\n        freq (`int`, *optional*):\n            The frequency at which to log completions. If not provided, defaults to the trainer's `eval_steps`.\n    \"\"\"\n\n    def __init__(\n        self,\n        trainer: Trainer,\n        generation_config: Optional[GenerationConfig] = None,\n        num_prompts: int = None,\n        freq: int = None,\n    ):\n        super().__init__()\n        self.trainer = trainer\n        self.generation_config = generation_config\n        self.freq = freq\n        self.table = []\n        self._last_logged_step = -1\n\n        if self.trainer.eval_dataset is None:\n            raise ValueError(\"Trainer must have an evaluation dataset to use the LogCompletionsCallback.\")\n        else:\n            self.eval_dataset = self.trainer.eval_dataset\n\n        if num_prompts is not None:\n            self.eval_dataset = self.eval_dataset.select(range(num_prompts))\n\n    def on_step_end(self, args, state, control, **kwargs):\n        # Only log once per step (this method may be called multiple times)\n        if state.global_step == self._last_logged_step:\n            return\n\n        # Only log every `freq` steps (if no `freq` is provided, log every `eval_steps` steps)\n        freq = self.freq or state.eval_steps\n        if state.global_step % freq != 0:\n            return\n\n        tokenizer = kwargs[\"tokenizer\"]\n        tokenizer.padding_side = \"left\"\n        accelerator = self.trainer.accelerator\n        model = self.trainer.model_wrapped\n        with accelerator.split_between_processes(self.eval_dataset[\"prompt\"]) as prompts:\n            completions = _generate_completions(\n                prompts,\n                model=model,\n                tokenizer=tokenizer,\n                accelerator=accelerator,\n                generation_config=self.generation_config,\n                batch_size=args.per_device_eval_batch_size,\n            )\n            completions = gather_object(completions)\n            prompts = gather_object(prompts)\n\n        # Build the data to log\n        if self.trainer.accelerator.is_main_process:\n            # prompts = self.eval_dataset[\"prompt\"][:]\n            global_step = [str(state.global_step)] * len(prompts)\n            data = list(zip(global_step, prompts, completions))\n            self.table.extend(data)\n            table = self._wandb.Table(columns=[\"step\", \"prompt\", \"completion\"], data=self.table)\n            self._wandb.log({\"completions\": table})\n\n        # Save the last logged step, so we don't log the same completions multiple times\n        self._last_logged_step = state.global_step\n\n\n# CPO Authors: Haoran Xu, Amr Sharaf, Yunmo Chen, Weiting Tan, Lingfeng Shen, Benjamin Van Durme, Kenton Murray, Young Jin Kim\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom contextlib import nullcontext\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.amp as amp\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom accelerate import PartialState\nfrom datasets import Dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    AutoModelForCausalLM,\n    DataCollator,\n    PreTrainedModel,\n    PreTrainedTokenizerBase,\n    Trainer,\n    is_wandb_available,\n)\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_utils import EvalLoopOutput\nfrom transformers.utils import is_peft_available, is_torch_fx_proxy\n\nfrom .cpo_config import CPOConfig\nfrom .utils import (\n    DPODataCollatorWithPadding,\n    add_bos_token_if_needed,\n    add_eos_token_if_needed,\n    disable_dropout_in_model,\n    pad_to_length,\n    peft_module_casting_to_bf16,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\n\nif is_wandb_available():\n    import wandb\n\n\nclass CPOTrainer(Trainer):\n    r\"\"\"\n    Initialize CPOTrainer.\n\n    Args:\n        model (`transformers.PreTrainedModel`):\n            The model to train, preferably an `AutoModelForSequenceClassification`.\n        args (`CPOConfig`):\n            The CPO config arguments to use for training.\n        data_collator (`transformers.DataCollator`):\n            The data collator to use for training. If None is specified, the default data collator (`DPODataCollatorWithPadding`) will be used\n            which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n        train_dataset (`datasets.Dataset`):\n            The dataset to use for training.\n        eval_dataset (`datasets.Dataset`):\n            The dataset to use for evaluation.\n        tokenizer (`transformers.PreTrainedTokenizerBase`):\n            The tokenizer to use for training. This argument is required if you want to use the default data collator.\n        model_init (`Callable[[], transformers.PreTrainedModel]`):\n            The model initializer to use for training. If None is specified, the default model initializer will be used.\n        callbacks (`List[transformers.TrainerCallback]`):\n            The callbacks to use for training.\n        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n            The optimizer and scheduler to use for training.\n        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n            The function to use to preprocess the logits before computing the metrics.\n        peft_config (`Dict`, defaults to `None`):\n            The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):\n            The function to use to compute the metrics. Must take a `EvalPrediction` and return\n            a dictionary string to metric values.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"cpo\"]\n\n    def __init__(\n        self,\n        model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,\n        args: Optional[CPOConfig] = None,\n        data_collator: Optional[DataCollator] = None,\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        peft_config: Optional[Dict] = None,\n        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,\n    ):\n        if args.model_init_kwargs is None:\n            model_init_kwargs = {}\n        elif not isinstance(model, str):\n            raise ValueError(\"You passed model_kwargs to the CPOTrainer. But your model is already instantiated.\")\n        else:\n            model_init_kwargs = args.model_init_kwargs\n            torch_dtype = model_init_kwargs.get(\"torch_dtype\")\n            if torch_dtype is not None:\n                # Convert to `torch.dtype` if an str is passed\n                if isinstance(torch_dtype, str) and torch_dtype != \"auto\":\n                    torch_dtype = getattr(torch, torch_dtype)\n                if torch_dtype != \"auto\" and not isinstance(torch_dtype, torch.dtype):\n                    raise ValueError(\n                        f\"Invalid `torch_dtype` passed to the CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}.\"\n                    )\n                model_init_kwargs[\"torch_dtype\"] = torch_dtype\n\n        if isinstance(model, str):\n            warnings.warn(\n                \"You passed a model_id to the CPOTrainer. This will automatically create an \"\n                \"`AutoModelForCausalLM` or a `PeftModel` (if you passed a `peft_config`) for you.\"\n            )\n            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)\n\n        # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16`\n        # has been called in order to properly call autocast if needed.\n        self._peft_has_been_casted_to_bf16 = False\n\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            # if model is a peft model and we have a peft_config, we merge and unload it first\n            if isinstance(model, PeftModel):\n                model = model.merge_and_unload()\n\n            if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_loaded_in_4bit\", False):\n                _support_gc_kwargs = hasattr(\n                    args, \"gradient_checkpointing_kwargs\"\n                ) and \"gradient_checkpointing_kwargs\" in list(\n                    inspect.signature(prepare_model_for_kbit_training).parameters\n                )\n\n                prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                if _support_gc_kwargs:\n                    prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n            elif getattr(args, \"gradient_checkpointing\", False):\n                # For backward compatibility with older versions of transformers\n                if hasattr(model, \"enable_input_require_grads\"):\n                    model.enable_input_require_grads()\n                else:\n\n                    def make_inputs_require_grad(module, input, output):\n                        output.requires_grad_(True)\n\n                    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n            # get peft model with the given config\n            model = get_peft_model(model, peft_config)\n            if args.bf16 and getattr(model, \"is_loaded_in_4bit\", False):\n                peft_module_casting_to_bf16(model)\n                # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager\n                self._peft_has_been_casted_to_bf16 = True\n\n        # For models that use gradient_checkpointing, we need to attach a hook that enables input\n        # to explicitly have `requires_grad=True`, otherwise training will either silently\n        # fail or completely fail.\n        elif getattr(args, \"gradient_checkpointing\", False):\n            # For backward compatibility with older versions of transformers\n            if hasattr(model, \"enable_input_require_grads\"):\n                model.enable_input_require_grads()\n            else:\n\n                def make_inputs_require_grad(module, input, output):\n                    output.requires_grad_(True)\n\n                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n\n        if args.generate_during_eval and not is_wandb_available():\n            raise ValueError(\n                \"`generate_during_eval=True` requires Weights and Biases to be installed.\"\n                \" Please install `wandb` to resolve.\"\n            )\n\n        if model is not None:\n            self.is_encoder_decoder = model.config.is_encoder_decoder\n        elif args.is_encoder_decoder is None:\n            raise ValueError(\"When no model is provided, you need to pass the parameter is_encoder_decoder.\")\n        else:\n            self.is_encoder_decoder = args.is_encoder_decoder\n\n        if self.is_encoder_decoder:\n            self.decoder_start_token_id = model.config.decoder_start_token_id\n            self.pad_token_id = model.config.pad_token_id\n\n        if tokenizer is None:\n            raise ValueError(\"tokenizer must be specified to tokenize a CPO dataset.\")\n        if args.max_length is None:\n            warnings.warn(\n                \"`max_length` is not set in the CPOConfig's init\"\n                \" it will default to `512` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_length = 512\n        else:\n            max_length = args.max_length\n        if args.max_prompt_length is None:\n            warnings.warn(\n                \"`max_prompt_length` is not set in the CPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_prompt_length = 128\n        else:\n            max_prompt_length = args.max_prompt_length\n\n        if args.max_completion_length is None and self.is_encoder_decoder:\n            warnings.warn(\n                \"When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init\"\n                \" it will default to `128` by default, but you should do it yourself in the future.\",\n                UserWarning,\n            )\n            max_completion_length = 128\n        else:\n            max_completion_length = args.max_completion_length\n\n        if data_collator is None:\n            data_collator = DPODataCollatorWithPadding(\n                pad_token_id=tokenizer.pad_token_id,\n                label_pad_token_id=args.label_pad_token_id,\n                is_encoder_decoder=self.is_encoder_decoder,\n            )\n\n            if args.remove_unused_columns:\n                args.remove_unused_columns = False\n                # warn users\n                warnings.warn(\n                    \"When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_dpo_data_collator = True\n        else:\n            self.use_dpo_data_collator = False\n\n        if args.disable_dropout:\n            disable_dropout_in_model(model)\n\n        self.max_length = max_length\n        self.generate_during_eval = args.generate_during_eval\n        self.label_pad_token_id = args.label_pad_token_id\n        self.padding_value = args.padding_value if args.padding_value is not None else tokenizer.pad_token_id\n        self.max_prompt_length = max_prompt_length\n        self.truncation_mode = args.truncation_mode\n        self.max_completion_length = max_completion_length\n        self.tokenizer = tokenizer\n\n        if args.loss_type in [\"hinge\", \"ipo\"] and args.label_smoothing > 0:\n            warnings.warn(\n                \"You are using a loss type that does not support label smoothing. Ignoring label_smoothing parameter.\"\n            )\n        if args.loss_type == \"kto_pair\":\n            raise ValueError(\"Support for kto_pair has been removed in CPOTrainer. Please use KTOTrainer.\")\n\n        self.beta = args.beta\n        self.label_smoothing = args.label_smoothing\n        self.loss_type = args.loss_type\n        self.cpo_alpha = args.cpo_alpha\n        self.aux_loss_enabled = getattr(model.config, \"output_router_logits\", False)\n\n        if args.loss_type == \"simpo\":\n            self.simpo_gamma = args.simpo_gamma\n            if self.cpo_alpha > 0:\n                warnings.warn(\n                    \"You are using CPO-SimPO method because you set a non-zero cpo_alpha. \"\n                    \"This will result in the CPO-SimPO method \"\n                    \"(https://github.com/fe1ixxu/CPO_SIMPO/tree/main). \"\n                    \"If you want to use a pure SimPO method, please set cpo_alpha to 0.\"\n                )\n\n        self._stored_metrics = defaultdict(lambda: defaultdict(list))\n\n        # Compute that only on the main process for faster data processing.\n        # see: https://github.com/huggingface/trl/pull/1255\n        with PartialState().local_main_process_first():\n            # tokenize the dataset\n            train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)\n            if eval_dataset is not None:\n                eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc)\n\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n        if not hasattr(self, \"accelerator\"):\n            raise AttributeError(\n                \"Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`.\"\n            )\n\n    def build_tokenized_answer(self, prompt, answer):\n        \"\"\"\n        Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`.\n        It does ensure `enc(a + b) = enc(a) + enc(a + b)[len(enc(a)):]`.\n        Reference:\n            https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n        \"\"\"\n\n        full_tokenized = self.tokenizer(prompt + answer, add_special_tokens=False)\n        prompt_input_ids = self.tokenizer(prompt, add_special_tokens=False)[\"input_ids\"]\n\n        answer_input_ids = full_tokenized[\"input_ids\"][len(prompt_input_ids) :]\n        answer_attention_mask = full_tokenized[\"attention_mask\"][len(prompt_input_ids) :]\n\n        # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]`\n        full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids])\n\n        # Prepare input tokens for token by token comparison\n        full_input_ids = np.array(full_tokenized[\"input_ids\"])\n\n        if len(full_input_ids) != len(full_concat_input_ids):\n            raise ValueError(\"Prompt input ids and answer input ids should have the same length.\")\n\n        # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens\n        # can be merged together when tokenizing prompt+answer. This could result\n        # on the last token from the prompt being different when tokenized on its own\n        # vs when done as prompt+answer.\n        response_token_ids_start_idx = len(prompt_input_ids)\n\n        # If tokenized prompt is different than both prompt+answer, then it means the\n        # last token has changed due to merging.\n        if prompt_input_ids != full_tokenized[\"input_ids\"][:response_token_ids_start_idx]:\n            response_token_ids_start_idx -= 1\n\n        prompt_input_ids = full_tokenized[\"input_ids\"][:response_token_ids_start_idx]\n        prompt_attention_mask = full_tokenized[\"attention_mask\"][:response_token_ids_start_idx]\n\n        if len(prompt_input_ids) != len(prompt_attention_mask):\n            raise ValueError(\"Prompt input ids and attention mask should have the same length.\")\n\n        answer_input_ids = full_tokenized[\"input_ids\"][response_token_ids_start_idx:]\n        answer_attention_mask = full_tokenized[\"attention_mask\"][response_token_ids_start_idx:]\n\n        return dict(\n            prompt_input_ids=prompt_input_ids,\n            prompt_attention_mask=prompt_attention_mask,\n            input_ids=answer_input_ids,\n            attention_mask=answer_attention_mask,\n        )\n\n    def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> Dict:\n        \"\"\"Tokenize a single row from a CPO specific dataset.\n\n        At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation\n        in case the prompt + chosen or prompt + rejected responses is/are too long. First\n        we truncate the prompt; if we're still too long, we truncate the chosen/rejected.\n\n        We also create the labels for the chosen/rejected responses, which are of length equal to\n        the sum of the length of the prompt and the chosen/rejected response, with\n        label_pad_token_id  for the prompt tokens.\n        \"\"\"\n        batch = {}\n        prompt = feature[\"prompt\"]\n        chosen = feature[\"chosen\"]\n        rejected = feature[\"rejected\"]\n\n        if not self.is_encoder_decoder:\n            # Check issues below for more details\n            #  1. https://github.com/huggingface/trl/issues/907\n            #  2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257\n            #  3. https://github.com/LianjiaTech/BELLE/issues/337\n\n            if not isinstance(prompt, str):\n                raise ValueError(f\"prompt should be an str but got {type(prompt)}\")\n            prompt_tokens = self.tokenizer(prompt, add_special_tokens=False)\n            prompt_tokens = {f\"prompt_{k}\": v for k, v in prompt_tokens.items()}\n\n            if not isinstance(chosen, str):\n                raise ValueError(f\"chosen should be an str but got {type(chosen)}\")\n            chosen_tokens = self.build_tokenized_answer(prompt, chosen)\n\n            if not isinstance(rejected, str):\n                raise ValueError(f\"rejected should be an str but got {type(rejected)}\")\n            rejected_tokens = self.build_tokenized_answer(prompt, rejected)\n\n            # Last prompt token might get merged by tokenizer and\n            # it should not be included for generation if that happens\n            prompt_len_input_ids = len(prompt_tokens[\"prompt_input_ids\"])\n\n            chosen_prompt_len_input_ids = len(chosen_tokens[\"prompt_input_ids\"])\n            rejected_prompt_len_input_ids = len(rejected_tokens[\"prompt_input_ids\"])\n            prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids)\n\n            for k, v in prompt_tokens.items():\n                prompt_tokens[k] = v[:prompt_len_input_ids]\n\n            # Make sure prompts only have one different token at most an\n            # and length only differs by 1 at most\n            num_diff_tokens = sum(\n                [a != b for a, b in zip(chosen_tokens[\"prompt_input_ids\"], rejected_tokens[\"prompt_input_ids\"])]\n            )\n            num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids)\n            if num_diff_tokens > 1 or num_diff_len > 1:\n                raise ValueError(\n                    \"Chosen and rejected prompt_input_ids might only differ on the \"\n                    \"last token due to tokenizer merge ops.\"\n                )\n\n            # add BOS token to head of prompt. Avoid adding if it's already there\n            prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed(\n                self.tokenizer.bos_token_id,\n                prompt_len_input_ids,\n                prompt_tokens,\n                chosen_prompt_len_input_ids,\n                chosen_tokens,\n                rejected_prompt_len_input_ids,\n                rejected_tokens,\n            )\n\n            # add EOS token to end of answer. Avoid adding if it's already there\n            chosen_tokens, rejected_tokens = add_eos_token_if_needed(\n                self.tokenizer.eos_token_id, chosen_tokens, rejected_tokens\n            )\n\n            longer_response_length = max(len(chosen_tokens[\"input_ids\"]), len(rejected_tokens[\"input_ids\"]))\n\n            # if combined sequence is too long, truncate the prompt\n            for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]:\n                if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > self.max_length:\n                    if self.truncation_mode == \"keep_start\":\n                        for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                            answer_tokens[k] = answer_tokens[k][: self.max_prompt_length]\n                    elif self.truncation_mode == \"keep_end\":\n                        for k in [\"prompt_input_ids\", \"prompt_attention_mask\"]:\n                            answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :]\n                    else:\n                        raise ValueError(f\"Unknown truncation mode: {self.truncation_mode}\")\n\n            # if that's still too long, truncate the response\n            for answer_tokens in [chosen_tokens, rejected_tokens]:\n                if len(answer_tokens[\"prompt_input_ids\"]) + longer_response_length > self.max_length:\n                    for k in [\"input_ids\", \"attention_mask\"]:\n                        answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length]\n\n            # Create labels\n            chosen_sequence_tokens = {\n                k: chosen_tokens[f\"prompt_{k}\"] + chosen_tokens[k] for k in [\"input_ids\", \"attention_mask\"]\n            }\n            rejected_sequence_tokens = {\n                k: rejected_tokens[f\"prompt_{k}\"] + rejected_tokens[k] for k in [\"input_ids\", \"attention_mask\"]\n            }\n            chosen_sequence_tokens[\"labels\"] = chosen_sequence_tokens[\"input_ids\"][:]\n            chosen_sequence_tokens[\"labels\"][: len(chosen_tokens[\"prompt_input_ids\"])] = [\n                self.label_pad_token_id\n            ] * len(chosen_tokens[\"prompt_input_ids\"])\n            rejected_sequence_tokens[\"labels\"] = rejected_sequence_tokens[\"input_ids\"][:]\n            rejected_sequence_tokens[\"labels\"][: len(rejected_tokens[\"prompt_input_ids\"])] = [\n                self.label_pad_token_id\n            ] * len(rejected_tokens[\"prompt_input_ids\"])\n\n            for k, toks in {\n                \"chosen_\": chosen_sequence_tokens,\n                \"rejected_\": rejected_sequence_tokens,\n                \"\": prompt_tokens,\n            }.items():\n                for type_key, tokens in toks.items():\n                    if type_key == \"token_type_ids\":\n                        continue\n                    batch[f\"{k}{type_key}\"] = tokens\n\n        else:\n            chosen_tokens = self.tokenizer(\n                chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True\n            )\n            rejected_tokens = self.tokenizer(\n                rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True\n            )\n            prompt_tokens = self.tokenizer(\n                prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True\n            )\n\n            batch[\"chosen_labels\"] = chosen_tokens[\"input_ids\"]\n            batch[\"rejected_labels\"] = rejected_tokens[\"input_ids\"]\n            batch[\"prompt_input_ids\"] = prompt_tokens[\"input_ids\"]\n            batch[\"prompt_attention_mask\"] = prompt_tokens[\"attention_mask\"]\n\n            if model is not None and hasattr(model, \"prepare_decoder_input_ids_from_labels\"):\n                batch[\"rejected_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                    labels=torch.tensor(batch[\"rejected_labels\"])\n                )\n                batch[\"chosen_decoder_input_ids\"] = model.prepare_decoder_input_ids_from_labels(\n                    labels=torch.tensor(batch[\"chosen_labels\"])\n                )\n\n        return batch\n\n    @staticmethod\n    def concatenated_inputs(\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        is_encoder_decoder: bool = False,\n        label_pad_token_id: int = -100,\n        padding_value: int = 0,\n        device: Optional[torch.device] = None,\n    ) -> Dict[str, torch.LongTensor]:\n        \"\"\"Concatenate the chosen and rejected inputs into a single tensor.\n\n        Args:\n            batch: A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors of shape (batch_size, sequence_length).\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n            label_pad_token_id: The label pad token id.\n            padding_value: The padding value to use for the concatenated inputs_ids.\n            device: The device for the concatenated inputs.\n\n        Returns:\n            A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'.\n        \"\"\"\n        concatenated_batch = {}\n\n        if is_encoder_decoder:\n            max_length = max(batch[\"chosen_labels\"].shape[1], batch[\"rejected_labels\"].shape[1])\n        else:\n            max_length = max(batch[\"chosen_input_ids\"].shape[1], batch[\"rejected_input_ids\"].shape[1])\n\n        for k in batch:\n            if k.startswith(\"chosen\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"chosen\", \"concatenated\")\n                concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value)\n        for k in batch:\n            if k.startswith(\"rejected\") and isinstance(batch[k], torch.Tensor):\n                if \"labels\" in k or is_encoder_decoder:\n                    pad_value = label_pad_token_id\n                elif k.endswith(\"_input_ids\"):\n                    pad_value = padding_value\n                elif k.endswith(\"_attention_mask\"):\n                    pad_value = 0\n                concatenated_key = k.replace(\"rejected\", \"concatenated\")\n                concatenated_batch[concatenated_key] = torch.cat(\n                    (\n                        concatenated_batch[concatenated_key],\n                        pad_to_length(batch[k], max_length, pad_value=pad_value),\n                    ),\n                    dim=0,\n                ).to(device=device)\n\n        if is_encoder_decoder:\n            concatenated_batch[\"concatenated_input_ids\"] = batch[\"prompt_input_ids\"].repeat(2, 1).to(device=device)\n            concatenated_batch[\"concatenated_attention_mask\"] = (\n                batch[\"prompt_attention_mask\"].repeat(2, 1).to(device=device)\n            )\n\n        return concatenated_batch\n\n    def cpo_loss(\n        self,\n        policy_chosen_logps: torch.FloatTensor,\n        policy_rejected_logps: torch.FloatTensor,\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Compute the CPO loss for a batch of policy and reference model log probabilities.\n\n        Args:\n            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)\n            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)\n\n        Returns:\n            A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).\n            The losses tensor contains the CPO loss for each example in the batch.\n            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.\n        \"\"\"\n        logits = (policy_chosen_logps - policy_rejected_logps).to(self.accelerator.device)\n\n        # The beta is a temperature parameter for the CPO loss, typically something in the range of 0.1 to 0.5.\n        # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and\n        # calculates a conservative CPO loss.\n\n        if self.loss_type == \"simpo\":\n            gamma_logratios = self.simpo_gamma / self.beta\n            logits = logits - gamma_logratios\n            # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0.\n            losses = (\n                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)\n                - F.logsigmoid(-self.beta * logits) * self.label_smoothing\n            )\n        elif self.loss_type == \"sigmoid\":\n            # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0.\n            losses = (\n                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)\n                - F.logsigmoid(-self.beta * logits) * self.label_smoothing\n            )\n        elif self.loss_type == \"hinge\":\n            losses = torch.relu(1 - self.beta * logits)\n        elif self.loss_type == \"ipo\":\n            # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper.\n            losses = (logits - 1 / (2 * self.beta)) ** 2\n        else:\n            raise ValueError(\n                f\"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'simpo']\"\n            )\n\n        chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach()\n        rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach()\n\n        return losses, chosen_rewards, rejected_rewards\n\n    @staticmethod\n    def get_batch_logps(\n        logits: torch.FloatTensor,\n        labels: torch.LongTensor,\n        average_log_prob: bool = False,\n        label_pad_token_id: int = -100,\n        is_encoder_decoder: bool = False,\n    ) -> torch.FloatTensor:\n        \"\"\"Compute the log probabilities of the given labels under the given logits.\n\n        Args:\n            logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size)\n            labels: Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are ignored. Shape: (batch_size, sequence_length)\n            average_log_prob: If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the log probabilities of the (non-masked) tokens.\n            label_pad_token_id: The label pad token id.\n            is_encoder_decoder: Whether the model is an encoder-decoder model.\n\n        Returns:\n            A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the given logits.\n        \"\"\"\n        if logits.shape[:-1] != labels.shape:\n            raise ValueError(\"Logits (batch and sequence length dim) and labels must have the same shape.\")\n\n        if not is_encoder_decoder:\n            labels = labels[:, 1:].clone()\n            logits = logits[:, :-1, :]\n        loss_mask = labels != label_pad_token_id\n\n        # dummy token; we'll ignore the losses on these tokens later\n        labels[labels == label_pad_token_id] = 0\n\n        per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n\n        if average_log_prob:\n            return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n        else:\n            return (per_token_logps * loss_mask).sum(-1)\n\n    def concatenated_forward(\n        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]\n    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:\n        \"\"\"Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.\n\n        We do this to avoid doing two forward passes, because it's faster for FSDP.\n        \"\"\"\n        concatenated_batch = self.concatenated_inputs(\n            batch,\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n            padding_value=self.padding_value,\n            device=self.accelerator.device,\n        )\n        len_chosen = batch[\"chosen_labels\"].shape[0]\n\n        model_kwargs = (\n            {\n                \"decoder_input_ids\": self._shift_right(concatenated_batch[\"concatenated_labels\"]),\n            }\n            if self.is_encoder_decoder\n            else {}\n        )\n\n        if self.aux_loss_enabled:\n            model_kwargs[\"output_router_logits\"] = True\n\n        outputs = model(\n            concatenated_batch[\"concatenated_input_ids\"],\n            attention_mask=concatenated_batch[\"concatenated_attention_mask\"],\n            use_cache=False,\n            **model_kwargs,\n        )\n        all_logits = outputs.logits\n\n        def cross_entropy_loss(logits, labels):\n            if not self.is_encoder_decoder:\n                # Shift so that tokens < n predict n\n                logits = logits[..., :-1, :].contiguous()\n                labels = labels[..., 1:].contiguous()\n            # Flatten the tokens\n            loss_fct = nn.CrossEntropyLoss()\n            logits = logits.view(-1, logits.shape[-1])\n            labels = labels.view(-1)\n            # Enable model parallelism\n            labels = labels.to(logits.device)\n            loss = loss_fct(logits, labels)\n            return loss\n\n        labels = concatenated_batch[\"concatenated_labels\"].clone()\n\n        if self.cpo_alpha == 0:\n            nll_loss = torch.tensor(0.0).to(self.accelerator.device)\n        else:\n            nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen])\n\n        all_logps = self.get_batch_logps(\n            all_logits,\n            concatenated_batch[\"concatenated_labels\"],\n            average_log_prob=self.loss_type in [\"ipo\", \"simpo\"],\n            is_encoder_decoder=self.is_encoder_decoder,\n            label_pad_token_id=self.label_pad_token_id,\n        )\n\n        chosen_logps = all_logps[:len_chosen]\n        rejected_logps = all_logps[len_chosen:]\n\n        chosen_logits = all_logits[:len_chosen]\n        rejected_logits = all_logits[len_chosen:]\n\n        if self.aux_loss_enabled:\n            return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss)\n\n        return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss)\n\n    def get_batch_loss_metrics(\n        self,\n        model,\n        batch: Dict[str, Union[List, torch.LongTensor]],\n        train_eval: Literal[\"train\", \"eval\"] = \"train\",\n    ):\n        \"\"\"Compute the CPO loss and other metrics for the given batch of inputs for train or test.\"\"\"\n        metrics = {}\n\n        forward_output = self.concatenated_forward(model, batch)\n        (\n            policy_chosen_logps,\n            policy_rejected_logps,\n            policy_chosen_logits,\n            policy_rejected_logits,\n            policy_nll_loss,\n        ) = forward_output[:5]\n        if self.aux_loss_enabled:\n            aux_loss = forward_output[5]\n\n        losses, chosen_rewards, rejected_rewards = self.cpo_loss(\n            policy_chosen_logps,\n            policy_rejected_logps,\n        )\n\n        loss = losses.mean() + self.cpo_alpha * policy_nll_loss\n        reward_accuracies = (chosen_rewards > rejected_rewards).float()\n\n        prefix = \"eval_\" if train_eval == \"eval\" else \"\"\n        metrics[f\"{prefix}rewards/chosen\"] = chosen_rewards.mean().cpu()\n        metrics[f\"{prefix}rewards/rejected\"] = rejected_rewards.mean().cpu()\n        metrics[f\"{prefix}rewards/accuracies\"] = reward_accuracies.mean().cpu()\n        metrics[f\"{prefix}rewards/margins\"] = (chosen_rewards - rejected_rewards).mean().cpu()\n        metrics[f\"{prefix}logps/rejected\"] = policy_rejected_logps.detach().mean().cpu()\n        metrics[f\"{prefix}logps/chosen\"] = policy_chosen_logps.detach().mean().cpu()\n        metrics[f\"{prefix}logits/rejected\"] = policy_rejected_logits.detach().mean().cpu()\n        metrics[f\"{prefix}logits/chosen\"] = policy_chosen_logits.detach().mean().cpu()\n        metrics[f\"{prefix}nll_loss\"] = policy_nll_loss.detach().mean().cpu()\n\n        if self.aux_loss_enabled:\n            loss += getattr(model.config, \"router_aux_loss_coef\", 0.0) * aux_loss\n\n        return loss, metrics\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n\n        compute_loss_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with compute_loss_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"train\")\n\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"train\")\n\n        if return_outputs:\n            return (loss, metrics)\n        return loss\n\n    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:\n        \"\"\"Generate samples from the model and reference model for the given batch of inputs.\"\"\"\n\n        # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with\n        # the torch cuda amp context manager as some hidden states are silently casted to full precision.\n        generate_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with generate_context_manager:\n            policy_output = model.generate(\n                input_ids=batch[\"prompt_input_ids\"],\n                attention_mask=batch[\"prompt_attention_mask\"],\n                max_length=self.max_length,\n                do_sample=True,\n                pad_token_id=self.tokenizer.pad_token_id,\n            )\n\n        policy_output = pad_to_length(policy_output, self.max_length, self.tokenizer.pad_token_id)\n        policy_output_decoded = self.tokenizer.batch_decode(policy_output, skip_special_tokens=True)\n\n        return policy_output_decoded\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ):\n        if not self.use_dpo_data_collator:\n            warnings.warn(\n                \"prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than \"\n                \"DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator\"\n            )\n        if ignore_keys is None:\n            if hasattr(model, \"config\"):\n                ignore_keys = getattr(model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        prediction_context_manager = amp.autocast(\"cuda\") if self._peft_has_been_casted_to_bf16 else nullcontext()\n\n        with torch.no_grad(), prediction_context_manager:\n            loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval=\"eval\")\n\n        # force log the metrics\n        self.store_metrics(metrics, train_eval=\"eval\")\n\n        if prediction_loss_only:\n            return (loss.detach(), None, None)\n\n        # logits for the chosen and rejected samples from model\n        logits_dict = {\n            \"eval_logits/chosen\": metrics[\"eval_logits/chosen\"],\n            \"eval_logits/rejected\": metrics[\"eval_logits/rejected\"],\n        }\n        logits = tuple(v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys)\n        logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)\n        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)\n\n        return (loss.detach(), logits, labels)\n\n    def store_metrics(self, metrics: Dict[str, float], train_eval: Literal[\"train\", \"eval\"] = \"train\") -> None:\n        for key, value in metrics.items():\n            self._stored_metrics[train_eval][key].append(value)\n\n    def evaluation_loop(\n        self,\n        dataloader: DataLoader,\n        description: str,\n        prediction_loss_only: Optional[bool] = None,\n        ignore_keys: Optional[List[str]] = None,\n        metric_key_prefix: str = \"eval\",\n    ) -> EvalLoopOutput:\n        \"\"\"\n        Overriding built-in evaluation loop to store metrics for each batch.\n        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.\n\n        Works both with or without labels.\n        \"\"\"\n\n        # Sample and save to game log if requested (for one batch to save time)\n        if self.generate_during_eval:\n            # Generate random indices within the range of the total number of samples\n            num_samples = len(dataloader.dataset)\n            random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size)\n\n            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader\n            random_batch_dataset = dataloader.dataset.select(random_indices)\n            random_batch = self.data_collator(random_batch_dataset)\n            random_batch = self._prepare_inputs(random_batch)\n\n            policy_output_decoded = self.get_batch_samples(self.model, random_batch)\n\n            self.log(\n                {\n                    \"game_log\": wandb.Table(\n                        columns=[\"Prompt\", \"Policy\"],\n                        rows=[\n                            [prompt, pol[len(prompt) :]]\n                            for prompt, pol in zip(random_batch[\"prompt\"], policy_output_decoded)\n                        ],\n                    )\n                }\n            )\n            self.state.log_history.pop()\n\n        # Base evaluation\n        initial_output = super().evaluation_loop(\n            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix\n        )\n\n        return initial_output\n\n    def log(self, logs: Dict[str, float]) -> None:\n        \"\"\"\n        Log `logs` on the various objects watching training, including stored metrics.\n\n        Args:\n            logs (`Dict[str, float]`):\n                The values to log.\n        \"\"\"\n        # logs either has 'loss' or 'eval_loss'\n        train_eval = \"train\" if \"loss\" in logs else \"eval\"\n        # Add averaged stored metrics to logs\n        for key, metrics in self._stored_metrics[train_eval].items():\n            logs[key] = torch.tensor(metrics).mean().item()\n        del self._stored_metrics[train_eval]\n        return super().log(logs)\n\n    def _shift_right(self, input_ids):\n        if self.decoder_start_token_id is None:\n            raise ValueError(\n                \"model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id.\"\n            )\n\n        # shift inputs to the right\n        if is_torch_fx_proxy(input_ids):\n            # Item assignment is not supported natively for proxies.\n            shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id)\n            shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)\n        else:\n            shifted_input_ids = input_ids.new_zeros(input_ids.shape)\n            shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()\n            shifted_input_ids[..., 0] = self.decoder_start_token_id\n\n        if self.pad_token_id is None:\n            raise ValueError(\"model.config.pad_token_id has to be defined.\")\n        # replace possible -100 values in labels by `pad_token_id`\n        shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id)\n\n        return shifted_input_ids\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"cpo\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport warnings\nfrom collections import defaultdict\nfrom dataclasses import FrozenInstanceError, replace\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom accelerate.utils import gather_object\nfrom datasets import Dataset\nfrom transformers import DataCollator, PreTrainedModel, PreTrainedTokenizerBase, Trainer, TrainingArguments\nfrom transformers.trainer_callback import TrainerCallback\nfrom transformers.trainer_pt_utils import nested_detach\nfrom transformers.trainer_utils import EvalPrediction\nfrom transformers.utils import is_peft_available\n\nfrom .reward_config import RewardConfig\nfrom .utils import (\n    RewardDataCollatorWithPadding,\n    compute_accuracy,\n    decode_and_strip_padding,\n    print_rich_table,\n    trl_sanitze_kwargs_for_tagging,\n)\n\n\nif is_peft_available():\n    from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training\n\n\nclass RewardTrainer(Trainer):\n    r\"\"\"\n    The RewardTrainer can be used to train your custom Reward Model. It is a subclass of the\n    `transformers.Trainer` class and inherits all of its attributes and methods. It is recommended to use\n    an `AutoModelForSequenceClassification` as the reward model. The reward model should be trained on a dataset\n    of paired examples, where each example is a tuple of two sequences. The reward model should be trained to\n    predict which example in the pair is more relevant to the task at hand.\n\n    The reward trainer expects a very specific format for the dataset. The dataset should contain two 4 entries at least\n    if you don't use the default `RewardDataCollatorWithPadding` data collator. The entries should be named\n    - `input_ids_chosen`\n    - `attention_mask_chosen`\n    - `input_ids_rejected`\n    - `attention_mask_rejected`\n\n    Optionally, you can also pass a `margin` entry to the dataset. This entry should contain the margin used to modulate the\n    loss of the reward model as outlined in https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/.\n    If you don't pass a margin, no margin will be used.\n    \"\"\"\n\n    _tag_names = [\"trl\", \"reward-trainer\"]\n\n    def __init__(\n        self,\n        model: Optional[Union[PreTrainedModel, nn.Module]] = None,\n        args: Optional[RewardConfig] = None,\n        data_collator: Optional[DataCollator] = None,\n        train_dataset: Optional[Dataset] = None,\n        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n        model_init: Optional[Callable[[], PreTrainedModel]] = None,\n        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,\n        callbacks: Optional[List[TrainerCallback]] = None,\n        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (\n            None,\n            None,\n        ),\n        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n        max_length: Optional[int] = None,\n        peft_config: Optional[Dict] = None,\n    ):\n        \"\"\"\n        Initialize RewardTrainer.\n\n        Args:\n            model (`transformers.PreTrainedModel`):\n                The model to train, preferably an `AutoModelForSequenceClassification`.\n            args (`RewardConfig`):\n                The arguments to use for training.\n            data_collator (`transformers.DataCollator`):\n                The data collator to use for training. If None is specified, the default data collator (`RewardDataCollatorWithPadding`) will be used\n                which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.\n            train_dataset (`datasets.Dataset`):\n                The dataset to use for training.\n            eval_dataset (`datasets.Dataset`):\n                The dataset to use for evaluation.\n            tokenizer (`transformers.PreTrainedTokenizerBase`):\n                The tokenizer to use for training. This argument is required if you want to use the default data collator.\n            model_init (`Callable[[], transformers.PreTrainedModel]`):\n                The model initializer to use for training. If None is specified, the default model initializer will be used.\n            compute_metrics (`Callable[[transformers.EvalPrediction], Dict]`, *optional* defaults to `compute_accuracy`):\n                The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) will be used.\n            callbacks (`List[transformers.TrainerCallback]`):\n                The callbacks to use for training.\n            optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):\n                The optimizer and scheduler to use for training.\n            preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`):\n                The function to use to preprocess the logits before computing the metrics.\n            max_length (`int`, defaults to `None`):\n                The maximum length of the sequences in the batch. This argument is required if you want to use the default data collator.\n            peft_config (`Dict`, defaults to `None`):\n                The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.\n        \"\"\"\n        if type(args) is TrainingArguments:\n            warnings.warn(\n                \"Using `transformers.TrainingArguments` for `args` is deprecated and will be removed in a future version. Please use `RewardConfig` instead.\",\n                FutureWarning,\n            )\n            if max_length is not None:\n                warnings.warn(\n                    \"The `max_length` argument is deprecated and will be removed in a future version. Please use the `RewardConfig` to set `max_length` instead.\",\n                    FutureWarning,\n                )\n        else:\n            if max_length is not None and args.max_length is not None:\n                raise ValueError(\n                    \"You cannot specify both `max_length` and `args.max_length`. Please use the `RewardConfig` to set `max_length` once.\"\n                )\n            if max_length is not None and args.max_length is None:\n                warnings.warn(\n                    \"The `max_length` argument is deprecated and will be removed in a future version. Please use the `RewardConfig` to set `max_length` instead.\",\n                    FutureWarning,\n                )\n        if not is_peft_available() and peft_config is not None:\n            raise ValueError(\n                \"PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models\"\n            )\n        elif is_peft_available() and peft_config is not None:\n            if not isinstance(model, PeftModel):\n                if getattr(model, \"is_loaded_in_8bit\", False) or getattr(model, \"is_quantized\", False):\n                    _supports_gc_kwargs = \"gradient_checkpointing_kwargs\" in list(\n                        inspect.signature(prepare_model_for_kbit_training).parameters\n                    )\n\n                    prepare_model_kwargs = {\"use_gradient_checkpointing\": args.gradient_checkpointing}\n\n                    if not _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None:\n                        warnings.warn(\n                            \"You passed `gradient_checkpointing_kwargs` in the trainer's kwargs, but your peft version does not support it. \"\n                            \"please update to the latest version of peft to use `gradient_checkpointing_kwargs`.\"\n                        )\n                    elif _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None:\n                        prepare_model_kwargs[\"gradient_checkpointing_kwargs\"] = args.gradient_checkpointing_kwargs\n\n                    model = prepare_model_for_kbit_training(model, **prepare_model_kwargs)\n\n                model = get_peft_model(model, peft_config)\n\n        if compute_metrics is None:\n            compute_metrics = compute_accuracy\n\n        if data_collator is None:\n            if tokenizer is None:\n                raise ValueError(\n                    \"max_length or a tokenizer must be specified when using the default RewardDataCollatorWithPadding\"\n                )\n            if type(args) is TrainingArguments:\n                if max_length is None:\n                    warnings.warn(\n                        \"When using RewardDataCollatorWithPadding, you should set `max_length` in RewardConfig.\"\n                        \" It will be set to `512` by default, but you should do it yourself in the future.\",\n                        UserWarning,\n                    )\n                    max_length = 512\n            else:\n                if max_length is None and args.max_length is None:\n                    warnings.warn(\n                        \"When using RewardDataCollatorWithPadding, you should set `max_length` in RewardConfig.\"\n                        \" It will be set to `512` by default, but you should do it yourself in the future.\",\n                        UserWarning,\n                    )\n                    max_length = 512\n                if max_length is None and args.max_length is not None:\n                    max_length = args.max_length\n\n            data_collator = RewardDataCollatorWithPadding(tokenizer, max_length=max_length)\n\n            if args.remove_unused_columns:\n                try:  # for bc before https://github.com/huggingface/transformers/pull/25435\n                    args.remove_unused_columns = False\n                except FrozenInstanceError:\n                    args = replace(args, remove_unused_columns=False)\n                # warn users\n                warnings.warn(\n                    \"When using RewardDataCollatorWithPadding, you should set `remove_unused_columns=False` in your RewardConfig\"\n                    \" we have set it for you, but you should do it yourself in the future.\",\n                    UserWarning,\n                )\n\n            self.use_reward_data_collator = True\n        else:\n            self.use_reward_data_collator = False\n        super().__init__(\n            model=model,\n            args=args,\n            data_collator=data_collator,\n            train_dataset=train_dataset,\n            eval_dataset=eval_dataset,\n            tokenizer=tokenizer,\n            model_init=model_init,\n            compute_metrics=compute_metrics,\n            callbacks=callbacks,\n            optimizers=optimizers,\n            preprocess_logits_for_metrics=preprocess_logits_for_metrics,\n        )\n\n        # Add tags for models that have been loaded with the correct transformers version\n        if hasattr(self.model, \"add_model_tags\"):\n            self.model.add_model_tags(self._tag_names)\n\n    def compute_loss(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        return_outputs=False,\n    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:\n        if not self.use_reward_data_collator:\n            warnings.warn(\n                \"The current compute_loss is implemented for RewardDataCollatorWithPadding,\"\n                \" if you are using a custom data collator make sure you know what you are doing or\"\n                \" implement your own compute_loss method.\"\n            )\n        rewards_chosen = model(\n            input_ids=inputs[\"input_ids_chosen\"],\n            attention_mask=inputs[\"attention_mask_chosen\"],\n            return_dict=True,\n        )[\"logits\"]\n        rewards_rejected = model(\n            input_ids=inputs[\"input_ids_rejected\"],\n            attention_mask=inputs[\"attention_mask_rejected\"],\n            return_dict=True,\n        )[\"logits\"]\n        # calculate loss, optionally modulate with margin\n        if \"margin\" in inputs:\n            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs[\"margin\"]).mean()\n        else:\n            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()\n\n        if self.args.center_rewards_coefficient is not None:\n            loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2)\n\n        if return_outputs:\n            return loss, {\n                \"rewards_chosen\": rewards_chosen,\n                \"rewards_rejected\": rewards_rejected,\n            }\n        return loss\n\n    def prediction_step(\n        self,\n        model: Union[PreTrainedModel, nn.Module],\n        inputs: Dict[str, Union[torch.Tensor, Any]],\n        prediction_loss_only: bool,\n        ignore_keys: Optional[List[str]] = None,\n    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:\n        inputs = self._prepare_inputs(inputs)\n        if ignore_keys is None:\n            if hasattr(self.model, \"config\"):\n                ignore_keys = getattr(self.model.config, \"keys_to_ignore_at_inference\", [])\n            else:\n                ignore_keys = []\n\n        with torch.no_grad():\n            loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True)\n\n        if prediction_loss_only:\n            return (loss, None, None)\n\n        loss = loss.detach()\n        logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys)\n        logits = nested_detach(logits)\n        # Stack accepted against rejected, mean over logits\n        # and softmax to get preferences between accepted and rejected to sum to 1\n        logits = torch.stack(logits).mean(dim=2).softmax(dim=0).T\n\n        labels = torch.zeros(logits.shape[0])\n        labels = self._prepare_inputs(labels)\n\n        return loss, logits, labels\n\n    def evaluate(self, *args, **kwargs):\n        num_print_samples = kwargs.pop(\"num_print_samples\", 4)\n        self.visualize_samples(num_print_samples)\n        return super().evaluate(*args, **kwargs)\n\n    def visualize_samples(self, num_print_samples: int):\n        \"\"\"\n        Visualize the reward model logits prediction\n\n        Args:\n            num_print_samples (`int`, defaults to `4`):\n                The number of samples to print. Set to `-1` to print all samples.\n        \"\"\"\n        eval_dataloader = self.get_eval_dataloader()\n        table = defaultdict(list)\n        for _, inputs in enumerate(eval_dataloader):\n            _, logits, _ = self.prediction_step(self.model, inputs, prediction_loss_only=False)\n            chosen_text = decode_and_strip_padding(inputs[\"input_ids_chosen\"], self.tokenizer)\n            rejected_text = decode_and_strip_padding(inputs[\"input_ids_rejected\"], self.tokenizer)\n            table[\"chosen_text\"].extend(gather_object(chosen_text))\n            table[\"rejected_text\"].extend(gather_object(rejected_text))\n            table[\"logits\"].extend(\n                gather_object([[round(inner_item, 4) for inner_item in item] for item in logits.tolist()])\n            )\n            if num_print_samples >= 0 and len(table[\"chosen_text\"]) >= num_print_samples:\n                break\n        df = pd.DataFrame(table)\n        if self.accelerator.process_index == 0:\n            print_rich_table(df[:num_print_samples])\n            if \"wandb\" in self.args.report_to:\n                import wandb\n\n                if wandb.run is not None:\n                    wandb.log({\"completions\": wandb.Table(dataframe=df)})\n\n    @wraps(Trainer.push_to_hub)\n    def push_to_hub(\n        self,\n        commit_message: Optional[str] = \"End of training\",\n        blocking: bool = True,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Overwrite the `push_to_hub` method in order to force-add the tag \"reward-trainer\" when pushing the\n        model on the Hub. Please refer to `~transformers.Trainer.push_to_hub` for more details.\n        Unlike the parent class, we don't use the `token` argument to mitigate security risks.\n        \"\"\"\n        kwargs = trl_sanitze_kwargs_for_tagging(model=self.model, tag_names=self._tag_names, kwargs=kwargs)\n        return super().push_to_hub(commit_message=commit_message, blocking=blocking, **kwargs)\n\n\n# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Literal, Optional\n\nfrom transformers import TrainingArguments\n\n\n@dataclass\nclass CPOConfig(TrainingArguments):\n    r\"\"\"\n    Configuration class for the [`CPOTrainer`].\n\n    Using [`~transformers.HfArgumentParser`] we can turn this class into\n    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the\n    command line.\n\n    Parameters:\n        max_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want\n            to use the default data collator.\n        max_prompt_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the prompt. This argument is required if you want to use the default data collator.\n        max_completion_length (`Optional[int]`, *optional*, defaults to `None`):\n            Maximum length of the completion. This argument is required if you want to use the default data collator\n            and your model is an encoder-decoder.\n        beta (`float`, *optional*, defaults to `0.1`):\n            Parameter controlling the deviation from the reference model. Higher β means less deviation from the\n            reference model. For the IPO loss (`loss_type=\"ipo\"`), β is the regularization parameter denoted by τ in\n            the [paper](https://huggingface.co/papers/2310.12036).\n        label_smoothing (`float`, *optional*, defaults to `0.0`):\n            Label smoothing factor. This argument is required if you want to use the default data collator.\n        loss_type (`str`, *optional*, defaults to `\"sigmoid\"`):\n            Type of loss to use. Possible values are:\n\n                - `\"sigmoid\"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper.\n                - `\"hinge\"`: hinge loss on the normalized likelihood from the [SLiC](https://huggingface.co/papers/2305.10425) paper.\n                - `\"ipo\"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper.\n                - `\"simpo\"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper.\n\n        disable_dropout (`bool`, *optional*, defaults to `True`):\n            Whether to disable dropout in the model.\n        cpo_alpha (`float`, *optional*, defaults to `1.0`):\n            Weight of the BC regularizer in CPO training.\n        simpo_gamma (`float`, *optional*, defaults to `0.5`):\n            Target reward margin for the SimPO loss, used only when the `loss_type=\"simpo\"`.\n        label_pad_token_id (`int`, *optional*, defaults to `-100`):\n            Label pad token id. This argument is required if you want to use the default data collator.\n        padding_value (`Optional[int]`, *optional*, defaults to `None`):\n            Padding value to use. If `None`, the padding value of the tokenizer is used.\n        truncation_mode (`str`,*optional*,  defaults to `\"keep_end\"`):\n            Truncation mode to use when the prompt is too long. Possible values are `\"keep_end\"` or `\"keep_start\"`.\n            This argument is required if you want to use the default data collator.\n        generate_during_eval (`bool`, *optional*, defaults to `False`):\n            If `True`, generates and logs completions from the model to W&B during evaluation.\n        is_encoder_decoder (`Optional[bool]`, *optional*, defaults to `None`):\n            When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument,\n            you need to specify if the model returned by the callable is an encoder-decoder model.\n        model_init_kwargs (`Optional[Dict[str, Any]]`, *optional*, defaults to `None`):\n            Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a\n            string.\n        dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):\n            Number of processes to use for processing the dataset.\n    \"\"\"\n\n    max_length: Optional[int] = None\n    max_prompt_length: Optional[int] = None\n    max_completion_length: Optional[int] = None\n    beta: float = 0.1\n    label_smoothing: float = 0.0\n    loss_type: Literal[\"sigmoid\", \"hinge\", \"ipo\", \"simpo\"] = \"sigmoid\"\n    disable_dropout: bool = True\n    cpo_alpha: float = 1.0\n    simpo_gamma: float = 0.5\n    label_pad_token_id: int = -100\n    padding_value: Optional[int] = None\n    truncation_mode: str = \"keep_end\"\n    generate_during_eval: bool = False\n    is_encoder_decoder: Optional[bool] = None\n    model_init_kwargs: Optional[Dict[str, Any]] = None\n    dataset_num_proc: Optional[int] = None","difficulty":"hard","domain":"Code Repository Understanding","length":"long","question":"In the May 20, 2023 commit of the trl repository, a new PPOTrainer class was introduced for Proximal Policy Optimization (PPO) training. This version significantly improved the trainer by introducing parallel training support using the Accelerate library. If you want to ensure that the policy gradient of the actor-critic model is not reset across epochs (preserving gradients between epochs) when enabling the Accelerator.split_between_epochs feature, what code modifications are required? Can you explain the code changes needed in ppo_trainer.py and accelerator.py (or related auxiliary files), particularly focusing on how to handle dependencies between the forward pass and backward pass?","sub_domain":"Code repo QA"}

Source: https://huggingface.co/datasets/zai-org/LongBench-v2

initial import

Posting: /agents

GET /api/v1/write?intent=publish&task_id=588a6997-9931-5318-9176-8e639f0ecb9f&body={url_encoded_text}&agent_name={optional_name}&nonce={optional_random_id}