We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
StackLLAMA examples uses sentiment model as reward model
The script to train the StackLLAMA model uses a sentiment pipeline for training the reward model while it should use the model that was trained before. 

https://github.com/lvwerra/trl/blob/main/examples/stack_llama/scripts/rl_training.py#L216

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/research_projects/stack_llama/scripts/rl_training.py`
Content:
```
1 # coding=utf-8
2 # Copyright 2022 The HuggingFace Inc. team. All rights reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 from dataclasses import dataclass, field
16 from typing import Optional
17 
18 import torch
19 from accelerate import Accelerator
20 from datasets import load_dataset
21 from peft import LoraConfig
22 from tqdm import tqdm
23 from transformers import Adafactor, AutoTokenizer, HfArgumentParser, pipeline
24 
25 from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
26 from trl.core import LengthSampler
27 
28 
29 tqdm.pandas()
30 
31 
32 @dataclass
33 class ScriptArguments:
34     """
35     The name of the Casual LM model we wish to fine with PPO
36     """
37 
38     # NOTE: gpt2 models use Conv1D instead of Linear layers which are not yet supported in 8 bit mode
39     # models like gpt-neo* models are more suitable.
40     model_name: Optional[str] = field(default="", metadata={"help": "the model name"})
41     tokenizer_name: Optional[str] = field(default="", metadata={"help": "the tokenizer name"})
42     reward_model_name: Optional[str] = field(default="", metadata={"help": "the reward model name"})
43     log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
44     learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
45     output_max_length: Optional[int] = field(default=128, metadata={"help": "maximum length for generation"})
46     mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
47     batch_size: Optional[int] = field(default=32, metadata={"help": "the batch size"})
48     ppo_epochs: Optional[int] = field(default=4, metadata={"help": "the number of ppo epochs"})
49     gradient_accumulation_steps: Optional[int] = field(
50         default=4, metadata={"help": "the number of gradient accumulation steps"}
51     )
52     adafactor: Optional[bool] = field(default=False, metadata={"help": "whether to use the adafactor optimizer"})
53     early_stopping: Optional[bool] = field(default=False, metadata={"help": "whether to early stop"})
54     target_kl: Optional[float] = field(default=0.1, metadata={"help": "kl target for early stopping"})
55     reward_baseline: Optional[float] = field(
56         default=0.0,
57         metadata={"help": "a baseline value that is subtracted from the reward"},
58     )
59     batched_gen: Optional[bool] = field(default=False, metadata={"help": "whether to use the batched text gen"})
60     save_freq: Optional[int] = field(default=None, metadata={"help": "n steps to save the model"})
61     output_dir: Optional[str] = field(default="runs/", metadata={"help": "n steps to save the model"})
62     seed: Optional[int] = field(default=0, metadata={"help": "the seed"})
63     steps: Optional[int] = field(default=20000, metadata={"help": "number of epochs"})
64     init_kl_coef: Optional[float] = field(
65         default=0.2,
66         metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},
67     )
68 
69     adap_kl_ctrl: Optional[bool] = field(default=True, metadata={"help": "Use adaptive KL control, otherwise linear"})
70 
71 
72 parser = HfArgumentParser(ScriptArguments)
73 script_args: ScriptArguments = parser.parse_args_into_dataclasses()[0]
74 reward_model_name = script_args.reward_model_name
75 dataset_name = "lvwerra/stack-exchange-paired"
76 config = PPOConfig(
77     steps=script_args.steps,
78     model_name=script_args.model_name,
79     learning_rate=script_args.learning_rate,
80     log_with=script_args.log_with,
81     batch_size=script_args.batch_size,
82     mini_batch_size=script_args.mini_batch_size,
83     gradient_accumulation_steps=script_args.gradient_accumulation_steps,
84     optimize_cuda_cache=True,
85     early_stopping=script_args.early_stopping,
86     target_kl=script_args.target_kl,
87     ppo_epochs=script_args.ppo_epochs,
88     seed=script_args.seed,
89     init_kl_coef=script_args.init_kl_coef,
90     adap_kl_ctrl=script_args.adap_kl_ctrl,
91 )
92 
93 train_dataset = load_dataset("lvwerra/stack-exchange-paired", data_dir="data/rl", split="train")
94 train_dataset = train_dataset.select(range(100000))
95 # We then define the arguments to pass to the sentiment analysis pipeline.
96 # We set `return_all_scores` to True to get the sentiment score for each token.
97 sent_kwargs = {
98     "return_all_scores": True,
99     "function_to_apply": "none",
100     "batch_size": 16,
101     "truncation": True,
102 }
103 
104 tokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name)
105 # GPT-2 tokenizer has a pad token, but it is not eos_token by default. We need to set it to eos_token.
106 # only for this model.
107 
108 if getattr(tokenizer, "pad_token", None) is None:
109     tokenizer.pad_token = tokenizer.eos_token
110 
111 
112 # Below is an example function to build the dataset. In our case, we use the IMDB dataset
113 # from the `datasets` library. One should customize this function to train the model on
114 # its own dataset.
115 def build_dataset(
116     tokenizer,
117     dataset_name="lvwerra/stack-exchange-paired",
118 ):
119     """
120     Build dataset for training. This builds the dataset from `load_dataset`, one should
121     customize this function to train the model on its own dataset.
122 
123     Args:
124         dataset_name (`str`):
125             The name of the dataset to be loaded.
126 
127     Returns:
128         dataloader (`torch.utils.data.DataLoader`):
129             The dataloader for the dataset.
130     """
131 
132     # load imdb with datasets
133     ds = load_dataset(dataset_name, data_dir="data/rl", split="train")
134     original_columns = ds.column_names
135     num_proc = 24
136 
137     def preprocess_function(examples):
138         new_examples = {
139             "query": [],
140             "input_ids": [],
141         }
142         for question in examples["question"]:
143             query = "Question: " + question + "\n\nAnswer: "
144             tokenized_question = tokenizer(query, truncation=True)
145             new_examples["query"].append(query)
146             new_examples["input_ids"].append(tokenized_question["input_ids"])
147 
148         return new_examples
149 
150     ds = train_dataset.map(
151         preprocess_function,
152         batched=True,
153         num_proc=num_proc,
154         remove_columns=original_columns,
155     )
156     ds = ds.filter(lambda x: len(x["input_ids"]) < 512, batched=False)
157 
158     ds.set_format(type="torch")
159     return ds
160 
161 
162 # We retrieve the dataloader by calling the `build_dataset` function.
163 dataset = build_dataset(tokenizer)
164 
165 
166 def collator(data):
167     return dict((key, [d[key] for d in data]) for key in data[0])
168 
169 
170 # set seed before initializing value head for deterministic eval
171 set_seed(config.seed)
172 
173 # Now let's build the model, the reference model, and the tokenizer.
174 current_device = Accelerator().local_process_index
175 
176 lora_config = LoraConfig(
177     r=16,
178     lora_alpha=32,
179     lora_dropout=0.05,
180     bias="none",
181     task_type="CAUSAL_LM",
182 )
183 model = AutoModelForCausalLMWithValueHead.from_pretrained(
184     config.model_name,
185     load_in_8bit=True,
186     device_map={"": current_device},
187     peft_config=lora_config,
188 )
189 
190 optimizer = None
191 if script_args.adafactor:
192     optimizer = Adafactor(
193         filter(lambda p: p.requires_grad, model.parameters()),
194         scale_parameter=False,
195         relative_step=False,
196         warmup_init=False,
197         lr=config.learning_rate,
198     )
199 # We then build the PPOTrainer, passing the model, the reference model, the tokenizer
200 ppo_trainer = PPOTrainer(
201     config,
202     model,
203     ref_model=None,
204     tokenizer=tokenizer,
205     dataset=dataset,
206     data_collator=collator,
207     optimizer=optimizer,
208 )
209 
210 # We then build the sentiment analysis pipeline, passing the model name and the
211 # sentiment analysis pipeline arguments. Let's also make sure to set the device
212 # to the same device as the PPOTrainer.
213 device = ppo_trainer.accelerator.device
214 if ppo_trainer.accelerator.num_processes == 1:
215     device = 0 if torch.cuda.is_available() else "cpu"  # to avoid a ` pipeline` bug
216 sentiment_pipe = pipeline(
217     "sentiment-analysis",
218     model=reward_model_name,
219     device_map={"": current_device},
220     model_kwargs={"load_in_8bit": True},
221     tokenizer=tokenizer,
222     return_token_type_ids=False,
223 )
224 
225 # We then define the arguments to pass to the `generate` function. These arguments
226 # are passed to the `generate` function of the PPOTrainer, which is a wrapper around
227 # the `generate` function of the trained model.
228 generation_kwargs = {
229     # "min_length": -1,
230     "top_k": 0.0,
231     "top_p": 1.0,
232     "do_sample": True,
233     "pad_token_id": tokenizer.pad_token_id,
234     "eos_token_id": 100_000,
235 }
236 output_min_length = 32
237 output_max_length = script_args.output_max_length
238 output_length_sampler = LengthSampler(output_min_length, output_max_length)
239 
240 for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
241     if epoch >= config.total_ppo_epochs:
242         break
243 
244     question_tensors = batch["input_ids"]
245 
246     response_tensors = ppo_trainer.generate(
247         question_tensors,
248         return_prompt=False,
249         length_sampler=output_length_sampler,
250         **generation_kwargs,
251     )
252     batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)
253 
254     # Compute sentiment score
255     texts = [q + r for q, r in zip(batch["query"], batch["response"])]
256     pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
257     rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]
258 
259     # Run PPO step
260     stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
261     ppo_trainer.log_stats(stats, batch, rewards)
262 
263     if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
264         ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
265
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/research_projects/stack_llama/scripts/rl_training.py b/examples/research_projects/stack_llama/scripts/rl_training.py
--- a/examples/research_projects/stack_llama/scripts/rl_training.py
+++ b/examples/research_projects/stack_llama/scripts/rl_training.py
@@ -92,6 +92,7 @@
 
 train_dataset = load_dataset("lvwerra/stack-exchange-paired", data_dir="data/rl", split="train")
 train_dataset = train_dataset.select(range(100000))
