| import json | |
| import os | |
| from typing import List | |
| import datasets | |
| _HF_ENDPOINT = os.getenv("HF_ENDPOINT", "https://huggingface.co") | |
| _DESCRIPTION = "Human preference data about helpfulness and harmlessness." | |
| _CITATION = "" | |
| _HOMEPAGE = "{}/datasets/Anthropic/hh-rlhf".format(_HF_ENDPOINT) | |
| _LICENSE = "mit" | |
| _URL = "{}/datasets/Anthropic/hh-rlhf/resolve/main/".format(_HF_ENDPOINT) | |
| _URLS = { | |
| "train": [ | |
| _URL + "harmless-base/train.jsonl.gz", | |
| _URL + "helpful-base/train.jsonl.gz", | |
| _URL + "helpful-online/train.jsonl.gz", | |
| _URL + "helpful-rejection-sampled/train.jsonl.gz", | |
| ], | |
| "test": [ | |
| _URL + "harmless-base/test.jsonl.gz", | |
| _URL + "helpful-base/test.jsonl.gz", | |
| _URL + "helpful-online/test.jsonl.gz", | |
| _URL + "helpful-rejection-sampled/test.jsonl.gz", | |
| ], | |
| } | |
| class HhRlhfEn(datasets.GeneratorBasedBuilder): | |
| VERSION = datasets.Version("0.0.0") | |
| def _info(self) -> datasets.DatasetInfo: | |
| features = datasets.Features( | |
| { | |
| "instruction": datasets.Value("string"), | |
| "chosen": datasets.Value("string"), | |
| "rejected": datasets.Value("string"), | |
| "history": datasets.Sequence(datasets.Sequence(datasets.Value("string"))), | |
| } | |
| ) | |
| return datasets.DatasetInfo( | |
| description=_DESCRIPTION, features=features, homepage=_HOMEPAGE, license=_LICENSE, citation=_CITATION | |
| ) | |
| def _split_generators(self, dl_manager: datasets.DownloadManager): | |
| file_path = dl_manager.download_and_extract(_URLS) | |
| return [ | |
| datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepaths": file_path["train"]}), | |
| datasets.SplitGenerator(name=datasets.Split.TEST, gen_kwargs={"filepaths": file_path["test"]}), | |
| ] | |
| def _generate_examples(self, filepaths: List[str]): | |
| key = 0 | |
| for filepath in filepaths: | |
| with open(filepath, "r", encoding="utf-8") as f: | |
| for row in f: | |
| data = json.loads(row) | |
| chosen = data["chosen"] | |
| rejected = data["rejected"] | |
| assist_idx = rejected.rfind("\n\nAssistant: ") | |
| r_reject = rejected[assist_idx + 13 :].strip() | |
| assist_idx = chosen.rfind("\n\nAssistant: ") | |
| r_accept = chosen[assist_idx + 13 :].strip() | |
| human_idx = chosen.rfind("\n\nHuman: ") | |
| query = chosen[human_idx + 9 : assist_idx].strip() | |
| prompt = chosen[:human_idx] | |
| history = [] | |
| while prompt.rfind("\n\nAssistant: ") != -1: | |
| assist_idx = prompt.rfind("\n\nAssistant: ") | |
| human_idx = prompt.rfind("\n\nHuman: ") | |
| if human_idx != -1: | |
| old_query = prompt[human_idx + 9 : assist_idx].strip() | |
| old_resp = prompt[assist_idx + 13 :].strip() | |
| history.insert(0, (old_query, old_resp)) | |
| else: | |
| break | |
| prompt = prompt[:human_idx] | |
| yield key, {"instruction": query, "chosen": r_accept, "rejected": r_reject, "history": history} | |
| key += 1 | |