import json
import numpy as np

np.random.seed(0)

data = []
with open("chat_train_data.json", "r") as fr:
    for line in fr.readlines():
        data.append(line)

test_size = 200
np.random.shuffle(data)

train_data = data[:-test_size]
test_data = data[-test_size:]

with open("train_data.json", "w") as fw:
    for d in train_data:
        fw.write(f"{d}")

with open("test_data.json", "w") as fw:
    for d in test_data:
        fw.write(f"{d}")