trans_data.py

import os
import json


def main():
    """
    主函数，处理JSONL格式的自我认知数据，替换模板变量并分割为训练集和验证集
    """
    # 直接在代码中定义名称和作者
    name = "千问宝宝"  # 模型名称
    author = "小侯团队"  # 模型作者

    mlx_data = []  # 存储处理后的数据列表

    # 读取原始JSONL文件
    with open("self-cognition/self_cognition.jsonl", "r") as fread:
        data_list = fread.readlines()  # 读取所有行

        # 处理每一行数据
        for data in data_list:
            data = json.loads(data)  # 解析JSON字符串
            user_text = data["query"]  # 获取用户输入文本

            # 根据语言标签处理助手回复
            if data["tag"] == "zh":  # 中文处理
                assistant_text = (
                    data["response"]
                    .replace("{{NAME}}", name)  # 使用定义的name
                    .replace("{{AUTHOR}}", author)  # 使用定义的author
                )
            else:  # 非中文处理(英文)
                assistant_text = (
                    data["response"]
                    .replace("{{NAME}}", "Little-" + name[:1])  # 英文名处理
                    .replace("{{AUTHOR}}", author + "Fans")  # 英文作者处理
                )

            # 构建格式化后的对话数据
            mlx_data.append(
                {
                    "messages": [
                        {"role": "user", "content": user_text},
                        {"role": "assistant", "content": assistant_text},
                    ]
                }
            )

    # 分割数据集: 验证集占20%(1/5)，训练集占80%
    val_data_num = len(mlx_data) // 5  # 计算验证集数量
    mlx_train_data = mlx_data[val_data_num:]  # 训练集(后80%)
    mlx_val_data = mlx_data[:val_data_num]  # 验证集(前20%)

    # 创建输出目录(如果不存在)
    os.makedirs("./mlx_data/", exist_ok=True)

    # 写入训练集文件
    with open("./mlx_data/train.jsonl", "w", encoding="utf-8") as fwrite:
        for data in mlx_train_data:
            fwrite.write(json.dumps(data, ensure_ascii=False) + "\n")  # 保持中文不转义

    # 写入验证集文件
    with open("./mlx_data/valid.jsonl", "w", encoding="utf-8") as fwrite:
        for data in mlx_val_data:
            fwrite.write(json.dumps(data, ensure_ascii=False) + "\n")  # 保持中文不转义


if __name__ == "__main__":
    # 直接调用主函数，无需参数
    main()