kh4dien/hh-rlhf-helpful-only

dataset_info:
features:

• name: chosen
  list:
  • name: content
    dtype: string
  • name: role
    dtype: string
• name: rejected
  list:
  • name: content
    dtype: string
  • name: role
    dtype: string
• name: source
  dtype: string
  splits:
• name: train
  num_bytes: 188419627
  num_examples: 74428
• name: sft
  num_bytes: 50539545
  num_examples: 20000
• name: test
  num_bytes: 9934147
  num_examples: 3886
  download_size: 137739490
  dataset_size: 248893319
  configs:
• config_name: default
  data_files:
  • split: train
    path: data/train-*
  • split: sft
    path: data/sft-*
  • split: test
    path: data/test-*

from datasets import Dataset, DatasetDict
from collections import defaultdict
import re
import random
import json

data = defaultdict(list)

paths = {
    "test_online": "./test_online.jsonl",
    "train_online": "./train_online.jsonl",
    "test_rejection": "./test_rejection.jsonl",
    "train_rejection": "./train_rejection.jsonl",
}

for name, path in paths.items():
    with open(path, "r") as f:
        for line in f:
            data[name].append(json.loads(line))

def split_data(text):
    delimiter = r"\n\nHuman: |\n\nAssistant: "
    messages = re.split(delimiter, text)[1:]

    conversation = []
    for i, message in enumerate(messages):
        if i % 2 == 0:
            conversation.append({"role": "user", "content": message})
        else:
            conversation.append({"role": "assistant", "content": message})

    return conversation

def fix(row, source):
    return {
        "chosen": split_data(row["chosen"]),
        "rejected": split_data(row["rejected"]),
        "source": source,
    }

train_rejection = [fix(d, "rejection") for d in data["train_rejection"]]
train_online = [fix(d, "online") for d in data["train_online"]]
test_rejection = [fix(d, "rejection") for d in data["test_rejection"]]
test_online = [fix(d, "online") for d in data["test_online"]]

train = train_rejection + train_online
test = test_rejection + test_online

sft_mix = random.sample(train, 20_000)

dataset = DatasetDict({
    "train": Dataset.from_list(train),
    "sft" : Dataset.from_list(sft_mix),
    "test": Dataset.from_list(test)
})

dataset.push_to_hub(<repo-id>)