The dataset is a relabel dataset of the 22k randomly sampled hhrlhf dataset (11k from helpful-base, 11k from harmless-base).
The annotators are Reward Model trained on this 22k original dataset based on the llama-2-7b .
The annotation python script is as follows:
Unless you fully understand the significance of this dataset, I do not recommend using it lightly.
The re-annotation of the dataset entirely depends on a llama-2-7b reward model, which I have not uploaded yet.
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
from datasets import load_dataset
from scipy.special import expit
import numpy as np
import random
import json
from tqdm import tqdm
def sample_probability(A, B, num_samples):
probabilities = expit(A - B)
samples = np.random.rand(num_samples) < probabilities
return samples.astype(int)
def set_seed(seed=42):
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
def get_prompt(split="train", data_dir="helpful-base"):
dataset_helpful = load_dataset("PKU-Alignment/processed-hh-rlhf", split=split, data_dir=data_dir)
dataset = [d for d in dataset_helpful]
prompts = [d['context'] for d in dataset]
prompts_temp = []
for p in prompts:
prompt = ""
for i in range(len(p)):
if p[i]['role'] == "human":
prompt += "\n\nHuman: " + p[i]['text']
else:
prompt += "\n\nAssistant: " + p[i]['text']
prompt += "\n\nAssistant: "
prompts_temp.append(prompt)
prompts = prompts_temp
chosen = [d['chosen']['text'] for d in dataset]
rejected = [d['rejected']['text'] for d in dataset]
return prompts, chosen, rejected
def get_prompt_from_path(path):
with open(path, "r") as f:
dataset = json.load(f)
prompts = [d['context'] for d in dataset]
prompts_temp = []
for p in prompts:
prompt = ""
for i in range(len(p)):
if p[i]['role'] == "human":
prompt += "\n\nHuman: " + p[i]['text']
else:
prompt += "\n\nAssistant: " + p[i]['text']
prompt += "\n\nAssistant: "
prompts_temp.append(prompt)
prompts = prompts_temp
chosen = [d['chosen']['text'] for d in dataset]
rejected = [d['rejected']['text'] for d in dataset]
return prompts, chosen, rejected, dataset
if __name__ == "__main__":
model_path = "/mnt/petrelfs/liuzhixuan/safe-rlhf/output/RM_llama-2-7b-base_mix_train_hhrlhf_5e-6_1e-5_condition_train_baseline"
# bf16
model = AutoModelForSequenceClassification.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)
set_seed()
splits = ["train", "train", "test", "test"]
data_dirs = ["helpful-base", "harmless-base", "helpful-base", "harmless-base"]
paths = [
("../dataset_PKU-Alignment_processed-hh-rlhf_test_harmless-base.json", "test", "harmless",),
("../dataset_PKU-Alignment_processed-hh-rlhf_train_harmless-base.json", "train", "harmless",),
("../dataset_PKU-Alignment_processed-hh-rlhf_test_helpful-base.json", "test", "helpful",),
("../dataset_PKU-Alignment_processed-hh-rlhf_train_helpful-base.json", "train", "helpful",),
]
for path, split, data_dir in paths:
prompts, chosen, rejected, dataset = get_prompt_from_path(path)
log_dict = []
step = 0
for prompt, chosen_response, rejected_response in tqdm(zip(prompts, chosen, rejected), total=len(prompts)):
chosen_text = prompt + chosen_response + tokenizer.eos_token
rejected_text = prompt + rejected_response + tokenizer.eos_token
input_chosen = tokenizer(chosen_text, return_tensors="pt", truncation=True, max_length=1024).to(model.device)
input_rejected = tokenizer(rejected_text, return_tensors="pt", truncation=True, max_length=1024).to(model.device)
outputs_chosen = model(**input_chosen).logits.item()
outputs_rejected = model(**input_rejected).logits.item()
# the prob of chosen is better than rejected is sigmoid(outputs_chosen - outputs_rejected)
new_label = sample_probability(outputs_chosen, outputs_rejected, 1).item()
log_dict.append(
{
"context": dataset[step]['context'],
"chosen": dataset[step]['chosen'],
"rejected": dataset[step]['rejected'],
"scores_chosen": outputs_chosen,
"scores_rejected": outputs_rejected,
"new_chosen_label": not new_label, # if new_label is 1, then chosen is better, so reverse the label. Now: False : Chosen is better. True: Rejected is better
}
)
step += 1
with open(f"hhrlhf_relabel_{split}_{data_dir}.json", "w") as f:
json.dump(log_dict, f, indent=4)