File size: 1,410 Bytes
2daf3c7 4463ade 2daf3c7 4463ade 2daf3c7 4463ade 2daf3c7 4463ade 2daf3c7 4463ade 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 50 51 52 53 54 55 56 57 58 |
import json
import os
import random
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]))
|