(args: argparse.Namespace)
| 105 | |
| 106 | |
| 107 | def load_source_dataset(args: argparse.Namespace): |
| 108 | from datasets import load_dataset |
| 109 | |
| 110 | dataset = load_dataset(args.dataset_name, split=args.split) |
| 111 | if args.sample_size is not None and args.sample_size < len(dataset): |
| 112 | dataset = dataset.select(range(args.sample_size)) |
| 113 | dataset = dataset.map(add_index, with_indices=True) |
| 114 | return dataset |
| 115 | |
| 116 | |
| 117 | def normalize_conversations(row: dict) -> dict: |