def METHOD_NAME(data_ds, aug_type, num_aug=10, percent=0.1, aug_base="mlm", example_keys=None):
    """
    Extend train dataset with augmentation.
    """
    if example_keys is None:
        return data_ds
    if aug_type is None or aug_type == "None":
        return data_ds
    if aug_type == "delete":
        aug = WordDelete(create_n=num_aug, aug_percent=percent)
    elif aug_type == "substitute":
        aug = WordSubstitute(aug_base, create_n=num_aug, aug_percent=percent)
    elif aug_type == "insert":
        aug = WordInsert(aug_base, create_n=num_aug, aug_percent=percent)
    elif aug_type == "swap":
        aug = WordSwap(create_n=num_aug, aug_percent=percent)
    else:
        raise ValueError("Unsupported data augment strategy `{}`".format(aug_type))
    aug_data = []
    for example in data_ds:
        for key in example_keys:
            text_aug = aug.augment(example[key])
            for text in text_aug:
                new_example = example.copy()
                example[key] = text
                aug_data.append(new_example)
    data_ds = MapDataset([x for x in data_ds] + aug_data)
    return data_ds