| |
| """ |
| Preprocess the llamafactory/pokemon-gpt4o-captions dataset to parquet format |
| """ |
|
|
| import argparse |
| import os |
|
|
| import datasets |
|
|
| from verl.utils.hdfs_io import copy, makedirs |
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--local_dir", default=None) |
| parser.add_argument("--hdfs_dir", default=None) |
| parser.add_argument("--local_dataset_path", default=None, help="The local path to the raw dataset, if it exists.") |
| parser.add_argument( |
| "--local_save_dir", |
| default="~/data/pokemon-gpt4o-captions", |
| help="The save directory for the preprocessed dataset.", |
| ) |
|
|
| args = parser.parse_args() |
| local_dataset_path = args.local_dataset_path |
|
|
| data_source = "llamafactory/pokemon-gpt4o-captions" |
|
|
| if local_dataset_path is not None: |
| dataset = datasets.load_dataset( |
| local_dataset_path, |
| ) |
| else: |
| dataset = datasets.load_dataset( |
| data_source, |
| ) |
|
|
| def map_fn(row: dict): |
| messages = [] |
| conversation = row.pop("conversations") |
| for conv in conversation: |
| if conv["from"] == "gpt": |
| role = "assistant" |
| elif conv["from"] == "human": |
| role = "user" |
| else: |
| raise ValueError(f"Unknown role: {conv['from']}") |
| messages.append( |
| { |
| "role": role, |
| "content": conv["value"], |
| } |
| ) |
|
|
| row["messages"] = messages |
| return row |
|
|
| dataset = dataset["train"].map(map_fn, num_proc=16) |
| dataset = dataset.train_test_split(test_size=0.1) |
| train_dataset = dataset["train"] |
| test_dataset = dataset["test"] |
|
|
| hdfs_dir = args.hdfs_dir |
| local_save_dir = args.local_dir |
| if local_save_dir is not None: |
| print("Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.") |
| else: |
| local_save_dir = args.local_save_dir |
|
|
| train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet")) |
| test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet")) |
|
|
| if hdfs_dir is not None: |
| makedirs(hdfs_dir) |
| copy(src=local_save_dir, dst=hdfs_dir) |
|
|
