| from datasets import Dataset, DatasetDict |
| from transformers import AutoTokenizer |
|
|
| def preprocess_dataset(dataset: Dataset | DatasetDict, tokenizer: AutoTokenizer) -> Dataset | DatasetDict: |
|
|
| tokenized_dataset = dataset.map( |
| lambda examples: tokenize_function(examples, tokenizer), batched=True, remove_columns=["text", 'text_clean', 'language'] |
| ) |
|
|
| return tokenized_dataset.map(group_texts, batched=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def tokenize_function(examples, tokenizer: AutoTokenizer): |
|
|
| result = tokenizer(examples["text"]) |
|
|
| if tokenizer.is_fast: |
| result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))] |
| |
| return result |
|
|
|
|
| def group_texts(examples, chunk_size: int = 128): |
|
|
| concatinated_examples = {k : sum(examples[k], []) for k in examples.keys()} |
|
|
| total_length = len(concatinated_examples["input_ids"]) |
| total_length = (total_length // chunk_size) * chunk_size |
|
|
| result = {k : [t[i : i+chunk_size] for i in range(0, total_length, chunk_size)] for k, t in concatinated_examples.items()} |
| result["labels"] = result["input_ids"].copy() |
|
|
| return result |
|
|
|
|
