Spaces:
Sleeping
Sleeping
File size: 1,144 Bytes
3fbe048 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """Loads the TweetEval irony split and encodes it for DistilBERT."""
from datasets import DatasetDict, load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding
from training.config import BASE_CHECKPOINT, DATASET_CONFIG, DATASET_ID, MAX_LENGTH
def load_splits() -> DatasetDict:
"""Returns the official train/validation/test splits, unmodified."""
return load_dataset(DATASET_ID, DATASET_CONFIG)
def build_tokenizer(checkpoint: str = BASE_CHECKPOINT):
return AutoTokenizer.from_pretrained(checkpoint)
def encode(splits: DatasetDict, tokenizer) -> DatasetDict:
"""
Tokenises without padding so the collator can pad each batch to its own
longest sequence: on CPU, padding everything to 96 tokens wastes most of the
attention computation on a corpus whose median tweet is under 30.
"""
def apply(batch: dict) -> dict:
return tokenizer(batch["text"], truncation=True, max_length=MAX_LENGTH)
return splits.map(apply, batched=True, remove_columns=["text"])
def build_collator(tokenizer) -> DataCollatorWithPadding:
return DataCollatorWithPadding(tokenizer=tokenizer)
|