We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Potential bug of model output dim in reward trainer example
Hi!

I've been exploring the `examples/scripts/reward_trainer.py` and believe I've found a potential bug that I'd like to address.

## Potential bug.
To my understanding, the reward model should produce a scalar output. However, when we load the pre-trained model in the [example script](https://github.com/lvwerra/trl/blob/main/examples/scripts/reward_trainer.py#L75-L80), the model will be a binary classifier by default, with a 2-dim logit output:

```python
import torch
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
model = model.cpu()
with torch.no_grad():
    out = model(**inputs)
print(out.logits.shape) # torch.Size([1, 2])
```

During training, this issue was hidden by the [the loss computation part](https://github.com/lvwerra/trl/blob/main/trl/trainer/reward_trainer.py#L169-L176) of the reward trainer:

```python
    rewards_chosen = model(
        input_ids=inputs["input_ids_chosen"],
        attention_mask=inputs["attention_mask_chosen"],
    )[0]

    rewards_rejected = model(
        input_ids=inputs["input_ids_rejected"],
        attention_mask=inputs["attention_mask_rejected"],
    )[0]

    loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
```
There, the `rewards_rejected` and `rewards_chosen` are dim [sample, 2] matrices, which I think should be dim [sample, 1] instead. The `.mean()` operator hides the issue as it returns a scalar value.

Here is a [colab](https://colab.research.google.com/drive/1szzGC8H1DQ7329_idVnnOnPIPE6D7J4j?usp=sharing) that shows the dimensionality problem.

## Proposed fix

To resolve this, in the model training example, we can use:

```python
model = AutoModelForSequenceClassification.from_pretrained(
    script_args.model_name,
    quantization_config=quantization_config,
    device_map=device_map,
    trust_remote_code=script_args.trust_remote_code,
    num_labels=1 # <---- make sure that the model output is a scalar
)
```
I'll submit a pull request if it is OK. 

In addition, while exploring `examples/scripts/reward_trainer.py`, I noticed a few other places that can be improved. For example, there is currently no validation steps that show the genearlization performance of the reward model; the `log_with` option in `ScriptArguments` is currently not being used anywhere and therefore has no effect. I'm happy to include the validation and fix the logging glitch in my pull request as well.

Tianlin

--- 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/scripts/reward_trainer.py`
Content:
```
1 # coding=utf-8
2 # Copyright 2023 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 from datasets import load_dataset
19 from peft import LoraConfig
20 from tqdm import tqdm
21 from transformers import (
22     AutoModelForSequenceClassification,
23     AutoTokenizer,
24     BitsAndBytesConfig,
25     HfArgumentParser,
26     TrainingArguments,
27 )
28 
29 from trl import RewardTrainer
30 
31 
32 tqdm.pandas()
33 
34 
35 # Define and parse arguments.
36 @dataclass
37 class ScriptArguments:
38     """
39     The name of the Casual LM model we wish to fine with SFTTrainer
40     """
41 
42     model_name: Optional[str] = field(default="facebook/opt-350m", metadata={"help": "the model name"})
43     dataset_name: Optional[str] = field(default="Anthropic/hh-rlhf", metadata={"help": "the model name"})
44     dataset_text_field: Optional[str] = field(default="text", metadata={"help": "the text field of the dataset"})
45     log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
46     learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
47     batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
48     seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
49     gradient_accumulation_steps: Optional[int] = field(
50         default=16, metadata={"help": "the number of gradient accumulation steps"}
51     )
52     load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
53     load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
54     use_peft: Optional[bool] = field(default=False, metadata={"help": "Wether to use PEFT or not to train adapters"})
55     trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
56     output_dir: Optional[str] = field(default="output", metadata={"help": "the output directory"})
57 
58 
59 parser = HfArgumentParser(ScriptArguments)
60 script_args = parser.parse_args_into_dataclasses()[0]
61 
62 # Step 1: Load the model
63 if script_args.load_in_8bit and script_args.load_in_4bit:
64     raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
65 elif script_args.load_in_8bit or script_args.load_in_4bit:
66     quantization_config = BitsAndBytesConfig(
67         load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit
68     )
69     # This means: fit the entire model on the GPU:0
70     device_map = {"": 0}
71 else:
72     device_map = None
73     quantization_config = None
74 
75 model = AutoModelForSequenceClassification.from_pretrained(
76     script_args.model_name,
77     quantization_config=quantization_config,
78     device_map=device_map,
79     trust_remote_code=script_args.trust_remote_code,
80 )
81 
82 # Step 2: Load the dataset and pre-process it
83 tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)
84 dataset = load_dataset(script_args.dataset_name, split="train")
85 # Turn the dataset into pairs of post + summaries, where text_j is the preferred question + answer and text_k is the other.
86 # Then tokenize the dataset.
87 # Adapt this section to your needs for custom datasets
88 
89 
90 def preprocess_function(examples):
91     new_examples = {
92         "input_ids_chosen": [],
93         "attention_mask_chosen": [],
94         "input_ids_rejected": [],
95         "attention_mask_rejected": [],
96     }
97     for chosen, rejected in zip(examples["chosen"], examples["rejected"]):
98         tokenized_j = tokenizer(chosen, truncation=True)
99         tokenized_k = tokenizer(rejected, truncation=True)
100 
101         new_examples["input_ids_chosen"].append(tokenized_j["input_ids"])
102         new_examples["attention_mask_chosen"].append(tokenized_j["attention_mask"])
103         new_examples["input_ids_rejected"].append(tokenized_k["input_ids"])
104         new_examples["attention_mask_rejected"].append(tokenized_k["attention_mask"])
105 
106     return new_examples
107 
108 
109 # preprocess the dataset and filter out QAs that are longer than script_args.max_length
110 original_columns = dataset.column_names
111 train_dataset = dataset.map(
112     preprocess_function,
113     batched=True,
114     num_proc=4,
115 )
116 train_dataset = train_dataset.filter(
117     lambda x: len(x["input_ids_chosen"]) <= script_args.seq_length
118     and len(x["input_ids_rejected"]) <= script_args.seq_length
119 )
120 
121 
122 # Step 3: Define the training arguments
123 training_args = TrainingArguments(
124     output_dir=script_args.output_dir,
125     per_device_train_batch_size=script_args.batch_size,
126     gradient_accumulation_steps=script_args.gradient_accumulation_steps,
127     learning_rate=script_args.learning_rate,
128 )
129 
130 # Step 4: Define the LoraConfig
131 if script_args.use_peft:
132     peft_config = LoraConfig(r=16, lora_alpha=16, bias="none", task_type="SEQ_CLS", modules_to_save=["scores"])
133 else:
134     peft_config = None
135 
136 # Step 5: Define the Trainer
137 trainer = RewardTrainer(
138     model=model,
139     tokenizer=tokenizer,
140     args=training_args,
141     train_dataset=train_dataset,
142     peft_config=peft_config,
143     max_length=script_args.seq_length,
144 )
145 
146 trainer.train()
147
```
--- 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/scripts/reward_trainer.py b/examples/scripts/reward_trainer.py
--- a/examples/scripts/reward_trainer.py
+++ b/examples/scripts/reward_trainer.py
@@ -43,8 +43,13 @@
     dataset_name: Optional[str] = field(default="Anthropic/hh-rlhf", metadata={"help": "the model name"})
     dataset_text_field: Optional[str] = field(default="text", metadata={"help": "the text field of the dataset"})
     log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
+    logging_steps: Optional[int] = field(default=500, metadata={"help": "the number of update steps between two logs"})
+    eval_split: Optional[str] = field(
+        default="none", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
+    )
     learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
     batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
+    num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
     seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
     gradient_accumulation_steps: Optional[int] = field(
         default=16, metadata={"help": "the number of gradient accumulation steps"}
@@ -77,11 +82,14 @@
     quantization_config=quantization_config,
     device_map=device_map,
     trust_remote_code=script_args.trust_remote_code,
+    num_labels=1,
 )
 
 # Step 2: Load the dataset and pre-process it
 tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)
-dataset = load_dataset(script_args.dataset_name, split="train")
+train_dataset = load_dataset(script_args.dataset_name, split="train")
+
+
 # Turn the dataset into pairs of post + summaries, where text_j is the preferred question + answer and text_k is the other.
 # Then tokenize the dataset.
 # Adapt this section to your needs for custom datasets
@@ -107,8 +115,7 @@
 
 
 # preprocess the dataset and filter out QAs that are longer than script_args.max_length
-original_columns = dataset.column_names
-train_dataset = dataset.map(
+train_dataset = train_dataset.map(
     preprocess_function,
     batched=True,
     num_proc=4,
@@ -118,13 +125,34 @@
     and len(x["input_ids_rejected"]) <= script_args.seq_length
 )
 
+if script_args.eval_split == "none":
+    eval_dataset = None
+else:
+    eval_dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)
+
+    eval_dataset = eval_dataset.map(
+        preprocess_function,
+        batched=True,
+        num_proc=4,
+    )
+    eval_dataset = eval_dataset.filter(
+        lambda x: len(x["input_ids_chosen"]) <= script_args.seq_length
+        and len(x["input_ids_rejected"]) <= script_args.seq_length
+    )
+
 
 # Step 3: Define the training arguments
 training_args = TrainingArguments(
     output_dir=script_args.output_dir,
     per_device_train_batch_size=script_args.batch_size,
+    num_train_epochs=script_args.num_train_epochs,
     gradient_accumulation_steps=script_args.gradient_accumulation_steps,
     learning_rate=script_args.learning_rate,
+    report_to="wandb" if script_args.log_with == "wandb" else "tensorboard",
+    remove_unused_columns=False,
+    optim="adamw_torch",
+    logging_steps=script_args.logging_steps,
+    evaluation_strategy="steps" if script_args.eval_split != "none" else "no",
 )
 
 # Step 4: Define the LoraConfig
@@ -139,6 +167,7 @@
     tokenizer=tokenizer,
     args=training_args,
     train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
     peft_config=peft_config,
     max_length=script_args.seq_length,
 )

{"golden_diff": "diff --git a/examples/scripts/reward_trainer.py b/examples/scripts/reward_trainer.py\n--- a/examples/scripts/reward_trainer.py\n+++ b/examples/scripts/reward_trainer.py\n@@ -43,8 +43,13 @@\n     dataset_name: Optional[str] = field(default=\"Anthropic/hh-rlhf\", metadata={\"help\": \"the model name\"})\n     dataset_text_field: Optional[str] = field(default=\"text\", metadata={\"help\": \"the text field of the dataset\"})\n     log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n+    logging_steps: Optional[int] = field(default=500, metadata={\"help\": \"the number of update steps between two logs\"})\n+    eval_split: Optional[str] = field(\n+        default=\"none\", metadata={\"help\": \"the dataset split to evaluate on; default to 'none' (no evaluation)\"}\n+    )\n     learning_rate: Optional[float] = field(default=1.41e-5, metadata={\"help\": \"the learning rate\"})\n     batch_size: Optional[int] = field(default=64, metadata={\"help\": \"the batch size\"})\n+    num_train_epochs: Optional[int] = field(default=1, metadata={\"help\": \"the number of training epochs\"})\n     seq_length: Optional[int] = field(default=512, metadata={\"help\": \"Input sequence length\"})\n     gradient_accumulation_steps: Optional[int] = field(\n         default=16, metadata={\"help\": \"the number of gradient accumulation steps\"}\n@@ -77,11 +82,14 @@\n     quantization_config=quantization_config,\n     device_map=device_map,\n     trust_remote_code=script_args.trust_remote_code,\n+    num_labels=1,\n )\n \n # Step 2: Load the dataset and pre-process it\n tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)\n-dataset = load_dataset(script_args.dataset_name, split=\"train\")\n+train_dataset = load_dataset(script_args.dataset_name, split=\"train\")\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.\n # Adapt this section to your needs for custom datasets\n@@ -107,8 +115,7 @@\n \n \n # preprocess the dataset and filter out QAs that are longer than script_args.max_length\n-original_columns = dataset.column_names\n-train_dataset = dataset.map(\n+train_dataset = train_dataset.map(\n     preprocess_function,\n     batched=True,\n     num_proc=4,\n@@ -118,13 +125,34 @@\n     and len(x[\"input_ids_rejected\"]) <= script_args.seq_length\n )\n \n+if script_args.eval_split == \"none\":\n+    eval_dataset = None\n+else:\n+    eval_dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)\n+\n+    eval_dataset = eval_dataset.map(\n+        preprocess_function,\n+        batched=True,\n+        num_proc=4,\n+    )\n+    eval_dataset = eval_dataset.filter(\n+        lambda x: len(x[\"input_ids_chosen\"]) <= script_args.seq_length\n+        and len(x[\"input_ids_rejected\"]) <= script_args.seq_length\n+    )\n+\n \n # Step 3: Define the training arguments\n training_args = TrainingArguments(\n     output_dir=script_args.output_dir,\n     per_device_train_batch_size=script_args.batch_size,\n+    num_train_epochs=script_args.num_train_epochs,\n     gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n     learning_rate=script_args.learning_rate,\n+    report_to=\"wandb\" if script_args.log_with == \"wandb\" else \"tensorboard\",\n+    remove_unused_columns=False,\n+    optim=\"adamw_torch\",\n+    logging_steps=script_args.logging_steps,\n+    evaluation_strategy=\"steps\" if script_args.eval_split != \"none\" else \"no\",\n )\n \n # Step 4: Define the LoraConfig\n@@ -139,6 +167,7 @@\n     tokenizer=tokenizer,\n     args=training_args,\n     train_dataset=train_dataset,\n+    eval_dataset=eval_dataset,\n     peft_config=peft_config,\n     max_length=script_args.seq_length,\n )\n", "issue": "Potential bug of model output dim in reward trainer example\nHi!\r\n\r\nI've been exploring the `examples/scripts/reward_trainer.py` and believe I've found a potential bug that I'd like to address.\r\n\r\n## Potential bug.\r\nTo my understanding, the reward model should produce a scalar output. However, when we load the pre-trained model in the [example script](https://github.com/lvwerra/trl/blob/main/examples/scripts/reward_trainer.py#L75-L80), the model will be a binary classifier by default, with a 2-dim logit output:\r\n\r\n```python\r\nimport torch\r\ninputs = tokenizer(\"Hello, my dog is cute\", return_tensors=\"pt\")\r\nmodel = model.cpu()\r\nwith torch.no_grad():\r\n    out = model(**inputs)\r\nprint(out.logits.shape) # torch.Size([1, 2])\r\n```\r\n\r\nDuring training, this issue was hidden by the [the loss computation part](https://github.com/lvwerra/trl/blob/main/trl/trainer/reward_trainer.py#L169-L176) of the reward trainer:\r\n\r\n```python\r\n    rewards_chosen = model(\r\n        input_ids=inputs[\"input_ids_chosen\"],\r\n        attention_mask=inputs[\"attention_mask_chosen\"],\r\n    )[0]\r\n\r\n    rewards_rejected = model(\r\n        input_ids=inputs[\"input_ids_rejected\"],\r\n        attention_mask=inputs[\"attention_mask_rejected\"],\r\n    )[0]\r\n\r\n    loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()\r\n```\r\nThere, the `rewards_rejected` and `rewards_chosen` are dim [sample, 2] matrices, which I think should be dim [sample, 1] instead. The `.mean()` operator hides the issue as it returns a scalar value.\r\n\r\nHere is a [colab](https://colab.research.google.com/drive/1szzGC8H1DQ7329_idVnnOnPIPE6D7J4j?usp=sharing) that shows the dimensionality problem.\r\n\r\n## Proposed fix\r\n\r\nTo resolve this, in the model training example, we can use:\r\n\r\n```python\r\nmodel = AutoModelForSequenceClassification.from_pretrained(\r\n    script_args.model_name,\r\n    quantization_config=quantization_config,\r\n    device_map=device_map,\r\n    trust_remote_code=script_args.trust_remote_code,\r\n    num_labels=1 # <---- make sure that the model output is a scalar\r\n)\r\n```\r\nI'll submit a pull request if it is OK. \r\n\r\nIn addition, while exploring `examples/scripts/reward_trainer.py`, I noticed a few other places that can be improved. For example, there is currently no validation steps that show the genearlization performance of the reward model; the `log_with` option in `ScriptArguments` is currently not being used anywhere and therefore has no effect. I'm happy to include the validation and fix the logging glitch in my pull request as well.\r\n\r\nTianlin\n", "before_files": [{"content": "# coding=utf-8\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\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import (\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    BitsAndBytesConfig,\n    HfArgumentParser,\n    TrainingArguments,\n)\n\nfrom trl import RewardTrainer\n\n\ntqdm.pandas()\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The name of the Casual LM model we wish to fine with SFTTrainer\n    \"\"\"\n\n    model_name: Optional[str] = field(default=\"facebook/opt-350m\", metadata={\"help\": \"the model name\"})\n    dataset_name: Optional[str] = field(default=\"Anthropic/hh-rlhf\", metadata={\"help\": \"the model name\"})\n    dataset_text_field: Optional[str] = field(default=\"text\", metadata={\"help\": \"the text field of the dataset\"})\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    batch_size: Optional[int] = field(default=64, metadata={\"help\": \"the batch size\"})\n    seq_length: Optional[int] = field(default=512, metadata={\"help\": \"Input sequence length\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=16, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    load_in_8bit: Optional[bool] = field(default=False, metadata={\"help\": \"load the model in 8 bits precision\"})\n    load_in_4bit: Optional[bool] = field(default=False, metadata={\"help\": \"load the model in 4 bits precision\"})\n    use_peft: Optional[bool] = field(default=False, metadata={\"help\": \"Wether to use PEFT or not to train adapters\"})\n    trust_remote_code: Optional[bool] = field(default=True, metadata={\"help\": \"Enable `trust_remote_code`\"})\n    output_dir: Optional[str] = field(default=\"output\", metadata={\"help\": \"the output directory\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\n# Step 1: Load the model\nif script_args.load_in_8bit and script_args.load_in_4bit:\n    raise ValueError(\"You can't load the model in 8 bits and 4 bits at the same time\")\nelif script_args.load_in_8bit or script_args.load_in_4bit:\n    quantization_config = BitsAndBytesConfig(\n        load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit\n    )\n    # This means: fit the entire model on the GPU:0\n    device_map = {\"\": 0}\nelse:\n    device_map = None\n    quantization_config = None\n\nmodel = AutoModelForSequenceClassification.from_pretrained(\n    script_args.model_name,\n    quantization_config=quantization_config,\n    device_map=device_map,\n    trust_remote_code=script_args.trust_remote_code,\n)\n\n# Step 2: Load the dataset and pre-process it\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name)\ndataset = load_dataset(script_args.dataset_name, split=\"train\")\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.\n# Adapt this section to your needs for custom datasets\n\n\ndef 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_j = tokenizer(chosen, truncation=True)\n        tokenized_k = tokenizer(rejected, truncation=True)\n\n        new_examples[\"input_ids_chosen\"].append(tokenized_j[\"input_ids\"])\n        new_examples[\"attention_mask_chosen\"].append(tokenized_j[\"attention_mask\"])\n        new_examples[\"input_ids_rejected\"].append(tokenized_k[\"input_ids\"])\n        new_examples[\"attention_mask_rejected\"].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\noriginal_columns = dataset.column_names\ntrain_dataset = dataset.map(\n    preprocess_function,\n    batched=True,\n    num_proc=4,\n)\ntrain_dataset = train_dataset.filter(\n    lambda x: len(x[\"input_ids_chosen\"]) <= script_args.seq_length\n    and len(x[\"input_ids_rejected\"]) <= script_args.seq_length\n)\n\n\n# Step 3: Define the training arguments\ntraining_args = TrainingArguments(\n    output_dir=script_args.output_dir,\n    per_device_train_batch_size=script_args.batch_size,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    learning_rate=script_args.learning_rate,\n)\n\n# Step 4: Define the LoraConfig\nif script_args.use_peft:\n    peft_config = LoraConfig(r=16, lora_alpha=16, bias=\"none\", task_type=\"SEQ_CLS\", modules_to_save=[\"scores\"])\nelse:\n    peft_config = None\n\n# Step 5: Define the Trainer\ntrainer = RewardTrainer(\n    model=model,\n    tokenizer=tokenizer,\n    args=training_args,\n    train_dataset=train_dataset,\n    peft_config=peft_config,\n    max_length=script_args.seq_length,\n)\n\ntrainer.train()\n", "path": "examples/scripts/reward_trainer.py"}], "after_files": [{"content": "# coding=utf-8\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\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom tqdm import tqdm\nfrom transformers import (\n    AutoModelForSequenceClassification,\n    AutoTokenizer,\n    BitsAndBytesConfig,\n    HfArgumentParser,\n    TrainingArguments,\n)\n\nfrom trl import RewardTrainer\n\n\ntqdm.pandas()\n\n\n# Define and parse arguments.\n@dataclass\nclass ScriptArguments:\n    \"\"\"\n    The name of the Casual LM model we wish to fine with SFTTrainer\n    \"\"\"\n\n    model_name: Optional[str] = field(default=\"facebook/opt-350m\", metadata={\"help\": \"the model name\"})\n    dataset_name: Optional[str] = field(default=\"Anthropic/hh-rlhf\", metadata={\"help\": \"the model name\"})\n    dataset_text_field: Optional[str] = field(default=\"text\", metadata={\"help\": \"the text field of the dataset\"})\n    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n    logging_steps: Optional[int] = field(default=500, metadata={\"help\": \"the number of update steps between two logs\"})\n    eval_split: Optional[str] = field(\n        default=\"none\", metadata={\"help\": \"the dataset split to evaluate on; default to 'none' (no evaluation)\"}\n    )\n    learning_rate: Optional[float] = field(default=1.41e-5, metadata={\"help\": \"the learning rate\"})\n    batch_size: Optional[int] = field(default=64, metadata={\"help\": \"the batch size\"})\n    num_train_epochs: Optional[int] = field(default=1, metadata={\"help\": \"the number of training epochs\"})\n    seq_length: Optional[int] = field(default=512, metadata={\"help\": \"Input sequence length\"})\n    gradient_accumulation_steps: Optional[int] = field(\n        default=16, metadata={\"help\": \"the number of gradient accumulation steps\"}\n    )\n    load_in_8bit: Optional[bool] = field(default=False, metadata={\"help\": \"load the model in 8 bits precision\"})\n    load_in_4bit: Optional[bool] = field(default=False, metadata={\"help\": \"load the model in 4 bits precision\"})\n    use_peft: Optional[bool] = field(default=False, metadata={\"help\": \"Wether to use PEFT or not to train adapters\"})\n    trust_remote_code: Optional[bool] = field(default=True, metadata={\"help\": \"Enable `trust_remote_code`\"})\n    output_dir: Optional[str] = field(default=\"output\", metadata={\"help\": \"the output directory\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\n# Step 1: Load the model\nif script_args.load_in_8bit and script_args.load_in_4bit:\n    raise ValueError(\"You can't load the model in 8 bits and 4 bits at the same time\")\nelif script_args.load_in_8bit or script_args.load_in_4bit:\n    quantization_config = BitsAndBytesConfig(\n        load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit\n    )\n    # This means: fit the entire model on the GPU:0\n    device_map = {\"\": 0}\nelse:\n    device_map = None\n    quantization_config = None\n\nmodel = AutoModelForSequenceClassification.from_pretrained(\n    script_args.model_name,\n    quantization_config=quantization_config,\n    device_map=device_map,\n    trust_remote_code=script_args.trust_remote_code,\n    num_labels=1,\n)\n\n# Step 2: Load the dataset and pre-process it\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name)\ntrain_dataset = load_dataset(script_args.dataset_name, split=\"train\")\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.\n# Adapt this section to your needs for custom datasets\n\n\ndef 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_j = tokenizer(chosen, truncation=True)\n        tokenized_k = tokenizer(rejected, truncation=True)\n\n        new_examples[\"input_ids_chosen\"].append(tokenized_j[\"input_ids\"])\n        new_examples[\"attention_mask_chosen\"].append(tokenized_j[\"attention_mask\"])\n        new_examples[\"input_ids_rejected\"].append(tokenized_k[\"input_ids\"])\n        new_examples[\"attention_mask_rejected\"].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=4,\n)\ntrain_dataset = train_dataset.filter(\n    lambda x: len(x[\"input_ids_chosen\"]) <= script_args.seq_length\n    and len(x[\"input_ids_rejected\"]) <= script_args.seq_length\n)\n\nif script_args.eval_split == \"none\":\n    eval_dataset = None\nelse:\n    eval_dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)\n\n    eval_dataset = eval_dataset.map(\n        preprocess_function,\n        batched=True,\n        num_proc=4,\n    )\n    eval_dataset = eval_dataset.filter(\n        lambda x: len(x[\"input_ids_chosen\"]) <= script_args.seq_length\n        and len(x[\"input_ids_rejected\"]) <= script_args.seq_length\n    )\n\n\n# Step 3: Define the training arguments\ntraining_args = TrainingArguments(\n    output_dir=script_args.output_dir,\n    per_device_train_batch_size=script_args.batch_size,\n    num_train_epochs=script_args.num_train_epochs,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    learning_rate=script_args.learning_rate,\n    report_to=\"wandb\" if script_args.log_with == \"wandb\" else \"tensorboard\",\n    remove_unused_columns=False,\n    optim=\"adamw_torch\",\n    logging_steps=script_args.logging_steps,\n    evaluation_strategy=\"steps\" if script_args.eval_split != \"none\" else \"no\",\n)\n\n# Step 4: Define the LoraConfig\nif script_args.use_peft:\n    peft_config = LoraConfig(r=16, lora_alpha=16, bias=\"none\", task_type=\"SEQ_CLS\", modules_to_save=[\"scores\"])\nelse:\n    peft_config = None\n\n# Step 5: Define the Trainer\ntrainer = RewardTrainer(\n    model=model,\n    tokenizer=tokenizer,\n    args=training_args,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n    peft_config=peft_config,\n    max_length=script_args.seq_length,\n)\n\ntrainer.train()\n", "path": "examples/scripts/reward_trainer.py"}]}