utils_dataset.py

import datasets

DEFAULT_HF_DATASETS_CACHE_DIR = "./dataset/hf_dataset_cache_dir/"

def get_dataset_multi30k(cache_dir: str = DEFAULT_HF_DATASETS_CACHE_DIR) -> datasets.dataset_dict.DatasetDict:
    return datasets.load_dataset("bentrevett/multi30k", cache_dir=cache_dir)


def encode_multi30k(examples, tokenizer, max_seq_len: int, strip_case: bool = False):
    """Given examples from a dataset, tokenize them (with padding+truncation).

    Args:
        examples (_type_): _description_
        tokenizer (_type_): _description_
        max_seq_len (int): _description_

    Returns:
        _type_: _description_
    """
    text_en = examples["en"]
    text_de = examples["de"]
    if strip_case:
        text_en = [text.lower() for text in text_en]
        text_de = [text.lower() for text in text_de]
    return tokenizer(
        text=text_en,
        text_target=text_de,
        truncation=True,
        padding="max_length",
        max_length=max_seq_len,
    )


def preprocess_dataset_multi30k(
        ds: datasets.dataset_dict.DatasetDict,
        tokenizer,
        max_seq_len: int,
        strip_case: bool = False,
    ) -> tuple[datasets.arrow_dataset.Dataset, datasets.arrow_dataset.Dataset, datasets.arrow_dataset.Dataset]:
    """Preprocesses the multi30k dataset, by doing:
    - tokenization (and padding+truncation)
    - convert to torch.Tensor

    Args:
        ds (datasets.dataset_dict.DatasetDict): _description_
        tokenizer (_type_): _description_
        max_seq_len (int): _description_

    Returns:
        tuple[datasets.arrow_dataset.Dataset, datasets.arrow_dataset.Dataset, datasets.arrow_dataset.Dataset]:
            ds_train:
            ds_val:
            ds_test:
    """
    ds_train = ds["train"]
    ds_val = ds["validation"]
    ds_test = ds["test"]

    return (
        preprocess_dataset_split_multi30k(ds_train, tokenizer, max_seq_len, strip_case=strip_case),
        preprocess_dataset_split_multi30k(ds_val, tokenizer, max_seq_len, strip_case=strip_case),
        preprocess_dataset_split_multi30k(ds_test, tokenizer, max_seq_len, strip_case=strip_case),
    )

def preprocess_dataset_split_multi30k(
        ds_split: datasets.arrow_dataset.Dataset,
        tokenizer,
        max_seq_len: int,
        strip_case: bool = False,
    ) -> datasets.arrow_dataset.Dataset:
    """Preprocesses a dataset split of multi30k.

    Args:
        ds_split (datasets.arrow_dataset.Dataset): _description_
        tokenizer (_type_): _description_
        max_seq_len (int): _description_

    Returns:
        datasets.arrow_dataset.Dataset: _description_
    """
    ds_tokenized = ds_split.map(lambda examples: encode_multi30k(examples, tokenizer, max_seq_len=max_seq_len, strip_case=strip_case), batched=True)
    # normally, we wouldn't select the raw text 'en', 'de', but to make visualization/debugging
    # easier we'll do this
    # input_ids.shape=[]
    dataset = ds_tokenized.select_columns(["input_ids", "attention_mask", "labels", "en", "de"])
    dataset = dataset.with_format(type="torch")
    # input_ids: torch.Tensor with shape=[max_seq_len], dtype=torch.int64
    #   Token ids for source sequence.
    # attention_mask: torch.Tensor with shape=[max_seq_len], dtype=torch.int64
    #   This is the padding mask.
    #   `1` means "valid", `0` means invalid (ie [PAD] token)
    # labels: torch.Tensor with shape=[max_seq_len], dtype=torch.int64
    #   Token ids for target sequence.
    # en: str.
    # de: str.
    return dataset


def visualize_translation_dataset(ds, tokenizer, n: int = 3):
    """Visualizes a dataset.
    Must be already preprocessed, ie output by `preprocess_dataset_multi30k()`.

    Args:
        ds (_type_): _description_
    """
    # print a few dataset examples
    for ind in range(min(n, len(ds))):
        input_ids = ds[ind]["input_ids"]
        attention_mask = ds[ind]["attention_mask"]
        labels = ds[ind]["labels"]

        input_ids_decoded = tokenizer.decode(input_ids, skip_special_tokens=True)
        labels_decoded = tokenizer.decode(labels, skip_special_tokens=True)

        # check that attention_mask is correct: should strip away all [PAD] output
        input_ids_decoded_2 = tokenizer.decode(input_ids[attention_mask == 1])

        print(f"[ind={ind}]")
        print(f"  input_ids={input_ids}")
        print(f"  attention_mask={attention_mask}")
        print(f"  labels={labels}")
        print(f"  input_ids_decoded={input_ids_decoded}")
        print(f"  labels_decoded={labels_decoded}")
        print(f"  (debug) input_ids_decoded_2={input_ids_decoded_2}")


def get_dataset_wmt14_en_fr(cache_dir: str = DEFAULT_HF_DATASETS_CACHE_DIR) -> datasets.dataset_dict.DatasetDict:
    return datasets.load_dataset("wmt/wmt14", "fr-en", cache_dir=cache_dir)