from datasets import DatasetDict, load_from_disk
import argparse

from openai_dataset_maker import features

def has_all_valid_labels(exp):
    for col, labels in exp.items():
        if col in {"text", "tokens"}:
            continue
        for label in labels:
            if label not in features[col]:
                return False
    return True

def is_evenly_shaped(exp):
    cnt_set = set()
    for col, labels in exp.items():
        if col == "text":
            continue
        cnt_set.add(len(labels))
    return len(cnt_set) == 1


if __name__ == '__main__':
    arg_parser = argparse.ArgumentParser(description="Train multi-task model.")
    arg_parser.add_argument("data_path", help="Load dataset from specified path.",
                            action="store")
    arg_parser.add_argument("--save-path", help="Save final dataset to specified path.",
                            action="store", default="./training_data")
    args = arg_parser.parse_args()

    loaded_dataset = load_from_disk(args.data_path)
    loaded_dataset = loaded_dataset.filter(is_evenly_shaped)
    loaded_dataset = loaded_dataset.filter(has_all_valid_labels)

    first_split = loaded_dataset.train_test_split(shuffle=True, seed=42, test_size=0.09)
    second_split = first_split["train"].train_test_split(test_size=0.1)

    new_ds = DatasetDict()
    new_ds["test"] = first_split["test"]
    new_ds["train"] = second_split["train"]
    new_ds["validation"] = second_split["test"]
    new_ds.save_to_disk(args.save_path)