| 79 | return {"system": DEFAULT_S2M_PROMPT, "messages": messages, "audio": audios} |
| 80 | |
| 81 | def main(): |
| 82 | parser = argparse.ArgumentParser(description="Process audio dataset for training") |
| 83 | parser.add_argument("--debug", action="store_true", help="Debug mode, only process limited samples") |
| 84 | parser.add_argument("--datapath", type=str, default="datasets/spoken-alpaca-gpt4", help="Path to dataset") |
| 85 | parser.add_argument("--num_workers", type=int, default=None, help="Number of workers (default: torch.cuda.device_count() * 4)") |
| 86 | parser.add_argument("--num_samples", type=int, default=1000, help="Number of samples to process in debug mode") |
| 87 | parser.add_argument("--delay_token_nums", type=int, default=10, help="Number of delay tokens") |
| 88 | args = parser.parse_args() |
| 89 | |
| 90 | set_start_method("spawn") |
| 91 | |
| 92 | num_workers = args.num_workers if args.num_workers is not None else torch.cuda.device_count() * 4 |
| 93 | |
| 94 | datapath = args.datapath |
| 95 | audio_dir = os.path.join(datapath, "audio") |
| 96 | os.makedirs(audio_dir, exist_ok=True) |
| 97 | |
| 98 | dataset = datasets.load_from_disk(datapath) |
| 99 | if args.debug: |
| 100 | dataset = dataset.select(range(args.num_samples)) |
| 101 | |
| 102 | dataset = dataset.map( |
| 103 | convert_format, |
| 104 | with_rank=True, |
| 105 | remove_columns=dataset.column_names, |
| 106 | num_proc=num_workers, |
| 107 | fn_kwargs={"audio_dir": audio_dir, "delay_token_nums": args.delay_token_nums}, |
| 108 | desc="Converting format of dataset", |
| 109 | ) |
| 110 | dataset = dataset.remove_columns([ dn for dn in dataset.column_names if dn not in ["system", "messages", "audio"]]) |
| 111 | |
| 112 | output_filename = "train.jsonl" if not args.debug else f"train_{args.num_samples//1000}k.jsonl" |
| 113 | output_path = os.path.join(datapath, output_filename) |
| 114 | dataset.to_json(output_path, orient="records", lines=True) |
| 115 | print(f"Saved to: {output_path}") |
| 116 | |
| 117 | if __name__ == "__main__": |
| 118 | main() |