scandl_module/scripts/sp_load_celer_zuco.py

import pandas as pd
import numpy as np
from tqdm import tqdm
import os
import random
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import Dataset as Dataset2
import datasets
from sklearn.model_selection import (
    train_test_split,
    GroupShuffleSplit,
    KFold,
    GroupKFold,
    StratifiedKFold,
)
from typing import Optional, List, Tuple, Union, Any, Dict
import sys

sys.path.append("./")
sys.path.append("../")
sys.path.append("../../")

from ScanDL2.CONSTANTS import PATH_TO_IA, PATH_TO_FIX, SUB_METADATA_PATH, path_to_zuco


def load_celer():
    path_to_fix = PATH_TO_FIX
    path_to_ia = PATH_TO_IA
    eyemovement_df = pd.read_csv(path_to_fix, delimiter="\t", low_memory=False)
    eyemovement_df["CURRENT_FIX_INTEREST_AREA_LABEL"] = (
        eyemovement_df.CURRENT_FIX_INTEREST_AREA_LABEL.replace("\t(.*)", "", regex=True)
    )
    word_info_df = pd.read_csv(path_to_ia, delimiter="\t")
    word_info_df["IA_LABEL"] = word_info_df.IA_LABEL.replace("\t(.*)", "", regex=True)
    return word_info_df, eyemovement_df


def load_celer_speakers(only_native_speakers: bool = True):
    sub_metadata_path = SUB_METADATA_PATH
    sub_info = pd.read_csv(sub_metadata_path, delimiter="\t")
    if only_native_speakers:
        readers_list = sub_info[sub_info.L1 == "English"].List.values
    else:
        readers_list = sub_info.List.values
    return readers_list.tolist()


def compute_word_length(arr):
    # length of a punctuation is 0, plus an epsilon to avoid division output inf
    arr = arr.astype("float64")
    arr[arr == 0] = 1 / (0 + 0.5)
    arr[arr != 0] = 1 / (arr[arr != 0])
    return arr


def compute_word_frequency(arr):
    arr[arr == np.inf] = np.nan
    arr[arr != np.inf] = np.log10(arr[arr != np.inf])
    return arr


def _collate_instance_helper(
    instance,
    pad_token_id,
    padding_steps,  # how many steps/dims to pad,
):
    padding_list = [pad_token_id] * padding_steps
    result = instance + padding_list
    return result


def _collate_batch_helper(
    examples,  # List of Lists of input IDs
    pad_token_id,
    max_length,
    return_mask=False,
):
    result = torch.full([len(examples), max_length], pad_token_id, dtype=torch.int64).tolist()
    mask_ = torch.full([len(examples), max_length], pad_token_id, dtype=torch.int64).tolist()
    for i, example in enumerate(examples):
        curr_len = min(len(example), max_length)
        result[i][:curr_len] = example[:curr_len]
        mask_[i][:curr_len] = [1] * curr_len
    if return_mask:
        return result, mask_
    return result


def _dummy_pad_words(
    examples,
    pad_token_id,
    max_length,
):
    padded_examples = list()
    for instance in examples:
        padded_examples.append(instance + (max_length - len(instance)) * [pad_token_id])
    return padded_examples


def infinite_loader(data_loader):
    while True:
        yield from data_loader


def celer_zuco_dataset_and_loader(
    data,
    data_args,
    split: str,
    deterministic=False,
    loop=True,
):

    dataset = CelerZucoDataset(
        dataset=data,
        data_args=data_args,
        split=split,
    )
    data_loader = DataLoader(
        dataset,
        batch_size=data_args.batch_size,
        shuffle=not deterministic,  # ??
        num_workers=0,
    )
    if loop:
        return infinite_loader(data_loader)
    else:
        return iter(data_loader)


def combined_split(data, reader_IDs, sn_IDs, test_size):
    """Splits the data so that the test data contains both unseen readers and unseen sentences."""
    random.seed(77)
    unique_reader_IDs = set(reader_IDs)
    unique_sn_IDs = set(sn_IDs)

    # sample the sentence and reader IDs that go into the test set
    unique_reader_IDs_test = random.sample(
        unique_reader_IDs, int(test_size * len(unique_reader_IDs))
    )
    unique_sn_IDs_test = random.sample(unique_sn_IDs, int(test_size * len(unique_sn_IDs)))

    train_data, test_data = [], []
    train_reader_IDs, test_reader_IDs = [], []
    train_sn_IDs, test_sn_IDs = [], []

    for i in range(len(data)):
        if reader_IDs[i] in unique_reader_IDs_test and sn_IDs[i] in unique_sn_IDs_test:
            test_data.append(data[i])
            test_reader_IDs.append(reader_IDs[i])
            test_sn_IDs.append(sn_IDs[i])
        elif reader_IDs[i] not in unique_reader_IDs_test and sn_IDs[i] not in unique_sn_IDs_test:
            train_data.append(data[i])
            train_reader_IDs.append(reader_IDs[i])
            train_sn_IDs.append(sn_IDs[i])
        else:
            continue

    return train_data, test_data, train_reader_IDs, test_reader_IDs, train_sn_IDs, test_sn_IDs


def load_zuco(task: str = None):  # 'zuco11', 'zuco12'
    dir = path_to_zuco
    if task.startswith("zuco1"):
        dir = dir + "zuco/"
    elif task == "zuco21":
        dir = dir + "zuco2/"
    dir = os.path.join(dir, f"task{task[-1]}", "Matlab_files")
    word_info_path = dir + "/Word_Infor.csv"
    word_info_df = pd.read_csv(word_info_path, sep="\t")
    scanpath_path = dir + "/scanpath.csv"
    eyemovement_df = pd.read_csv(scanpath_path, sep="\t")
    return word_info_df, eyemovement_df


