join_datasets_custom_split.py

import os
import json
import random

import pandas as pd


DATA_DIR = f"/home/{os.environ['USER']}/data/wit/all_jsons"
SEED = 0
PROPORTION_TRAIN = 0.98
PROPORTION_VALID = 0.01

random.seed(SEED)

all_files = [f"{DATA_DIR}/{file_}" for file_ in os.listdir(DATA_DIR) if ("all" not in file_)]

print(all_files)

examples = []
for file_ in all_files:
    print(file_)
    with open(file_) as f:
        file_examples = [json.dumps(json.loads(line), ensure_ascii=False) for line in f.readlines()]
    print(len(file_examples))
    examples.extend(file_examples)

print(f"Before dedup: {len(examples)}")
examples = list(set(examples))
print(f"After dedup: {len(examples)}")

print(examples[0])
# Shuffle examples
random.shuffle(examples)
print(examples[0])

split_dataset = {}
split_dataset["train"] = examples[:int(len(examples) * PROPORTION_TRAIN)]
split_dataset["valid"] = examples[int(len(examples) * PROPORTION_TRAIN): int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID))]
split_dataset["test"] = examples[int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)):]


for split in ["train", "valid", "test"]:
    print("-----")
    print(len(split_dataset[split]))
    print("-----")
    with open(f"/home/{os.environ['USER']}/data/wit/all_jsons/{split}_dataset_all_98_1_1_split.json", "w") as f:
        f.write("\n".join(split_dataset[split]))