#!/usr/bin/env python3
# convert_jsonl.py

import os
import json
from glob import glob
from tqdm import tqdm

# 1. 输入、输出目录
INPUT_DIR = "/cpfs/user/boyuan/verl_workspace/baidu/data/gsm8k/main"
OUTPUT_DIR = "/cpfs/user/boyuan/verl_workspace/baidu/ernie/ERNIE/examples/data"

# 2. 遍历所有 .jsonl 文件
for input_path in glob(os.path.join(INPUT_DIR, "*.jsonl")):
    filename = os.path.basename(input_path)
    output_filename = "gsm8k_" + filename
    output_path = os.path.join(OUTPUT_DIR, output_filename)

    # 先统计总行数（用于进度条）
    with open(input_path, "r", encoding="utf-8") as fcount:
        total_lines = sum(1 for _ in fcount)

    # 流式处理并显示进度
    with open(input_path, "r", encoding="utf-8") as fin, \
         open(output_path, "w", encoding="utf-8") as fout:

        for line in tqdm(fin, total=total_lines, desc=f"Processing {filename}", unit=" lines"):
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue

            if "question" not in obj or "answer" not in obj:
                continue

            new_obj = {
                "src": [obj["question"]],
                "tgt": [obj["answer"]]
            }
            fout.write(json.dumps(new_obj, ensure_ascii=False) + "\n")

    print(f"→ {filename} 处理完成，结果写入 {output_filename}")