| import os |
| import json |
| import sys |
| from datasets import load_dataset |
|
|
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) |
| import config |
|
|
| def extract_prompt_chosen_rejected(example): |
| prompt = None |
| chosen = None |
| rejected = None |
|
|
| if "chosen" in example and "rejected" in example: |
| c_list = example["chosen"] |
| r_list = example["rejected"] |
|
|
| if isinstance(c_list, list) and len(c_list) > 0: |
| last_c = c_list[-1] |
| chosen = last_c.get("content", "") if isinstance(last_c, dict) else str(last_c) |
| if len(c_list) > 1: |
| prompt_item = c_list[0] |
| prompt = prompt_item.get("content", "") if isinstance(prompt_item, dict) else str(prompt_item) |
|
|
| if isinstance(r_list, list) and len(r_list) > 0: |
| last_r = r_list[-1] |
| rejected = last_r.get("content", "") if isinstance(last_r, dict) else str(last_r) |
| if not prompt and len(r_list) > 1: |
| prompt_item = r_list[0] |
| prompt = prompt_item.get("content", "") if isinstance(prompt_item, dict) else str(prompt_item) |
|
|
| if not prompt and "prompt" in example: |
| prompt = example["prompt"] |
|
|
| if prompt and chosen and rejected: |
| return { |
| "prompt": prompt.strip(), |
| "chosen": chosen.strip(), |
| "rejected": rejected.strip() |
| } |
| return None |
|
|
| def main(): |
| print(f"Downloading RLHF preference dataset: {config.DATASET_ID}...") |
| dataset = load_dataset(config.DATASET_ID, split="train") |
|
|
| print(f"Loaded {len(dataset)} raw samples. Filtering and reformatting DPO preference pairs...") |
| formatted_items = [] |
| |
| for example in dataset: |
| item = extract_prompt_chosen_rejected(example) |
| if item and len(item["prompt"]) > 10 and len(item["chosen"]) > 20 and len(item["rejected"]) > 20: |
| formatted_items.append(item) |
| if len(formatted_items) >= 10000: |
| break |
|
|
| print(f"Successfully processed {len(formatted_items)} valid preference pairs.") |
|
|
| val_size = min(1000, int(len(formatted_items) * 0.1)) |
| train_size = len(formatted_items) - val_size |
|
|
| train_items = formatted_items[:train_size] |
| val_items = formatted_items[train_size:] |
|
|
| os.makedirs(config.DATA_DIR, exist_ok=True) |
|
|
| print(f"Saving {len(train_items)} training pairs to {config.TRAIN_FILE}...") |
| with open(config.TRAIN_FILE, "w", encoding="utf-8") as f: |
| for item in train_items: |
| f.write(json.dumps(item) + "\n") |
|
|
| print(f"Saving {len(val_items)} validation pairs to {config.VAL_FILE}...") |
| with open(config.VAL_FILE, "w", encoding="utf-8") as f: |
| for item in val_items: |
| f.write(json.dumps(item) + "\n") |
|
|
| print("RLHF Preference dataset preparation complete!") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
