Merge pull request #222 from theblackcat102/main

add training code for reward model

model/reward/instructor/README.md

@@ -0,0 +1,53 @@
+# Sections to train Reward Model (RM)
+
+Trainer code based on huggingface. Compatible with deepspeed or accelerate
+
+Requirements
+
+```
+wandb
+evaluate
+datasets
+transformers
+torch==1.12
+```
+
+Start training reward model
+
+```bash
+python trainer.py configs/electra-base-dis-webgpt.yml
+```
+
+Additional axis labeling, this outputs a 4 summary quality evaluation metrics (score are normalized to 0-1 )
+
+```bash
+python summary_quality_trainer.py configs/test-bloomz-560m-quality.yml
+```
+
+The four summary are :
+
+- overall
+
+- accuracy
+
+- coverage
+
+- coherence
+
+## Dataset
+
+For now we only supports webgpt and summary dataset from OpenAI. Once open-asisstant dataset are available it will be added here.
+
+## Model
+
+Check out configs
+
+```
+Open-Assistant/model/reward/instructor/configs/
+    bloomz-560m.yml
+    electra-base-dis-webgpt.yml
+    galactica-125m.yml
+    galactica-1b.yml
+```
+
+You can add new huggingface model as you want.

model/reward/instructor/TODO.md

@@ -0,0 +1,19 @@
+Some other reward features we can use
+
+0. Finish classifcation feature
+
+1. Summaries from human feedback
+
+- use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
+
+- each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
+
+- ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
+
+  - this should be placed under experimental_dataset.py
+
+2. Add support for anthropic dataset
+
+- anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
+
+  - this is basically a MCTS from alphazero.

model/reward/instructor/cls_dataset.py

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+"""
+
+    classification based ranking
+
+"""
+import json
+import os
+import random
+
+from dataset import load_dataset
+from torch.utils.data import Dataset
+
+from .utils import webgpt_return_format
+
+
+class WebGPTDataset(Dataset):
+    def __init__(self, mode="train", index_cache="dataset/webgpt_train_idx.pt", additional_dataset=None) -> None:
+        super().__init__()
+        """
+            mode : train or val, used for validation purpose, has nothing to do with original split
+            additional_dataset : a list of jsonline format with idx, question and texts (generate candidates)
+                idx : must match the index you iterate from comparison enumerate order
+                question : for validation purpose
+                texts : list of K generate results from the question prompt
+        """
+        os.makedirs("dataset", exist_ok=True)
+        dataset = load_dataset("openai/webgpt_comparisons")
+        self.dataset = []
+        self.dataset_index = []
+        for idx, row in enumerate(dataset["train"]):
+            self.dataset.append(webgpt_return_format(row))
+
+        # since this dataset was generated from 176B GPT-3
+        # we needed some more sample generated from the starting model
+        # since this model must rank model generated by GPT-3 being better than your starting model
+        self.sample_additional = False
+        if additional_dataset is not None:
+            self.sample_additional = True
+            self.additional = {}
+            with open(additional_dataset, "r") as f:
+                for line in f:
+                    row = json.loads(line)
+                    if row["idx"] in self.dataset_index:
+                        self.additional[row["idx"]] = row["negatives"]
+            if len(self.additional) != len(self.dataset_index):
+                for match_idx in self.dataset_index:
+                    if match_idx in self.additional:
+                        continue
+
+                    idx = match_idx - 900
+                    while idx not in self.additional:
+                        idx -= 1
+                    self.additional[match_idx] = self.additional[idx]
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        row = self.dataset[index]
+        if not self.sample_additional:
+            return row["question"], row["pos"], row["neg"]
+
+        gen_neg = random.choice(self.additional[self.dataset_index[index]])
+        return row["question"], row["pos"], row["neg"], gen_neg

model/reward/instructor/configs/bloomz-560m-summary.yml

@@ -0,0 +1,9 @@
+model_name: bigscience/bloomz-560m
+learning_rate: 3e-5
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 2
+max_length: 600
+freeze_layer: 12
+num_train_epochs: 2
+datasets:
+  - hfsummary

