"""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)