(
*,
dataset: str,
configuration: str | None,
split: str,
samples: int | None,
seed: int,
max_seq_len: int,
tokenizer: transformers.PreTrainedTokenizerBase,
add_extra_pad_token: bool = False,
apply_chat_template: bool = False,
)
| 92 | |
| 93 | |
| 94 | def load_preprocess_data( |
| 95 | *, |
| 96 | dataset: str, |
| 97 | configuration: str | None, |
| 98 | split: str, |
| 99 | samples: int | None, |
| 100 | seed: int, |
| 101 | max_seq_len: int, |
| 102 | tokenizer: transformers.PreTrainedTokenizerBase, |
| 103 | add_extra_pad_token: bool = False, |
| 104 | apply_chat_template: bool = False, |
| 105 | ): |
| 106 | ds = datasets.load_dataset(dataset, name=configuration, split=split) |
| 107 | ds = ds.shuffle(seed=seed) |
| 108 | if samples is not None: |
| 109 | ds = ds.select(range(samples)) |
| 110 | if tokenizer.pad_token_id is None: |
| 111 | tokenizer.pad_token_id = tokenizer.eos_token_id |
| 112 | if apply_chat_template: |
| 113 | ds = ds.map( |
| 114 | lambda x: { |
| 115 | "text": do_chat_template(x, tokenizer), |
| 116 | }, |
| 117 | num_proc=64, |
| 118 | ) |
| 119 | ds = ds.filter(lambda row: row["text"] and row["text"].strip(), num_proc=64) |
| 120 | ds = ds.map( |
| 121 | lambda x: { |
| 122 | "input_ids": truncate_tokens( |
| 123 | x["text"], tokenizer, max_seq_len, add_extra_pad_token |
| 124 | ) |
| 125 | }, |
| 126 | num_proc=64, |
| 127 | ).filter(lambda x: len(x["input_ids"]) > 0, num_proc=64) |
| 128 | return ds |
| 129 | |
| 130 | |
| 131 | def truncate_tokens( |
no test coverage detected