def get_kfold_indices_scanpath(
    splitting_IDs_dict: Dict[str, Union[int, str]],
    n_splits: int = 5,
):
    """Function to implement the 'random'/'scanpath' split, in which the test set contains both sentences and
    readers that were seen during training. Unfortunately, it is not possible to control for the ratios of both
    readers and sentences (as they are co-dependent), so I will make sure that the data is shuffled and at least
    the readers are stratified."""

    reader_IDs = splitting_IDs_dict["reader"]
    sn_IDs = splitting_IDs_dict["sentence"]

    tuple_ids = [
        (idx, reader_ID, sn_ID) for idx, (reader_ID, sn_ID) in enumerate(zip(reader_IDs, sn_IDs))
    ]

    # list of indices of the sn-reader pairs of unique sentences
    unique_sns_ids = [idx for (idx, reader_ID, sn_ID) in tuple_ids if not sn_ID.startswith("en")]
    # unique_sns_indices = {sn_ID: idx for (idx, reader_ID, sn_ID) in tuple_ids if not sn_ID.startswith('en')}
    universal_sn_ids = [idx for (idx, reader_ID, sn_ID) in tuple_ids if sn_ID.startswith("en")]

    # get the reader IDs for the unique and universal sns
    unique_reader_ids = [
        reader_ID for (idx, reader_ID, sn_ID) in tuple_ids if not sn_ID.startswith("en")
    ]
    universal_reader_ids = [
        reader_ID for (idx, reader_ID, sn_ID) in tuple_ids if sn_ID.startswith("en")
    ]

    # stratify both the readers for the universal sentences and the ones for the unique sentences
    kfold_unique = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=77)
    kfold_universal = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=77)
    train_idx_unique, test_idx_unique = list(), list()
    train_idx_universal, test_idx_universal = list(), list()
    kfold_unique.get_n_splits(unique_sns_ids, groups=unique_reader_ids)
    kfold_universal.get_n_splits(universal_sn_ids, groups=universal_reader_ids)
    for train_idx, test_idx in kfold_unique.split(
        X=unique_sns_ids, y=unique_reader_ids, groups=unique_reader_ids
    ):
        train_idx_unique.append(train_idx)
        test_idx_unique.append(test_idx)
    for train_idx, test_idx in kfold_universal.split(
        X=universal_sn_ids, y=universal_reader_ids, groups=universal_reader_ids
    ):
        train_idx_universal.append(train_idx)
        test_idx_universal.append(test_idx)

    # concatenate the respective train and test idx together
    all_train_idx, all_test_idx = list(), list()

    for train_idx_univ, train_idx_uniq in zip(train_idx_universal, train_idx_unique):
        train_idx = np.concatenate([train_idx_univ, train_idx_uniq])
        all_train_idx.append(train_idx)
    for test_idx_univ, test_idx_uniq in zip(test_idx_universal, test_idx_unique):
        test_idx = np.concatenate([test_idx_univ, test_idx_uniq])
        all_test_idx.append(test_idx)

    all_idx = [(train_idx, test_idx) for train_idx, test_idx in zip(all_train_idx, all_test_idx)]
    return all_idx


def get_kfold(
    data: List[Tuple[Any]],
    splitting_IDs_dict: Dict[str, Union[int, str]],
    splitting_criterion: str = "scanpath",  # 'scanpath', 'reader', 'sentence', 'combined',
    n_splits: int = 5,
):
    if splitting_criterion == "scanpath":
        # kfold = KFold(n_splits=n_splits, random_state=77, shuffle=True)
        # kfold.get_n_splits(data)
        # return kfold.split(data)
        return get_kfold_indices_scanpath(
            splitting_IDs_dict=splitting_IDs_dict,
            n_splits=n_splits,
        )
    elif splitting_criterion in ["reader", "sentence"]:
        splitting_group = splitting_IDs_dict[splitting_criterion]
        kfold = GroupKFold(n_splits=n_splits)
        kfold.get_n_splits(data, groups=splitting_group)
        return kfold.split(data, groups=splitting_group)
    else:  # combined split
        raise NotImplementedError


def get_kfold_indices_combined(
    data: List[Tuple[Any]],
    splitting_IDs_dict: Dict[str, Union[int, str]],
    n_splits: int = 5,
):
    kfold_reader = GroupKFold(n_splits=n_splits)
    kfold_sentence = GroupKFold(n_splits=n_splits)
    kfold_reader.get_n_splits(data, groups=splitting_IDs_dict["reader"])
    kfold_sentence.get_n_splits(data, groups=splitting_IDs_dict["sentence"])
    reader_indices, sentence_indices = list(), list()
    for train_idx, test_idx in kfold_reader.split(data, groups=splitting_IDs_dict["reader"]):
        reader_indices.append((train_idx, test_idx))
    for train_idx, test_idx in kfold_sentence.split(data, groups=splitting_IDs_dict["sentence"]):
        sentence_indices.append((train_idx, test_idx))
    return reader_indices, sentence_indices


def flatten_data(data: Dict[str, List[Any]]):
    flattened_data = list()
    for i in range(len(data["sn_sp_repr"])):
        flattened_data.append(
            (
                data["mask"][i],
                data["sn_sp_repr"][i],
                data["sn_input_ids"][i],
                data["indices_pos_enc"][i],
                data["sn_sp_fix_dur"][i],
                data["sn_repr_len"][i],
                data["words_for_mapping"][i],
                data["mask_sn_padding"][i],
                data["mask_transformer_att"][i],
                data["sn_ids"][i],
                data["reader_ids"][i],
            )
        )
    return flattened_data


def unflatten_data(flattened_data: List[Tuple[Any]], split: str):
    dataset = Dataset2.from_dict(
        {
            "mask": [instance[0] for instance in flattened_data],
            "sn_sp_repr": [instance[1] for instance in flattened_data],
            "sn_input_ids": [instance[2] for instance in flattened_data],
            "indices_pos_enc": [instance[3] for instance in flattened_data],
            "sn_sp_fix_dur": [instance[4] for instance in flattened_data],
            "sn_repr_len": [instance[5] for instance in flattened_data],
            "words_for_mapping": [instance[6] for instance in flattened_data],
            "mask_sn_padding": [instance[7] for instance in flattened_data],
            "mask_transformer_att": [instance[8] for instance in flattened_data],
            "sn_ids": [instance[9] for instance in flattened_data],
            "reader_ids": [instance[10] for instance in flattened_data],
        }
    )
    data_dict = datasets.DatasetDict()
    data_dict[split] = dataset
    return data_dict


