Load and preprocess the dataset.
(tokenizer, args)
| 236 | |
| 237 | |
| 238 | def load_and_preprocess_cr(tokenizer, args): |
| 239 | """Load and preprocess the dataset.""" |
| 240 | if args.data_path.endswith(".json"): |
| 241 | data = load_dataset("json", data_files=args.data_path) |
| 242 | else: |
| 243 | data = load_dataset(args.data_path) |
| 244 | |
| 245 | # Create a wrapper function that includes all necessary arguments |
| 246 | def generate_and_tokenize_prompt_wrapper(data_point): |
| 247 | return generate_and_tokenize_prompt_cr(data_point, tokenizer, args) |
| 248 | |
| 249 | train_dataset = data["train"].shuffle().map( |
| 250 | generate_and_tokenize_prompt_wrapper, |
| 251 | num_proc=8, |
| 252 | remove_columns=data["train"].column_names # Remove original columns |
| 253 | ) |
| 254 | |
| 255 | return train_dataset |
| 256 | |
| 257 | def generate_prompt_cr(data_point): |
| 258 | """Generate prompt from data point.""" |