From 7cd7740d5e7ed65e06e44d66b3ff49798fab83cc Mon Sep 17 00:00:00 2001 From: BAIKEMARK Date: Sat, 24 May 2025 10:50:52 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84clean=5Fby=5Fscore=E9=80=BB?= =?UTF-8?q?=E8=BE=91=E5=B9=B6=E5=B0=86=E5=85=B6=E6=95=B4=E5=90=88=E8=87=B3?= =?UTF-8?q?strategies?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- weclone/data/clean/clean_by_score.py | 38 --------------- weclone/data/clean/strategies.py | 65 +++++++++++++++++++++---- weclone/data/clean/strategies_online.py | 11 ----- weclone/train/train_sft.py | 17 ++++--- 4 files changed, 66 insertions(+), 65 deletions(-) delete mode 100644 weclone/data/clean/clean_by_score.py diff --git a/weclone/data/clean/clean_by_score.py b/weclone/data/clean/clean_by_score.py deleted file mode 100644 index c80c7cc..0000000 --- a/weclone/data/clean/clean_by_score.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -import json -from weclone.utils.config import load_config -from weclone.utils.log import logger - - -def clean_sft_data() -> str: - """ - 清洗 SFT 数据并返回清洗后的文件路径。 - 如果未启用清洗，则返回原始路径。 - """ - config = load_config(arg_type="make_dataset") - sft_json_path = os.path.join(config["dataset_dir"], "sft-my.json") - output_json_path = os.path.join(config["dataset_dir"], "sft-my-l.json") - accept_score = config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1) - - if not config.get("clean_dataset", {}).get("enable_clean"): - logger.info("未启用清洗功能") - return sft_json_path - - try: - with open(sft_json_path, 'r', encoding='utf-8') as f: - data = json.load(f) - filtered_data = [item for item in data if item.get("score", 1) >= accept_score] - - with open(output_json_path, 'w', encoding='utf-8') as f: - json.dump(filtered_data, f, ensure_ascii=False, indent=4) - - logger.success(f"已筛出低于{accept_score}分的数据，共保留 {len(filtered_data)} 条数据") - return output_json_path - - except Exception as e: - logger.error(f"清洗数据失败，使用原始数据: {str(e)}") - return sft_json_path - - -if __name__ == "__main__": - clean_sft_data() diff --git a/weclone/data/clean/strategies.py b/weclone/data/clean/strategies.py index 776321d..aab4373 100644 --- a/weclone/data/clean/strategies.py +++ b/weclone/data/clean/strategies.py @@ -6,7 +6,8 @@ from typing import Any, Dict, List, Union from langchain_core.prompts import PromptTemplate from weclone.data.models import QaPair, CutMessage, QaPairScore from weclone.prompts.clean_data import CLEAN_PROMPT - +import os +import commentjson from weclone.utils.log import logger @@ -86,13 +87,59 @@ class LLMCleaningStrategy(CleaningStrategy): printable_df_str = distribution_df.reset_index().to_string(index=False) logger.success(f"llm打分分数分布情况:\n{printable_df_str}") - def clean(self, data: List[QaPair]) -> List[QaPair]: + def clean(self) -> str: """ - 根据打分结果，删除分数低于阈值的数据。 + 清洗 SFT 数据并返回清洗后的文件路径。 + 如果未启用清洗，则返回原始路径。 """ - return [ - qa - for qa in data - if qa.score is not None - and qa.score >= self.make_dataset_config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1) - ] + config = self.make_dataset_config + dataset_dir = config["dataset_dir"] + dataset_info_path = os.path.join(dataset_dir, "dataset_info.json") + + sft_json_path = os.path.join(dataset_dir, "sft-my.json") + output_json_path = os.path.join(dataset_dir, "sft-my-l.json") + accept_score = config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1) + + if not config.get("clean_dataset", {}).get("enable_clean"): + logger.info("未启用清洗功能") + self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json") + return sft_json_path + + try: + with open(sft_json_path, 'r', encoding='utf-8') as f: + data = json.load(f) + filtered_data = [item for item in data if item.get("score", 0) >= accept_score] + + with open(output_json_path, 'w', encoding='utf-8') as f: + json.dump(filtered_data, f, ensure_ascii=False, indent=4) + + logger.success(f"已筛出低于{accept_score}分的数据，共保留 {len(filtered_data)} 条数据") + self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my-l.json") + return output_json_path + + except Exception as e: + logger.error(f"清洗数据失败，使用原始数据: {str(e)}") + self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json") + return sft_json_path + + def _update_dataset_info_file(self, dataset_info_path: str, new_file_name: str): + """ + 修改 dataset_info.json 文件中的 file_name 字段 + """ + try: + with open(dataset_info_path, "r", encoding="utf-8") as f: + dataset_info = commentjson.load(f) + + # 更新所有支持的数据集的 file_name + for key in ["wechat-sft", "wechat-sft-with-history"]: + if key in dataset_info: + dataset_info[key]["file_name"] = new_file_name + + # 写回文件 + with open(dataset_info_path, "w", encoding="utf-8") as f: + commentjson.dump(dataset_info, f, indent=4, ensure_ascii=False) + + logger.info(f"已更新 dataset_info.json 中的 file_name 为 {new_file_name}") + + except Exception as e: + logger.warning(f"无法更新 dataset_info.json: {e}") diff --git a/weclone/data/clean/strategies_online.py b/weclone/data/clean/strategies_online.py index aaa71d9..438214f 100644 --- a/weclone/data/clean/strategies_online.py +++ b/weclone/data/clean/strategies_online.py @@ -95,14 +95,3 @@ class OlineLLMCleaningStrategy(CleaningStrategy): distribution_df.index.name = "分数" printable_df_str = distribution_df.reset_index().to_string(index=False) logger.success(f"在线模型打分分数分布情况:\n{printable_df_str}") - - def clean(self, data: List[QaPair]) -> List[QaPair]: - """ - 根据打分结果，删除分数低于阈值的数据。 - """ - threshold = self.make_dataset_config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1) - return [ - qa - for qa in data - if qa.score is not None and qa.score >= threshold - ] diff --git a/weclone/train/train_sft.py b/weclone/train/train_sft.py index 7fc3316..78483d7 100644 --- a/weclone/train/train_sft.py +++ b/weclone/train/train_sft.py @@ -5,24 +5,27 @@ from llamafactory.train.tuner import run_exp from llamafactory.extras.misc import get_current_device from weclone.utils.config import load_config from weclone.utils.log import logger -from weclone.data.clean.clean_by_score import clean_sft_data +from weclone.data.clean.strategies import LLMCleaningStrategy def main(): - config = load_config(arg_type="train_sft") + train_config = load_config(arg_type="train_sft") + dataset_config = load_config(arg_type="make_dataset") device = get_current_device() if device == "cpu": logger.warning("请注意你正在使用CPU训练，非Mac设备可能会出现问题") - sft_json_path = clean_sft_data() - if not os.path.exists(sft_json_path): - logger.error(f"错误：文件 '{sft_json_path}' 不存在，请确保数据处理步骤已正确生成该文件。") + cleaner = LLMCleaningStrategy(make_dataset_config=dataset_config) + cleaned_data_path = cleaner.clean() + + if not os.path.exists(cleaned_data_path): + logger.error(f"错误：文件 '{cleaned_data_path}' 不存在，请确保数据处理步骤已正确生成该文件。") sys.exit(1) - formatted_config = json.dumps(config, indent=4, ensure_ascii=False) + formatted_config = json.dumps(train_config, indent=4, ensure_ascii=False) logger.info(f"微调配置：\n{formatted_config}") - run_exp(config) + run_exp(train_config) if __name__ == "__main__":