model/reward/instructor/configs/bloomz-560m.yml

@@ -0,0 +1,10 @@
+model_name: bigscience/bloomz-560m
+learning_rate: 3e-5
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 2
+max_length: 600
+freeze_layer: 12
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/electra-base-dis-webgpt.yml

@@ -0,0 +1,3 @@
+model_name: google/electra-large-discriminator
+learning_rate: 3e-5
+max_length: 300

model/reward/instructor/configs/galactica-125m.yml

@@ -0,0 +1,13 @@
+model_name: facebook/galactica-125m
+learning_rate: 1e-5
+gradient_checkpointing: false
+gradient_accumulation_steps: 32
+per_device_train_batch_size: 2
+warmup_steps: 600
+eval_steps: 200
+save_steps: 500
+max_length: 512
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/galactica-1b.yml

@@ -0,0 +1,14 @@
+model_name: facebook/galactica-1.3b
+learning_rate: 6e-6
+gradient_checkpointing: false
+gradient_accumulation_steps: 16
+per_device_train_batch_size: 2
+warmup_steps: 600
+freeze_layer: 20
+eval_steps: 200
+save_steps: 500
+max_length: 400
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/configs/test-galactica-125m-classification.yml

@@ -0,0 +1,14 @@
+model_name: facebook/galactica-125m
+learning_rate: 1e-5
+gradient_checkpointing: false
+gradient_accumulation_steps: 10
+per_device_train_batch_size: 6
+warmup_steps: 600
+loss: cls
+eval_steps: 200
+save_steps: 500
+max_length: 128
+num_train_epochs: 2
+datasets:
+  - webgpt
+  - hfsummary

