| 21 | import argparse |
| 22 | |
| 23 | def main(args): |
| 24 | |
| 25 | tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, use_fast=False, trust_remote_code=True) |
| 26 | |
| 27 | def tokenize_function(examples): |
| 28 | output = tokenizer(examples[args.text_field]) |
| 29 | return output |
| 30 | |
| 31 | block_size = args.block_size |
| 32 | # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size. |
| 33 | def group_texts(examples): |
| 34 | # Concatenate all texts. |
| 35 | concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()} |
| 36 | total_length = len(concatenated_examples[list(examples.keys())[0]]) |
| 37 | # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can |
| 38 | # customize this part to your needs. |
| 39 | if total_length >= block_size: |
| 40 | total_length = (total_length // block_size) * block_size |
| 41 | result = { |
| 42 | k: [t[i : i + block_size] for i in range(0, total_length, block_size)] |
| 43 | for k, t in concatenated_examples.items() |
| 44 | } |
| 45 | result["labels"] = result["input_ids"].copy() |
| 46 | return result |
| 47 | |
| 48 | filename = '.'.join(args.input_file.split("/")[-1].split(".")[:-1]) |
| 49 | os.makedirs(args.output_dir, exist_ok=True) |
| 50 | cache_dir = os.path.join(args.output_dir, filename) |
| 51 | tmp_cache_dir = os.path.join(args.output_dir, filename+"_text") |
| 52 | |
| 53 | if args.data_type == "jsonl": |
| 54 | raw_dataset = load_dataset("json", data_files=args.input_file, cache_dir=tmp_cache_dir, keep_in_memory=False, encoding="utf8") |
| 55 | elif args.data_type == 'text': |
| 56 | raw_dataset = load_dataset("text", data_files=args.input_file, cache_dir=tmp_cache_dir, keep_in_memory=False, encoding="utf8") |
| 57 | else: |
| 58 | raise NotImplementedError(f"data type should be in json,txt not {args.data_type}") |
| 59 | |
| 60 | print("remove_column_names:", raw_dataset.column_names['train']) |
| 61 | tokenized_dataset = raw_dataset.map( |
| 62 | tokenize_function, |
| 63 | batched=True, |
| 64 | num_proc=args.preprocessing_num_workers, |
| 65 | remove_columns=raw_dataset.column_names['train'], |
| 66 | load_from_cache_file=True, |
| 67 | keep_in_memory=False, |
| 68 | cache_file_names = {k: os.path.join(tmp_cache_dir, 'tokenized.arrow') for k in raw_dataset}, |
| 69 | desc="Running tokenizer on dataset", |
| 70 | ) |
| 71 | if args.filter_by_length is not None: |
| 72 | tokenized_dataset["train"] = tokenized_dataset["train"].filter( |
| 73 | lambda x: len(x["input_ids"]) >= args.filter_by_length |
| 74 | ) |
| 75 | grouped_datasets = tokenized_dataset.map( |
| 76 | group_texts, |
| 77 | batched=True, |
| 78 | num_proc=args.preprocessing_num_workers, |
| 79 | load_from_cache_file=True, |
| 80 | keep_in_memory=False, |