File size: 1,353 Bytes
2daf3c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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]))
|