model/reward/instructor/experimental_dataset.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+"""
+    HFSummary
+
+        I want to train a multi regression model on axis_evals dataset mainly we can estimate the score of these score
+
+         - {"overall": "6", "accuracy": "6", "coverage": "6", "coherence": "7"}
+
+        Should be better than just a preference score
+
+"""
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from torch.utils.data import Dataset
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+
+
+@dataclass
+class DataCollatorForSummaryScore:
+    """
+
+    Data collator that will dynamically pad the inputs for multiple choice received.
+
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    num_choices: int = 2
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    drop_token_type: bool = False  # galactica
+
+    def __call__(self, batch):
+
+        features = []
+        labels = []
+        for feature, label in batch:
+            features.append(feature)
+            labels.append(label)
+
+        batch_feature = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        if self.drop_token_type:
+            batch_feature.pop("token_type_ids")
+        # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
+        batch_feature["labels"] = torch.from_numpy(np.array(labels)).float()
+        return batch_feature
+
+
+class HFSummaryQuality(Dataset):
+    def __init__(self, split, tokenizer, max_length=300) -> None:
+        super().__init__()
+        assert split in ("validation", "test")
+        dataset = load_dataset("Tristan/summarize_from_feedback", "axis")[split]
+        self.max_length = max_length
+        mean_scores = defaultdict(list)
+        self.contexts = []
+        self.responses = []
+        self.labels = []
+        for data in dataset:
+
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
+            self.contexts.append(context)
+
+            response = data["summary"]["text"]
+            self.responses.append(response)
+            self.labels.append(data["summary"]["axes"])
+            for axis, score in data["summary"]["axes"].items():
+                if score is not None:
+                    mean_scores[axis].append(score)
+
+        self.label2idx = {key: idx for idx, key in enumerate(mean_scores.keys())}
+        self.label2mean = {key: np.mean(scores) for key, scores in mean_scores.items()}
+        self.tokenizer = tokenizer
+        print(self.label2idx)
+
+    def __len__(self):
+        return len(self.responses)
+
+    def __getitem__(self, index):
+        context = self.contexts[index]
+        # return pairs of comparison
+        response = self.responses[index]
+        labels = np.zeros(len(self.label2idx))
+        for key, score in self.labels[index].items():
+            labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score) / 10
+        return self.tokenizer(context, response, truncation=True, max_length=self.max_length), labels

model/reward/instructor/rank_datasets.py

@@ -0,0 +1,166 @@
+# -*- coding: utf-8 -*-
+"""
+    author: theblackcat102
+
+    Dataset output format from __getitem__
+
+     - question / prompt : string
+
+     - answers / rows : list of tuple pair. The first element in the tuple pair must be the positive pair (rank higher than the second element)
+
+    A list of rank based dataset for training using rank loss
+
+    Some nice features to have
+
+    [] support additional negative samples generated from other models.
+
+        For example we can use galactica-125m to generate a TLDR and assume it was
+        inferior than the human perference one
+
+
+"""
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+from datasets import load_dataset
+from torch.utils.data import Dataset
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+
+
+@dataclass
+class DataCollatorForPairRank:
+    """
+
+    Data collator that will dynamically pad the inputs for multiple choice received.
+
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    num_choices: int = 2
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    drop_token_type: bool = False  # galactica
+
+    def __call__(self, features):
+
+        flatten_features = []
+        batch_size = 0
+        for question, pairs in features:
+            for (pos, neg) in pairs:
+                flatten_features.append(self.tokenizer(question, pos, truncation=True, max_length=self.max_length))
+                flatten_features.append(self.tokenizer(question, neg, truncation=True, max_length=self.max_length))
+                batch_size += 1
+
+        batch = self.tokenizer.pad(
+            flatten_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        if self.drop_token_type:
+            batch.pop("token_type_ids")
+        # batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
+        return batch
+
+
+class WebGPT(Dataset):
+    def __init__(self) -> None:
+        super().__init__()
+
+        dataset = load_dataset("openai/webgpt_comparisons")
+        questions = {}
+        # using prompt as our index will allows us
+        # to add additional generated prompt later
+        self.index2question = {}
+        for row in dataset["train"]:
+            question = row["question"]["full_text"]
+            if question not in self.index2question:
+                self.index2question[len(self.index2question)] = question
+
+            if question not in questions:
+                questions[question] = []
+
+            if row["score_0"] > row["score_1"]:
+                # not going to risk it
+                questions[question].append((row["answer_0"], row["answer_1"]))
+            else:
+                questions[question].append((row["answer_1"], row["answer_0"]))
+
+        self.questions = questions
+
+    def __len__(self):
+        return len(self.index2question)
+
+    def __getitem__(self, index):
+        question = self.index2question[index]
+        rows = self.questions[question]
+        # optimize the format later
+        return question, rows
+
+
+class HFSummary(Dataset):
+    """
+    Human feedback data from OpenAI
+    https://github.com/openai/summarize-from-feedback
+
+    labeling method : pair comparison, 0 or 1
+
+    """
+
+    def __init__(self, split="train", conf_threshold=-1, max_comparison_per_sample=3) -> None:
+        super().__init__()
+        assert split in ("train", "valid1", "valid2", "test")
+        summaries = {}
+        # using prompt as our index will allows us
+        # to add additional generated prompt later
+        self.index2summary = {}
+        self.max_comparison_per_sample = max_comparison_per_sample
+        major_split = split if "train" == split else "validation"
+        dataset = load_dataset("Tristan/summarize_from_feedback", "comparisons")[major_split]
+        for data in dataset:
+            if (
+                "extra" in data
+                and "confidence" in data["extra"]
+                and data["extra"]["confidence"] is not None
+                and conf_threshold > data["extra"]["confidence"]
+            ):
+                print("skipping {}".format(data["info"]["id"]))
+                continue
+
+            if split != "train" and split != data["split"]:
+                continue
+
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
+
+            if context not in self.index2summary:
+                self.index2summary[len(self.index2summary)] = context
+
+            if context not in summaries:
+                summaries[context] = []
+
+            pos, neg = (0, 1) if data["choice"] == 0 else (1, 0)
+            summaries[context].append((data["summaries"][pos]["text"], data["summaries"][neg]["text"]))
+
+        self.summaries = summaries
+
+        self.postfix_prompt = " TLDR;"
+
+    def __len__(self):
+        return len(self.index2summary)
+
+    def __getitem__(self, index):
+        context = self.index2summary[index]
+        # return pairs of comparison
+        rows = self.summaries[context]
+        # pair very big
+        # we are going to do some sampling
+        # not optimal but good for now
+        valid_idx = np.random.choice(len(rows), self.max_comparison_per_sample)
+        # optimize the format later
+        return context + self.postfix_prompt, [r for idx, r in enumerate(rows) if idx in valid_idx]

model/reward/instructor/requirements.txt

@@ -0,0 +1,6 @@
+datasets==2.8.0
+evaluate==0.4.0
+scikit-learn==1.2.0
+torch==1.12.1+cu116
+transformers==4.25.1
+wandb==0.13.7

model/reward/instructor/summary_quality_trainer.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+import os
+from argparse import ArgumentParser
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
+import numpy as np
+import torch
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from torch import nn
+from torch.utils.data import Dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer
+
+os.environ["WANDB_PROJECT"] = "quality-scoring"
+
+parser = ArgumentParser()
+parser.add_argument("config", type=str)
+
+accuracy = evaluate.load("mse")
+
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    return accuracy.compute(predictions=predictions.flatten(), references=labels.flatten())
+
+
+class QualityTrainer(Trainer):
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.loss_fct = nn.L1Loss()
+        self.sigmoid = nn.Sigmoid()
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        labels = inputs.pop("labels")
+        # forward pass
+        outputs = model(**inputs)
+        logits = self.sigmoid(outputs.get("logits"))
+        loss = self.loss_fct(logits, labels)
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _compute_loss(self, model, inputs):
+        inputs = self._prepare_inputs(inputs)
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        logits = self.sigmoid(outputs.get("logits"))
+        loss = self.loss_fct(logits, labels)
+
+        return loss, logits
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+        with torch.no_grad():
+            # compute loss on predict data
+            loss, logits = self._compute_loss(model, inputs)
+
+        loss = loss.mean().detach()
+        labels = inputs["labels"]
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, logits, labels)
+
+
+if __name__ == "__main__":
+    training_conf = argument_parsing(parser)
+
+    model_name = training_conf["model_name"]
+    tokenizer = get_tokenizer(model_name)
+    collate_fn = DataCollatorForSummaryScore(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
+    train = HFSummaryQuality(split="validation", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    eval = HFSummaryQuality(split="test", tokenizer=tokenizer, max_length=training_conf["max_length"])
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=len(train.label2idx), problem_type="regression"
+    )
+
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
+        model = freeze_top_n_layers(model, num_layer)
+        model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        print("Number of trainable : {}M".format(int(params / 1e6)))
+
+    args = TrainingArguments(
+        output_dir=f"{model_name}-finetuned",
+        num_train_epochs=training_conf["num_train_epochs"],
+        warmup_steps=500,
+        learning_rate=training_conf["learning_rate"],
+        # half_precision_backend="apex",
+        fp16=True,
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
+        weight_decay=0.01,
+        max_grad_norm=2.0,
+        logging_steps=10,
+        save_total_limit=4,
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
+        save_steps=1000,
+        report_to="wandb",
+    )
+    trainer = QualityTrainer(
+        model,
+        args,
+        train_dataset=train,
+        eval_dataset=eval,
+        data_collator=collate_fn,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+    trainer.train()

model/reward/instructor/tests/test_dataset.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+
+def test_hfsummary():
+
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
+    dataset = HFSummary("train")
+    print(len(dataset))
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
+    for batch in dataloader:
+        batch["input_ids"].shape
+
+
+def test_webgpt():
+
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
+    dataset = WebGPT()
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
+    for batch in dataloader:
+        print(batch["input_ids"].shape)
+
+
+def test_hf_quality():
+
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
+    collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
+    dataset = HFSummaryQuality("validation", tokenizer)
+    dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
+    for batch in dataloader:
+        print(batch["input_ids"].shape)
+
+
+if __name__ == "__main__":
+    test_hf_quality()
+    # test_webgpt()

model/reward/instructor/trainer.py

@@ -0,0 +1,186 @@
+# -*- coding: utf-8 -*-
+import os
+from argparse import ArgumentParser
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import evaluate
+import numpy as np
+import torch
+from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
+from torch import nn
+from torch.utils.data import ConcatDataset, Dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from utils import argument_parsing, freeze_top_n_layers, get_tokenizer, train_val_dataset
+
+os.environ["WANDB_PROJECT"] = "reward-model"
+
+accuracy = evaluate.load("accuracy")
+parser = ArgumentParser()
+parser.add_argument("config", type=str)
+
+
+@dataclass
+class CustomTrainingArguments(TrainingArguments):
+    loss_function: str = "rank"
+
+
+def compute_metrics(eval_pred):
+    predictions, _ = eval_pred
+    predictions = np.argmax(predictions, axis=1)
+    return accuracy.compute(predictions=predictions, references=[0] * predictions.shape[0])
+
+
+class RankLoss(nn.Module):
+    def __init__(self, eps=1e-8) -> None:
+        super().__init__()
+        self.eps = eps
+        self.log_sigmoid = nn.LogSigmoid()
+
+    def forward(self, pos, neg):
+        return -self.log_sigmoid(pos - neg + self.eps).mean()
+
+
+class RankTrainer(Trainer):
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.loss_fct = RankLoss() if args.loss_function == "rank" else nn.CrossEntropyLoss()
+        self.loss_function = args.loss_function
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        # forward pass
+        outputs = model(**inputs)
+        logits = outputs.get("logits").view(-1, 2)
+        if self.loss_function == "rank":
+            loss = self.loss_fct(logits[:, 0], logits[:, 1])
+        else:
+            loss = self.loss_fct(logits, torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long))
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _compute_loss(self, model, inputs):
+        inputs = self._prepare_inputs(inputs)
+        outputs = model(**inputs)
+        logits = outputs.get("logits").view(-1, 2)
+        if self.loss_function == "rank":
+            loss = self.loss_fct(logits[:, 0], logits[:, 1])
+        else:
+            loss = self.loss_fct(logits, torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long))
+
+        return loss, logits
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+        with torch.no_grad():
+            # compute loss on predict data
+            loss, logits = self._compute_loss(model, inputs)
+
+        loss = loss.mean().detach()
+        labels = torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long)
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, logits, labels)
+
+
+if __name__ == "__main__":
+    training_conf = argument_parsing(parser)
+
+    model_name = training_conf["model_name"]
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type="regression")
+    if "freeze_layer" in training_conf:
+        num_layer = training_conf["freeze_layer"]
+        model = freeze_top_n_layers(model, num_layer)
+        model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        print("Number of trainable : {}M".format(int(params / 1e6)))
+
+    tokenizer = get_tokenizer(model_name)
+    args = CustomTrainingArguments(
+        output_dir=f"{model_name}-finetuned",
+        num_train_epochs=training_conf["num_train_epochs"],
+        warmup_steps=500,
+        loss_function=training_conf["loss"],
+        learning_rate=training_conf["learning_rate"],
+        # half_precision_backend="apex",
+        fp16=True,
+        gradient_checkpointing=training_conf["gradient_checkpointing"],
+        gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
+        per_device_train_batch_size=training_conf["per_device_train_batch_size"],
+        per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
+        weight_decay=0.01,
+        max_grad_norm=2.0,
+        logging_steps=10,
+        save_total_limit=4,
+        evaluation_strategy="steps",
+        eval_steps=training_conf["eval_steps"],
+        save_steps=1000,
+        report_to="wandb",
+    )
+    train_datasets, evals = [], {}
+    if "webgpt" in training_conf["datasets"]:
+        web_dataset = WebGPT()
+        train, eval = train_val_dataset(web_dataset)
+        train_datasets.append(train)
+        evals["webgpt"] = eval
+    if "hfsummary" in training_conf["datasets"]:
+        sum_train = HFSummary(split="train")
+        train_datasets.append(sum_train)
+        sum_eval = HFSummary(split="valid1")
+        assert len(sum_eval) > 0
+        evals["hfsummary"] = sum_eval
+    train = ConcatDataset(train_datasets)
+    collate_fn = DataCollatorForPairRank(
+        tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
+    )
+    assert len(evals) > 0
+    trainer = RankTrainer(
+        model,
+        args,
+        train_dataset=train,
+        eval_dataset=eval,
+        data_collator=collate_fn,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+    trainer.train()

model/reward/instructor/utils.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+import re
+
+import yaml
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Subset
+from transformers import AutoTokenizer
+
+re_reference_remove = re.compile(r"\[([0-9])+\]|\[([0-9])+,([0-9])+\]")
+
+
+def webgpt_return_format(row):
+    if row["score_0"] >= row["score_1"]:
+        # remove this to prevent information leak, since we are not using reference
+        return {
+            "question": row["question"]["full_text"],
+            "pos": re_reference_remove.sub("", row["answer_0"]),
+            "neg": re_reference_remove.sub("", row["answer_1"]),
+        }
+
+    return {
+        "question": row["question"]["full_text"],
+        "pos": re_reference_remove.sub("", row["answer_1"]),
+        "neg": re_reference_remove.sub("", row["answer_0"]),
+    }
+
+
+def get_tokenizer(tokenizer_name):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    if "galactica" in tokenizer_name:
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
+
+    return tokenizer
+
+
+def train_val_dataset(dataset, val_split=0.2):
+    train_idx, val_idx = train_test_split(
+        list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
+    )
+    # [3879, 11479, 8341, 9177, 10798, 18177, 5735, 15669, 4837, 2760]
+    print(val_idx[:10])
+    # [13582, 5919, 11875, 7373, 19135, 13706, 8555, 15788, 15005, 15209]
+    print(train_idx[:10])
+    return Subset(dataset, train_idx), Subset(dataset, val_idx)
+
+
+def freeze_top_n_layers(model, target_layers):
+    # its possible we can simply detect which module is a ModuleList
+    # and simply freeze the module without doing string parsing
+    for name, param in model.named_parameters():
+        if "embed" in name:
+            param.requires_grad = False
+        elif ".layer" in name or ".h." in name:
+            tokens = name.split(".")
+            idx = 0
+            for token in tokens:
+                if "layer" in token or token == "h":
+                    break
+                idx += 1
+            if idx >= len(tokens):
+                continue
+
+            layer_ = int(tokens[idx + 1])
+            if layer_ < target_layers:
+                # print('freeze ', layer_, name)
+                param.requires_grad = False
+    return model
+
+
+def argument_parsing(parser):
+    default_params = {
+        "num_train_epochs": 4,
+        "learning_rate": 3e-5,
+        "eval_steps": 500,
+        "loss": "rank",
+        "max_length": 440,
+        "per_device_eval_batch_size": 5,
+        "per_device_train_batch_size": 8,
+        "gradient_accumulation_steps": 8,
+        "gradient_checkpointing": False,
+        "datasets": ["webgpt"],
+    }
+    args = parser.parse_args()
+    with open(args.config, "r", encoding="utf-8") as f:
+        training_conf = yaml.safe_load(f.read())
+
+    params = {**default_params, **training_conf}
+    params["gradient_accumulation_steps"] = int(params["gradient_accumulation_steps"])
+    params["num_train_epochs"] = int(params["num_train_epochs"])
+    params["per_device_train_batch_size"] = int(params["per_device_train_batch_size"])
+    params["learning_rate"] = float(params["learning_rate"])
+    return params
+
+
+if __name__ == "__main__":
+    from transformers import AutoModelForSequenceClassification
+
+    model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
+    freeze_top_n_layers(model, 10)
+    print(model.state_dict().keys())