+
 # We then define the arguments to pass to the sentiment analysis pipeline.
 # We set `return_all_scores` to True to get the sentiment score for each token.
 sent_kwargs = {
@@ -207,9 +208,9 @@
     optimizer=optimizer,
 )
 
-# We then build the sentiment analysis pipeline, passing the model name and the
-# sentiment analysis pipeline arguments. Let's also make sure to set the device
-# to the same device as the PPOTrainer.
+# We then build the sentiment analysis pipeline using our reward model, passing the
+# model name and the sentiment analysis pipeline arguments. Let's also make sure to
+# set the device to the same device as the PPOTrainer.
 device = ppo_trainer.accelerator.device
 if ppo_trainer.accelerator.num_processes == 1:
     device = 0 if torch.cuda.is_available() else "cpu"  # to avoid a ` pipeline` bug
@@ -251,7 +252,7 @@
     )
     batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)
 
-    # Compute sentiment score
+    # Compute reward score (using the sentiment analysis pipeline)
     texts = [q + r for q, r in zip(batch["query"], batch["response"])]
     pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
     rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

{"golden_diff": "diff --git a/examples/research_projects/stack_llama/scripts/rl_training.py b/examples/research_projects/stack_llama/scripts/rl_training.py\n--- a/examples/research_projects/stack_llama/scripts/rl_training.py\n+++ b/examples/research_projects/stack_llama/scripts/rl_training.py\n@@ -92,6 +92,7 @@\n \n train_dataset = load_dataset(\"lvwerra/stack-exchange-paired\", data_dir=\"data/rl\", split=\"train\")\n train_dataset = train_dataset.select(range(100000))\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.\n sent_kwargs = {\n@@ -207,9 +208,9 @@\n     optimizer=optimizer,\n )\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.\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.\n device = ppo_trainer.accelerator.device\n if ppo_trainer.accelerator.num_processes == 1:\n     device = 0 if torch.cuda.is_available() else \"cpu\"  # to avoid a ` pipeline` bug\n@@ -251,7 +252,7 @@\n     )\n     batch[\"response\"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)\n \n-    # Compute sentiment score\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", "issue": "StackLLAMA examples uses sentiment model as reward model\nThe script to train the StackLLAMA model uses a sentiment pipeline for training the reward model while it should use the model that was trained before. \r\n\r\nhttps://github.com/lvwerra/trl/blob/main/examples/stack_llama/scripts/rl_training.py#L216\n", "before_files": [{"content": "# coding=utf-8\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 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\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(\"lvwerra/stack-exchange-paired\", data_dir=\"data/rl\", split=\"train\")\ntrain_dataset = train_dataset.select(range(100000))\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    # load imdb with datasets\n    ds = load_dataset(dataset_name, data_dir=\"data/rl\", split=\"train\")\n    original_columns = ds.column_names\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)\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 dict((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=True,\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, 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    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\": True},\n    tokenizer=tokenizer,\n    return_token_type_ids=False,\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.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 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[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", "path": "examples/research_projects/stack_llama/scripts/rl_training.py"}], "after_files": [{"content": "# coding=utf-8\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 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\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(\"lvwerra/stack-exchange-paired\", data_dir=\"data/rl\", split=\"train\")\ntrain_dataset = train_dataset.select(range(100000))\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    # load imdb with datasets\n    ds = load_dataset(dataset_name, data_dir=\"data/rl\", split=\"train\")\n    original_columns = ds.column_names\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)\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 dict((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=True,\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\": True},\n    tokenizer=tokenizer,\n    return_token_type_ids=False,\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.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", "path": "examples/research_projects/stack_llama/scripts/rl_training.py"}]}