2_2_split_train_test.py

import json
import random

# Path to your input JSON file
input_path = "/data/hongwei/xzh_workspace/patch_megatron_code_new/my_datasets/sampled_videos_w_wo_motion_stage1_85k.json"

# Load the data
with open(input_path, 'r') as f:
    data = json.load(f)

# Shuffle the data
random.shuffle(data)

# Split into training and testing
test_size = 100
test_data = data[:test_size]
train_data = data[test_size:]

# Save the split data
train_output_path = input_path.replace(".json", "_train.json")
test_output_path = input_path.replace(".json", "_test.json")

with open(train_output_path, 'w') as f:
    json.dump(train_data, f, indent=2)

with open(test_output_path, 'w') as f:
    json.dump(test_data, f, indent=2)

print(f"Training set saved to: {train_output_path}")
print(f"Test set saved to: {test_output_path}")