def process_celer(
    sn_list,
    reader_list,
    word_info_df,
    eyemovement_df,
    tokenizer,
    args,
    split: Optional[str] = "train",
    subset_size: Optional[int] = None,
    split_sizes: Optional[Dict[str, float]] = None,
    splitting_criterion: Optional[str] = "scanpath",  # 'reader', 'sentence', 'combined'
    inference: Optional[str] = None,  # cv, zuco
):
    """
    Process the Celer corpus so that it can be used as input to the Diffusion model, where the original sentence (sn)
    is the condition and the scan path (sp) is the target that will be noised.
    :param sn_list: list of unique sentence IDs in celer
    :param reader_list: list of reader IDs in celer
    :param word_info_df: pd Dataframe with sentence info
    :param eyemovement_df: pd Dataframe with fixation info
    :param tokenizer: BertTokenizer
    :param args:
    :param split: 'train', 'train-test', 'train-test-val', 'train-val'
    :param subset_size: for test runs: to not load whole dataset but specified no. of instances
    :param split_sizes: proportion of data going into train, test and val
    :param splitting_criterion: how the data should be split for testing and validation.
        'reader' = New Reader setting
        'sentence' = New Sentence setting
        'combined' = New Reader/New Sentence setting
        'scanpath' = data is split at random
    :param inference: if inference is cross-validation, the data is returned before splitting into train test val
    """
    SP_ordinal_pos = []
    SP_landing_pos = []
    SP_fix_dur = []

    data = {
        "mask": list(),  # 0 for sn, 1 for sp
        "sn_sp_repr": list(),  # word IDs of sn and corresponding word IDs of sp (fixated words, interest area IDs) padded with args.seq_len -1
        "sn_input_ids": list(),  # input IDs of tokenized sentence, padded with pad token ID
        "indices_pos_enc": list(),  # indices from 1 ... len(sn input ids) 1 ... (seq_len - len(sn input ids))
        "sn_repr_len": list(),  # length of sentence in subword tokens
        "words_for_mapping": list(),  # original words of sentence, padded with PAD
        "mask_sn_padding": list(),  # masks both the sentence and the padding, for the loss computations
        "mask_transformer_att": list(),  # masks only the padding, for the transformer attention,
        "sn_ids": list(),  # the sentence IDs
        "reader_ids": list(),  # the reader IDs
    }

    max_len = 0

    reader_IDs, sn_IDs = list(), list()

    for sn_id_idx, sn_id in tqdm(enumerate(sn_list), total=len(sn_list)):  # for text/sentence ID

        if subset_size is not None:
            if sn_id_idx == subset_size + 1:
                break

        # subset the fixations report DF to a DF containing only the current sentence/text ID (each sentence appears multiple times)
        sn_df = eyemovement_df[eyemovement_df.sentenceid == sn_id]
        # notice: Each sentence is recorded multiple times in file |word_info_df|.
        # subset the interest area report DF to a DF containing only the current sentence/text ID
        sn = word_info_df[word_info_df.sentenceid == sn_id]
        # sn is a dataframe containing only one sentence (the sentence with the current sentence ID)
        sn = sn[
            sn["list"] == sn.list.values.tolist()[0]
        ]  # list = experimental list number (unique to each participant).
        # compute word length and frequency features for each word
        sn_str = sn.sentence.iloc[-1]  # the whole sentence as string
        if (
            sn_id == "1987/w7_019/w7_019.295-3"
            or sn_id == "1987/w7_036/w7_036.147-43"
            or sn_id == "1987/w7_091/w7_091.360-6"
        ):
            # extra inverted commas at the end of the sentence
            sn_str = sn_str[:-3] + sn_str[-1:]
        if sn_id == "1987/w7_085/w7_085.200-18":
            sn_str = sn_str[:43] + sn_str[44:]

        # skip nan values bc they are of type float (np.isnan raises an error)
        if isinstance(sn_str, float):
            continue

        sn_len = len(sn_str.split())

        # add CLS and SEP 'manually' to the sentence so that they receive the word IDs 0 and len(sn)+1
        sn_str = "[CLS] " + sn_str + " [SEP]"

        tokenizer.padding_side = "right"

        for sub_id_idx, sub_id in enumerate(reader_list):

            if sub_id_idx == 5:
                continue

            sub_df = sn_df[sn_df.list == sub_id]
            # remove fixations on non-words
            sub_df = sub_df.loc[
                sub_df.CURRENT_FIX_INTEREST_AREA_LABEL != "."
            ]  # Label for the interest area to which the currentixation is assigned
            if len(sub_df) == 0:
                # no scanpath data found for the subject
                continue

            # prepare decoder input and output
            sp_word_pos, sp_fix_loc, sp_fix_dur = (
                sub_df.CURRENT_FIX_INTEREST_AREA_ID.values,
                sub_df.CURRENT_FIX_NEAREST_INTEREST_AREA_DISTANCE.values,
                sub_df.CURRENT_FIX_DURATION.values,
            )

            # check if recorded fixation duration are within reasonable limits
            # Less than 15ms attempt to merge with neighbouring fixation if fixate is on the same word, otherwise delete
            outlier_indx = np.where(sp_fix_dur < 50)[
                0
            ]  # gives indices of the fixations in the fixations list that were shorter than 50ms

            if outlier_indx.size > 0:
                for out_idx in range(len(outlier_indx)):
                    outlier_i = outlier_indx[out_idx]
                    merge_flag = False

                    # outliers are commonly found in the fixation of the last record and the first record, and are removed directly
                    if outlier_i == len(sp_fix_dur) - 1 or outlier_i == 0:
                        merge_flag = True

                    else:
                        if outlier_i - 1 >= 0 and not merge_flag:
                            # try to merge with the left fixation if they landed both on the same interest area
                            if (
                                sub_df.iloc[outlier_i].CURRENT_FIX_INTEREST_AREA_LABEL
                                == sub_df.iloc[outlier_i - 1].CURRENT_FIX_INTEREST_AREA_LABEL
                            ):
                                sp_fix_dur[outlier_i - 1] = (
                                    sp_fix_dur[outlier_i - 1] + sp_fix_dur[outlier_i]
                                )
                                merge_flag = True

                        if outlier_i + 1 < len(sp_fix_dur) and not merge_flag:
                            # try to merge with the right fixation
                            if (
                                sub_df.iloc[outlier_i].CURRENT_FIX_INTEREST_AREA_LABEL
                                == sub_df.iloc[outlier_i + 1].CURRENT_FIX_INTEREST_AREA_LABEL
                            ):
                                sp_fix_dur[outlier_i + 1] = (
                                    sp_fix_dur[outlier_i + 1] + sp_fix_dur[outlier_i]
                                )
                                merge_flag = True

                    # delete the position (interest area ID), the fixation location and the fixation duration from the respective arrays
                    sp_word_pos = np.delete(sp_word_pos, outlier_i)
                    sp_fix_loc = np.delete(sp_fix_loc, outlier_i)
                    sp_fix_dur = np.delete(sp_fix_dur, outlier_i)
                    sub_df.drop(sub_df.index[outlier_i], axis=0, inplace=True)
                    outlier_indx = outlier_indx - 1

            # sanity check
            # scanpath too long, remove outliers, speed up the inference; more than 50 fixations on sentence
            if len(sp_word_pos) > 50:  # 72/10684
                continue
            # scanpath too short for a normal length sentence
            if len(sp_word_pos) <= 1 and sn_len > 10:
                continue

            sp_ordinal_pos = sp_word_pos.astype(
                int
            )  # interest area index, i.e., word IDs in fixation report
            SP_ordinal_pos.append(sp_ordinal_pos)
            SP_fix_dur.append(sp_fix_dur)
            # preprocess landing position feature
            # assign missing value to 'nan'
            sp_fix_loc = np.where(sp_fix_loc == ".", np.nan, sp_fix_loc)
            # convert string of number of float type
            sp_fix_loc = [float(i) for i in sp_fix_loc]
            # Outliers in calculated landing positions due to lack of valid AOI data, assign to 'nan'
            if (
                np.nanmax(sp_fix_loc) > 35
            ):  # returns fixation outliers (coordinates very off); np.nanmax returns the max value while igonoring nans
                missing_idx = np.where(np.array(sp_fix_loc) > 5)[
                    0
                ]  # array with indices where fix loc greater than 5
                for miss in missing_idx:
                    if sub_df.iloc[miss].CURRENT_FIX_INTEREST_AREA_LEFT in [
                        "NONE",
                        "BEFORE",
                        "AFTER",
                        "BOTH",
                    ]:
                        sp_fix_loc[miss] = np.nan
                    else:
                        print(
                            "Landing position calculation error. Unknown cause, needs to be checked"
                        )
            SP_landing_pos.append(sp_fix_loc)

            encoded_sn = tokenizer.encode_plus(
                sn_str.split(),
                add_special_tokens=False,
                padding=False,
                return_attention_mask=False,
                is_split_into_words=True,
                truncation=False,
            )

            sn_word_ids = encoded_sn.word_ids()
            sp_word_ids = [0] + sp_ordinal_pos.tolist() + [max(encoded_sn.word_ids())]

            sn_input_ids = encoded_sn["input_ids"]
            assert len(sn_word_ids) == len(sn_input_ids)

            max_len = max(max_len, len(sn_word_ids) + len(sp_word_ids))

            # truncating
            sep_token_sn_word_ids = sn_word_ids[-1]
            sep_token_sp_word_ids = sp_word_ids[-1]
            sep_token_sn_input_ids = sn_input_ids[-1]

            sn_word_ids = sn_word_ids[:-1]
            sp_word_ids = sp_word_ids[:-1]
            sn_input_ids = sn_input_ids[:-1]

            while len(sn_word_ids) + len(sp_word_ids) > args.seq_len - 3:
                if len(sn_word_ids) > len(sp_word_ids):
                    sn_word_ids.pop()
                    sn_input_ids.pop()
                elif len(sp_word_ids) > len(sn_word_ids):
                    sp_word_ids.pop()
                else:
                    sn_word_ids.pop()
                    sn_input_ids.pop()
                    sp_word_ids.pop()

            # add the SEP token word ID and input ID again
            sn_word_ids.append(sep_token_sn_word_ids)
            sp_word_ids.append(sep_token_sp_word_ids)
            sn_input_ids.append(sep_token_sn_input_ids)

            sn_sp_repr = sn_word_ids + sp_word_ids

            mask = [0] * len(sn_word_ids)
            mask_sn_padding = (
                [0] * len(sn_word_ids)
                + [1] * len(sp_word_ids)
                + [0] * (args.seq_len - len(sn_word_ids) - len(sp_word_ids))
            )
            mask_transformer_att = (
                [1] * len(sn_word_ids)
                + [1] * len(sp_word_ids)
                + [0] * (args.seq_len - len(sn_word_ids) - len(sp_word_ids))
            )

            indices_pos_enc = list(range(0, len(sn_word_ids))) + list(
                range(0, args.seq_len - len(sn_word_ids))
            )

            sn_repr_len = len(sn_word_ids)
            words_for_mapping = sn_str.split() + (args.seq_len - len(sn_str.split())) * ["[PAD]"]

            data["mask"].append(mask)
            data["sn_sp_repr"].append(sn_sp_repr)
            data["sn_input_ids"].append(sn_input_ids)
            data["indices_pos_enc"].append(indices_pos_enc)
            data["sn_repr_len"].append(sn_repr_len)
            data["words_for_mapping"].append(" ".join(words_for_mapping))
            data["mask_sn_padding"].append(mask_sn_padding)
            data["mask_transformer_att"].append(mask_transformer_att)
            data["sn_ids"].append(sn_id)
            data["reader_ids"].append(sub_id)

            reader_IDs.append(sub_id)
            sn_IDs.append(sn_id)

    # padding
    data["mask"] = _collate_batch_helper(
        examples=data["mask"],
        pad_token_id=1,
        max_length=args.seq_len,
    )
    data["sn_sp_repr"] = _collate_batch_helper(
        examples=data["sn_sp_repr"],
        pad_token_id=args.seq_len - 1,
        max_length=args.seq_len,
    )
    data["sn_input_ids"] = _collate_batch_helper(
        examples=data["sn_input_ids"],
        pad_token_id=tokenizer.pad_token_id,
        max_length=args.seq_len,
    )

    splitting_IDs_dict = {
        "reader": reader_IDs,
        "sentence": sn_IDs,
    }

    if inference == "cv":
        return data, splitting_IDs_dict

    if split == "train":

        dataset = Dataset2.from_dict(data)
        train_dataset = datasets.DatasetDict()
        train_dataset["train"] = dataset
        return train_dataset, splitting_IDs_dict

    else:

        # flatten the data
        flattened_data = list()
        for i in range(len(data["sn_sp_repr"])):
            flattened_data.append(
                (
                    data["mask"][i],
                    data["sn_sp_repr"][i],
                    data["sn_input_ids"][i],
                    data["indices_pos_enc"][i],
                    data["sn_repr_len"][i],
                    data["words_for_mapping"][i],
                    data["mask_sn_padding"][i],
                    data["mask_transformer_att"][i],
                    data["sn_ids"][i],
                    data["reader_ids"][i],
                )
            )

        if split == "train-test":

            if split_sizes:
                test_size = split_sizes["test_size"]
            else:
                test_size = 0.25

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":

                    (
                        train_data,
                        test_data,
                        train_reader_IDs,
                        test_reader_IDs,
                        train_sn_IDs,
                        test_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=test_size,
                    )

                else:

                    splitting_IDs = splitting_IDs_dict[splitting_criterion]
                    gss = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=77)
                    for train_index, test_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        test_data = np.array(flattened_data)[test_index].tolist()

            else:
                train_data, test_data = train_test_split(
                    flattened_data, test_size=test_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            test_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in test_data],
                    "sn_sp_repr": [instance[1] for instance in test_data],
                    "sn_input_ids": [instance[2] for instance in test_data],
                    "indices_pos_enc": [instance[3] for instance in test_data],
                    "sn_repr_len": [instance[4] for instance in test_data],
                    "words_for_mapping": [instance[5] for instance in test_data],
                    "mask_sn_padding": [instance[6] for instance in test_data],
                    "mask_transformer_att": [instance[7] for instance in test_data],
                    "sn_ids": [instance[8] for instance in test_data],
                    "reader_ids": [instance[9] for instance in test_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            test_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            test_data_dict["test"] = test_dataset
            return train_data_dict, test_data_dict

        elif split == "train-val":

            if split_sizes:
                val_size = split_sizes["val_size"]
            else:
                val_size = 0.1

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":
                    (
                        train_data,
                        val_data,
                        train_reader_IDs,
                        val_reader_IDs,
                        train_sn_IDs,
                        val_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=val_size,
                    )
                else:
                    splitting_IDs = splitting_IDs_dict[splitting_criterion]
                    gss = GroupShuffleSplit(n_splits=1, test_size=val_size, random_state=77)
                    for train_index, val_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        val_data = np.array(flattened_data)[val_index].tolist()

            else:
                train_data, val_data = train_test_split(
                    flattened_data, test_size=val_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            # unflatten the data
            val_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in val_data],
                    "sn_sp_repr": [instance[1] for instance in val_data],
                    "sn_input_ids": [instance[2] for instance in val_data],
                    "indices_pos_enc": [instance[3] for instance in val_data],
                    "sn_repr_len": [instance[4] for instance in val_data],
                    "words_for_mapping": [instance[5] for instance in val_data],
                    "mask_sn_padding": [instance[6] for instance in val_data],
                    "mask_transformer_att": [instance[7] for instance in val_data],
                    "sn_ids": [instance[8] for instance in val_data],
                    "reader_ids": [instance[9] for instance in val_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            val_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            val_data_dict["val"] = val_dataset
            return train_data_dict, val_data_dict

        elif split == "train-val-test":

            if split_sizes:
                val_size = split_sizes["val_size"]
                test_size = split_sizes["test_size"]
            else:
                val_size = 0.1
                test_size = 0.25

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":
                    # split train and test data so that unseen readers and sentences are in the test data
                    (
                        train_data,
                        test_data,
                        train_reader_IDs,
                        test_reader_IDs,
                        train_sn_IDs,
                        test_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=test_size,
                    )
                    # randomly split train data into train and validation
                    train_data, val_data = train_test_split(
                        train_data, test_size=val_size, random_state=77, shuffle=True
                    )

                else:
                    splitting_IDs = splitting_IDs_dict[splitting_criterion]

                    # split into train and test
                    gss = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=77)
                    for train_index, test_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        test_data = np.array(flattened_data)[test_index].tolist()
                        train_ids = np.array(splitting_IDs)[train_index].tolist()

                    # split into train and val
                    gss = GroupShuffleSplit(n_splits=1, test_size=val_size, random_state=77)
                    for train_index, val_index in gss.split(train_data, groups=train_ids):
                        val_data = np.array(train_data)[val_index].tolist()
                        train_data = np.array(train_data)[train_index].tolist()

            else:
                train_data, test_data = train_test_split(
                    flattened_data, test_size=test_size, shuffle=True, random_state=77
                )
                train_data, val_data = train_test_split(
                    train_data, test_size=val_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            test_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in test_data],
                    "sn_sp_repr": [instance[1] for instance in test_data],
                    "sn_input_ids": [instance[2] for instance in test_data],
                    "indices_pos_enc": [instance[3] for instance in test_data],
                    "sn_repr_len": [instance[4] for instance in test_data],
                    "words_for_mapping": [instance[5] for instance in test_data],
                    "mask_sn_padding": [instance[6] for instance in test_data],
                    "mask_transformer_att": [instance[7] for instance in test_data],
                    "sn_ids": [instance[8] for instance in test_data],
                    "reader_ids": [instance[9] for instance in test_data],
                }
            )
            val_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in val_data],
                    "sn_sp_repr": [instance[1] for instance in val_data],
                    "sn_input_ids": [instance[2] for instance in val_data],
                    "indices_pos_enc": [instance[3] for instance in val_data],
                    "sn_repr_len": [instance[4] for instance in val_data],
                    "words_for_mapping": [instance[5] for instance in val_data],
                    "mask_sn_padding": [instance[6] for instance in val_data],
                    "mask_transformer_att": [instance[7] for instance in val_data],
                    "sn_ids": [instance[8] for instance in val_data],
                    "reader_ids": [instance[9] for instance in val_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            test_data_dict = datasets.DatasetDict()
            val_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            test_data_dict["test"] = test_dataset
            val_data_dict["val"] = val_dataset
            return train_data_dict, test_data_dict, val_data_dict


class CelerZucoDataset(Dataset):

    def __init__(
        self,
        dataset,
        data_args,
        split,  # 'train', 'test', 'val'
    ):
        super().__init__()
        self.dataset = dataset
        self.length = len(self.dataset[split])
        self.data_args = data_args
        self.split = split

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        sample = {
            "mask": np.array(self.dataset[self.split][idx]["mask"]),
            "sn_sp_repr": np.array(self.dataset[self.split][idx]["sn_sp_repr"]),
            "sn_input_ids": np.array(self.dataset[self.split][idx]["sn_input_ids"]),
            "indices_pos_enc": np.array(self.dataset[self.split][idx]["indices_pos_enc"]),
            "sn_sp_fix_dur": np.array(self.dataset[self.split][idx]["sn_sp_fix_dur"]),
            "sn_repr_len": np.array(self.dataset[self.split][idx]["sn_repr_len"]),
            "words_for_mapping": self.dataset[self.split][idx]["words_for_mapping"],
            "mask_sn_padding": np.array(self.dataset[self.split][idx]["mask_sn_padding"]),
            "mask_transformer_att": np.array(self.dataset[self.split][idx]["mask_transformer_att"]),
            "sn_ids": self.dataset[self.split][idx]["sn_ids"],
            "reader_ids": self.dataset[self.split][idx]["reader_ids"],
        }
        return sample


def process_zuco(
    sn_list,
    reader_list,
    word_info_df,
    eyemovement_df,
    tokenizer,
    args,
    split: Optional[str] = "train",
    subset_size: Optional[int] = None,
    split_sizes: Optional[Dict[str, float]] = None,
    splitting_criterion: Optional[str] = "scanpath",  # 'reader', 'sentence', 'combined'
):
    """
    Process the ZuCo corpus so that it can be used as input to the Diffusion model, where the original sentence (sn)
    is the condition and the scan path (sp) is the target that will be noised.
    :param sn_list: list of unique sentence IDs in zuco
    :param reader_list: list of reader IDs in zuco
    :param word_info_df: pd Dataframe with sentence info
    :param eyemovement_df: pd Dataframe with fixation info
    :param tokenizer: BertTokenizer
    :param args:
    :param split: 'train', 'train-test', 'train-test-val', 'train-val'
    :param subset_size: for test runs: to not load whole dataset but specified no. of instances
    :param split_sizes: proportion of data going into train, test and val
    :param splitting_criterion: how the data should be split for testing and validation.
        'reader' = New Reader setting
        'sentence' = New Sentence setting
        'combined' = New Reader/New Sentence setting
        'scanpath' = data is split at random
    """
    SP_ordinal_pos = []
    SP_landing_pos = []
    SP_fix_dur = []

    data = {
        "mask": list(),  # 0 for sn, 1 for sp
        "sn_sp_repr": list(),  # word IDs of sn and corresponding word IDs of sp (fixated words, interest area IDs),
        # padded with args.seq_len -1
        "sn_input_ids": list(),  # input IDs of tokenized sentence, padded with pad token ID
        "indices_pos_enc": list(),  # indices from 1 ... len(sn input ids) 1 ... (seq_len - len(sn input ids))
        "sn_repr_len": list(),  # length of sentence in subword tokens
        "words_for_mapping": list(),  # original words of sentence, padded with PAD
        "mask_sn_padding": list(),  # masks both the sentence and the padding, for the loss computations
        "mask_transformer_att": list(),  # masks only the padding, for the transformer attention
        "sn_ids": list(),  # sentence IDs
        "reader_ids": list(),  # reader IDs
    }

    max_len = 0
    all_lens = list()

    reader_IDs, sn_IDs = list(), list()

    for sn_id_idx, sn_id in tqdm(enumerate(sn_list), total=len(sn_list)):

        if subset_size is not None:
            if sn_id_idx == subset_size + 1:
                break

        sn_df = eyemovement_df[eyemovement_df.sn == sn_id]
        sn = word_info_df[word_info_df.SN == sn_id]
        sn_str = " ".join(sn.WORD.values)
        sn_len = len(sn_str.split())

        tokenizer.padding_side = "right"
        sn_str = "[CLS] " + sn_str + " [SEP]"

        for sub_id_idx, sub_id in enumerate(reader_list):

            sub_df = sn_df[sn_df.id == sub_id]
            # remove fixations on non-words
            sub_df = sub_df.loc[sub_df.CURRENT_FIX_INTEREST_AREA_LABEL != ""]
            if len(sub_df) == 0:
                # no scanpath data found for the subject
                continue

            sp_word_pos, sp_fix_loc, sp_fix_dur = (
                sub_df.wn.values,
                sub_df.fl.values,
                sub_df.dur.values,
            )

            # check if recorded fixation duration are within reasonable limits
            # Less than 50ms attempt to merge with neighbouring fixation if fixate is on the same word, otherwise delete
            outlier_indx = np.where(sp_fix_dur < 50)[0]
            if outlier_indx.size > 0:
                for out_idx in range(len(outlier_indx)):
                    outlier_i = outlier_indx[out_idx]
                    merge_flag = False
                    if outlier_i - 1 >= 0 and not merge_flag:
                        # try to merge with the left fixation
                        if (
                            sub_df.iloc[outlier_i].CURRENT_FIX_INTEREST_AREA_LABEL
                            == sub_df.iloc[outlier_i - 1].CURRENT_FIX_INTEREST_AREA_LABEL
                        ):
                            sp_fix_dur[outlier_i - 1] = (
                                sp_fix_dur[outlier_i - 1] + sp_fix_dur[outlier_i]
                            )
                            merge_flag = True

                    if outlier_i + 1 < len(sp_fix_dur) and not merge_flag:
                        # try to merge with the right fixation
                        if (
                            sub_df.iloc[outlier_i].CURRENT_FIX_INTEREST_AREA_LABEL
                            == sub_df.iloc[outlier_i + 1].CURRENT_FIX_INTEREST_AREA_LABEL
                        ):
                            sp_fix_dur[outlier_i + 1] = (
                                sp_fix_dur[outlier_i + 1] + sp_fix_dur[outlier_i]
                            )
                            merge_flag = True

                    sp_word_pos = np.delete(sp_word_pos, outlier_i)
                    sp_fix_loc = np.delete(sp_fix_loc, outlier_i)
                    sp_fix_dur = np.delete(sp_fix_dur, outlier_i)
                    sub_df.drop(sub_df.index[outlier_i], axis=0, inplace=True)
                    outlier_indx = outlier_indx - 1

            # sanity check
            # scanpath too short for a normal length sentence
            if len(sp_word_pos) <= 1 and sn_len > 10:
                continue

            sp_ordinal_pos = sp_word_pos.astype(int)
            SP_ordinal_pos.append(sp_ordinal_pos)
            SP_fix_dur.append(sp_fix_dur)

            # preprocess landing position feature
            # assign missing value to 'nan'
            # sp_fix_loc=np.where(sp_fix_loc=='.', np.nan, sp_fix_loc)
            # convert string of number of float type
            sp_fix_loc = [
                float(i) if isinstance(i, int) or isinstance(i, float) else np.nan
                for i in sp_fix_loc
                if isinstance(i, int) or isinstance(i, float)
            ]
            SP_landing_pos.append(sp_fix_loc)

            # encode the sentence
            encoded_sn = tokenizer.encode_plus(
                sn_str.split(),
                add_special_tokens=False,
                padding=False,
                return_attention_mask=False,
                is_split_into_words=True,
                truncation=False,
            )

            sn_word_ids = encoded_sn.word_ids()
            sp_word_ids = [0] + sp_ordinal_pos.tolist() + [max(encoded_sn.word_ids())]

            sn_input_ids = encoded_sn["input_ids"]
            assert len(sn_word_ids) == len(sn_input_ids)

            max_len = max(max_len, len(sn_word_ids) + len(sp_word_ids))
            all_lens.append(len(sn_word_ids) + len(sp_word_ids))

            # truncating
            sep_token_sn_word_ids = sn_word_ids[-1]
            sep_token_sp_word_ids = sp_word_ids[-1]
            sep_token_sn_input_ids = sn_input_ids[-1]

            sn_word_ids = sn_word_ids[:-1]
            sp_word_ids = sp_word_ids[:-1]
            sn_input_ids = sn_input_ids[:-1]

            while len(sn_word_ids) + len(sp_word_ids) > args.seq_len - 3:
                if len(sn_word_ids) > len(sp_word_ids):
                    sn_word_ids.pop()
                    sn_input_ids.pop()
                elif len(sp_word_ids) > len(sn_word_ids):
                    sp_word_ids.pop()
                else:
                    sn_word_ids.pop()
                    sn_input_ids.pop()
                    sp_word_ids.pop()

            # add the SEP token word ID and input ID again
            sn_word_ids.append(sep_token_sn_word_ids)
            sp_word_ids.append(sep_token_sp_word_ids)
            sn_input_ids.append(sep_token_sn_input_ids)

            sn_sp_repr = sn_word_ids + sp_word_ids

            mask = [0] * len(sn_word_ids)
            mask_sn_padding = (
                [0] * len(sn_word_ids)
                + [1] * len(sp_word_ids)
                + [0] * (args.seq_len - len(sn_word_ids) - len(sp_word_ids))
            )
            mask_transformer_att = (
                [1] * len(sn_word_ids)
                + [1] * len(sp_word_ids)
                + [0] * (args.seq_len - len(sn_word_ids) - len(sp_word_ids))
            )

            indices_pos_enc = list(range(0, len(sn_word_ids))) + list(
                range(0, args.seq_len - len(sn_word_ids))
            )

            sn_repr_len = len(sn_word_ids)
            words_for_mapping = sn_str.split() + (args.seq_len - len(sn_str.split())) * ["[PAD]"]

            data["mask"].append(mask)
            data["sn_sp_repr"].append(sn_sp_repr)
            data["sn_input_ids"].append(sn_input_ids)
            data["indices_pos_enc"].append(indices_pos_enc)
            data["sn_repr_len"].append(sn_repr_len)
            data["words_for_mapping"].append(" ".join(words_for_mapping))
            data["mask_sn_padding"].append(mask_sn_padding)
            data["mask_transformer_att"].append(mask_transformer_att)
            data["sn_ids"].append(sn_id)
            data["reader_ids"].append(sub_id)

            reader_IDs.append(sub_id)
            sn_IDs.append(sn_id)

    # padding
    data["mask"] = _collate_batch_helper(
        examples=data["mask"],
        pad_token_id=1,
        max_length=args.seq_len,
    )
    data["sn_sp_repr"] = _collate_batch_helper(
        examples=data["sn_sp_repr"],
        pad_token_id=args.seq_len - 1,
        max_length=args.seq_len,
    )
    data["sn_input_ids"] = _collate_batch_helper(
        examples=data["sn_input_ids"],
        pad_token_id=tokenizer.pad_token_id,
        max_length=args.seq_len,
    )

    splitting_IDs_dict = {
        "reader": reader_IDs,
        "sentence": sn_IDs,
    }

    if split == "train":

        dataset = Dataset2.from_dict(data)
        train_dataset = datasets.DatasetDict()
        train_dataset["train"] = dataset
        return train_dataset

    else:

        # flatten the data
        flattened_data = list()
        for i in range(len(data["sn_sp_repr"])):
            flattened_data.append(
                (
                    data["mask"][i],
                    data["sn_sp_repr"][i],
                    data["sn_input_ids"][i],
                    data["indices_pos_enc"][i],
                    data["sn_repr_len"][i],
                    data["words_for_mapping"][i],
                    data["mask_sn_padding"][i],
                    data["mask_transformer_att"][i],
                    data["sn_ids"][i],
                    data["reader_ids"][i],
                )
            )

        if split == "train-test":

            if split_sizes:
                test_size = split_sizes["test_size"]
            else:
                test_size = 0.25

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":

                    (
                        train_data,
                        test_data,
                        train_reader_IDs,
                        test_reader_IDs,
                        train_sn_IDs,
                        test_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=test_size,
                    )

                else:

                    splitting_IDs = splitting_IDs_dict[splitting_criterion]
                    gss = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=77)
                    for train_index, test_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        test_data = np.array(flattened_data)[test_index].tolist()

            else:
                train_data, test_data = train_test_split(
                    flattened_data, test_size=test_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            test_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in test_data],
                    "sn_sp_repr": [instance[1] for instance in test_data],
                    "sn_input_ids": [instance[2] for instance in test_data],
                    "indices_pos_enc": [instance[3] for instance in test_data],
                    "sn_repr_len": [instance[4] for instance in test_data],
                    "words_for_mapping": [instance[5] for instance in test_data],
                    "mask_sn_padding": [instance[6] for instance in test_data],
                    "mask_transformer_att": [instance[7] for instance in test_data],
                    "sn_ids": [instance[8] for instance in test_data],
                    "reader_ids": [instance[9] for instance in test_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            test_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            test_data_dict["test"] = test_dataset
            return train_data_dict, test_data_dict

        elif split == "train-val":

            if split_sizes:
                val_size = split_sizes["val_size"]
            else:
                val_size = 0.1

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":
                    (
                        train_data,
                        val_data,
                        train_reader_IDs,
                        val_reader_IDs,
                        train_sn_IDs,
                        val_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=val_size,
                    )
                else:
                    splitting_IDs = splitting_IDs_dict[splitting_criterion]
                    gss = GroupShuffleSplit(n_splits=1, test_size=val_size, random_state=77)
                    for train_index, val_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        val_data = np.array(flattened_data)[val_index].tolist()

            else:
                train_data, val_data = train_test_split(
                    flattened_data, test_size=val_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            # unflatten the data
            val_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in val_data],
                    "sn_sp_repr": [instance[1] for instance in val_data],
                    "sn_input_ids": [instance[2] for instance in val_data],
                    "indices_pos_enc": [instance[3] for instance in val_data],
                    "sn_repr_len": [instance[4] for instance in val_data],
                    "words_for_mapping": [instance[5] for instance in val_data],
                    "mask_sn_padding": [instance[6] for instance in val_data],
                    "mask_transformer_att": [instance[7] for instance in val_data],
                    "sn_ids": [instance[8] for instance in val_data],
                    "reader_ids": [instance[9] for instance in val_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            val_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            val_data_dict["val"] = val_dataset
            return train_data_dict, val_data_dict

        elif split == "train-val-test":

            if split_sizes:
                val_size = split_sizes["val_size"]
                test_size = split_sizes["test_size"]
            else:
                val_size = 0.1
                test_size = 0.25

            if splitting_criterion != "scanpath":

                if splitting_criterion == "combined":
                    # split train and test data so that unseen readers and sentences are in the test data
                    (
                        train_data,
                        test_data,
                        train_reader_IDs,
                        test_reader_IDs,
                        train_sn_IDs,
                        test_sn_IDs,
                    ) = combined_split(
                        data=flattened_data,
                        reader_IDs=splitting_IDs_dict["reader"],
                        sn_IDs=splitting_IDs_dict["sentence"],
                        test_size=test_size,
                    )
                    # randomly split train data into train and validation
                    train_data, val_data = train_test_split(
                        train_data, test_size=val_size, random_state=77, shuffle=True
                    )

                else:
                    splitting_IDs = splitting_IDs_dict[splitting_criterion]

                    # split into train and test
                    gss = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=77)
                    for train_index, test_index in gss.split(flattened_data, groups=splitting_IDs):
                        train_data = np.array(flattened_data)[train_index].tolist()
                        test_data = np.array(flattened_data)[test_index].tolist()
                        train_ids = np.array(splitting_IDs)[train_index].tolist()

                    # split into train and val
                    gss = GroupShuffleSplit(n_splits=1, test_size=val_size, random_state=77)
                    for train_index, val_index in gss.split(train_data, groups=train_ids):
                        val_data = np.array(train_data)[val_index].tolist()
                        train_data = np.array(train_data)[train_index].tolist()

            else:
                train_data, test_data = train_test_split(
                    flattened_data, test_size=test_size, shuffle=True, random_state=77
                )
                train_data, val_data = train_test_split(
                    train_data, test_size=val_size, shuffle=True, random_state=77
                )

            # unflatten the data
            train_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in train_data],
                    "sn_sp_repr": [instance[1] for instance in train_data],
                    "sn_input_ids": [instance[2] for instance in train_data],
                    "indices_pos_enc": [instance[3] for instance in train_data],
                    "sn_repr_len": [instance[4] for instance in train_data],
                    "words_for_mapping": [instance[5] for instance in train_data],
                    "mask_sn_padding": [instance[6] for instance in train_data],
                    "mask_transformer_att": [instance[7] for instance in train_data],
                    "sn_ids": [instance[8] for instance in train_data],
                    "reader_ids": [instance[9] for instance in train_data],
                }
            )
            test_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in test_data],
                    "sn_sp_repr": [instance[1] for instance in test_data],
                    "sn_input_ids": [instance[2] for instance in test_data],
                    "indices_pos_enc": [instance[3] for instance in test_data],
                    "sn_repr_len": [instance[4] for instance in test_data],
                    "words_for_mapping": [instance[5] for instance in test_data],
                    "mask_sn_padding": [instance[6] for instance in test_data],
                    "mask_transformer_att": [instance[7] for instance in test_data],
                    "sn_ids": [instance[8] for instance in test_data],
                    "reader_ids": [instance[9] for instance in test_data],
                }
            )
            val_dataset = Dataset2.from_dict(
                {
                    "mask": [instance[0] for instance in val_data],
                    "sn_sp_repr": [instance[1] for instance in val_data],
                    "sn_input_ids": [instance[2] for instance in val_data],
                    "indices_pos_enc": [instance[3] for instance in val_data],
                    "sn_repr_len": [instance[4] for instance in val_data],
                    "words_for_mapping": [instance[5] for instance in val_data],
                    "mask_sn_padding": [instance[6] for instance in val_data],
                    "mask_transformer_att": [instance[7] for instance in val_data],
                    "sn_ids": [instance[8] for instance in val_data],
                }
            )
            train_data_dict = datasets.DatasetDict()
            test_data_dict = datasets.DatasetDict()
            val_data_dict = datasets.DatasetDict()
            train_data_dict["train"] = train_dataset
            test_data_dict["test"] = test_dataset
            val_data_dict["val"] = val_dataset
            return train_data_dict, test_data_dict, val_data_dict