import json
import random

# --- 設定 ---
input_file = 'train.jsonl'
output_train_file = 'train.jsonl'  # 元のファイルを上書き
output_valid_file = 'valid.jsonl'
valid_split_ratio = 0.1  # 10%を検証データにする
# -------------

print(f"'{input_file}' を読み込んでいます...")

# 元のデータを読み込む
with open(input_file, 'r', encoding='utf-8') as f:
    lines = f.readlines()

# データをシャッフルしてランダム性を確保
random.shuffle(lines)

# 分割点を計算
split_index = int(len(lines) * (1 - valid_split_ratio))

# 学習データと検証データに分割
train_data = lines[:split_index]
valid_data = lines[split_index:]

# 新しい学習データとして書き出す
with open(output_train_file, 'w', encoding='utf-8') as f:
    f.writelines(train_data)
print(f"学習データを '{output_train_file}' に保存しました。({len(train_data)}行)")

# 検証データとして書き出す
with open(output_valid_file, 'w', encoding='utf-8') as f:
    f.writelines(valid_data)
print(f"検証データを '{output_valid_file}' に保存しました。({len(valid_data)}行)")

print("分割が完了しました。")