[fix] resolve conflict from main

model/reward/instructor/trainer.py

@@ -8,7 +8,15 @@ import torch
 from models import RankGenModel
 from rank_datasets import DataCollatorForPairRank, RankGenCollator
 from torch import nn
-from transformers import AutoModelForSequenceClassification, PreTrainedModel, Trainer, TrainingArguments
+from transformers import (
+    AdamW,
+    AutoModelForSequenceClassification,
+    PreTrainedModel,
+    Trainer,
+    TrainingArguments,
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+)
 from utils import argument_parsing, freeze_top_n_layers, get_datasets, get_tokenizer
 
 os.environ["WANDB_PROJECT"] = "reward-model"
@@ -144,7 +152,7 @@ if __name__ == "__main__":
         evaluation_strategy="steps",
         eval_steps=training_conf["eval_steps"],
         save_steps=1000,
-        report_to="local",
+        report_to="wandb",
     )
 
     tokenizer = get_tokenizer(training_conf["tokenizer_name"])
@@ -154,6 +162,21 @@ if __name__ == "__main__":
     else:
         collate_fn = DataCollatorForPairRank(tokenizer, max_length=training_conf["max_length"])
     assert len(evals) > 0
+
+    optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+    scheduler = None
+    if "scheduler" in training_conf:
+        if training_conf["scheduler"] == "linear":
+            scheduler = get_linear_schedule_with_warmup()
+        elif training_conf["scheduler"] == "cosine":
+            scheduler = get_cosine_schedule_with_warmup(
+                optimizer,
+                num_warmup_steps=args.warmup_steps,
+                num_training_steps=len(train)
+                * args.num_train_epochs
+                / (args.per_device_train_batch_size * args.gradient_accumulation_steps),
+            )
+
     trainer = RankTrainer(
         model=model,
         model_name=model_name,
@@ -164,6 +187,7 @@ if __name__ == "__main__":
         data_collator=collate_fn,
         tokenizer=tokenizer,
         compute_metrics=compute_metrics,
+        optimizers=(optimizer, scheduler),
     )
     # trainer.evaluate()
     trainer.train()