From bbe16fda730a6b7a468842229024d8cfafdb9924 Mon Sep 17 00:00:00 2001
From: xming521 <1223398803@qq.com>
Date: Thu, 22 May 2025 20:05:22 +0800
Subject: [PATCH] Add full-flow test pipeline
---
pyproject.toml | 11 +-
tests/README.md | 39 --
tests/__init__.py | 0
tests/data/clean/test_strategies.py | 176 ------
tests/test_full_pipe.py | 154 +++++
tests/test_full_pipeline.py | 902 ----------------------------
tests/test_get_sample_audio.py | 151 -----
tests/test_get_score.py | 72 ---
tests/test_old_csv_to_json.py | 312 ----------
tests/test_qa_generator.py | 353 -----------
tests/test_weclone_pipeline_mock.py | 306 ----------
weclone/data/qa_generator.py | 9 +-
12 files changed, 169 insertions(+), 2316 deletions(-)
delete mode 100644 tests/README.md
create mode 100644 tests/__init__.py
delete mode 100644 tests/data/clean/test_strategies.py
create mode 100644 tests/test_full_pipe.py
delete mode 100644 tests/test_full_pipeline.py
delete mode 100644 tests/test_get_sample_audio.py
delete mode 100644 tests/test_get_score.py
delete mode 100644 tests/test_old_csv_to_json.py
delete mode 100644 tests/test_qa_generator.py
delete mode 100644 tests/test_weclone_pipeline_mock.py
diff --git a/pyproject.toml b/pyproject.toml
index 7c3f453..d720eb7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,8 +43,12 @@ sparktts = [
"torchaudio>=2.6.0",
"tqdm>=4.66.5",
]
-main = ["llamafactory>=0.9.2", "openai==1.76.0", "vllm==0.8.2; platform_system == 'Linux'"]
-dev = ["pytest", "pyright", "ruff"]
+main = [
+ "llamafactory>=0.9.2",
+ "openai==1.76.0",
+ "vllm==0.8.2; platform_system == 'Linux'",
+]
+dev = ["pytest", "pytest-order", "pyright", "ruff"]
[project.scripts]
weclone-cli = "weclone.cli:cli"
@@ -115,3 +119,6 @@ lint.select = [
"Q", # flake8-quotes
]
target-version = "py310"
+
+[tool.pytest.ini_options]
+addopts = "-x"
diff --git a/tests/README.md b/tests/README.md
deleted file mode 100644
index d0e6163..0000000
--- a/tests/README.md
+++ /dev/null
@@ -1,39 +0,0 @@
-# WEClone 测试指南
-
-本目录包含WEClone项目的测试文件，用于确保项目各个组件正常工作。
-
-## 测试文件说明
-
-- `test_weclone_pipeline.py`: 全流程测试，按顺序测试数据生成、训练、API服务和模型评估
-- `test_qa_generator.py`: 测试QA生成器功能
-
-
-## 运行全流程测试
-
-要运行完整的测试流程，请执行以下命令：
-
-```bash
-# 在项目根目录下执行
-python -m tests.test_weclone_pipeline
-```
-
-## 测试流程说明
-
-全流程测试按照以下顺序测试项目的主要组件：
-
-1. **数据生成**：测试 `weclone/data/qa_generator.py` 模块，模拟微信聊天记录的处理和QA对的生成
-2. **模型训练**：测试 `weclone/train/train_sft.py` 模块，模拟使用生成的数据进行模型的SFT训练
-3. **API服务**：测试 `weclone/server/api_service.py` 模块，模拟启动API服务
-4. **模型评估**：测试 `weclone/eval/test_model.py` 模块，模拟对训练后的模型进行评估
-
-## 注意事项
-
-- 测试使用Python的unittest框架和mock库，模拟各个组件的运行环境和依赖
-- 测试不会修改实际的数据文件或模型文件，所有操作都在临时目录中进行
-- 要运行单独的测试方法，可以使用以下命令：
-
-```bash
-# 例如，只运行QA生成器测试
-python -m unittest tests.test_weclone_pipeline.TestWeclonePipeline.test_qa_generator
-```
-
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/data/clean/test_strategies.py b/tests/data/clean/test_strategies.py
deleted file mode 100644
index 74b465f..0000000
--- a/tests/data/clean/test_strategies.py
+++ /dev/null
@@ -1,176 +0,0 @@
-import pytest
-from unittest.mock import patch, MagicMock, call
-from langchain_core.prompts import PromptTemplate
-from datetime import datetime
-import pandas as pd # 导入 pandas
-
-# 确保可以正确导入被测试的模块和依赖项
-# 可能需要根据你的项目结构调整导入路径
-try:
- from weclone.data.clean.strategies import LLMCleaningStrategy
- from weclone.data.models import QaPair
- from weclone.prompts.clean_data import CLEAN_PROMPT
-except ImportError:
- # 如果直接运行脚本时找不到模块，尝试添加项目根目录到 sys.path
- import sys
- import os
- # 获取当前脚本文件所在的目录 (tests/data/clean)
- current_dir = os.path.dirname(os.path.abspath(__file__))
- # 获取 tests 目录
- tests_dir = os.path.dirname(os.path.dirname(current_dir))
- # 获取项目根目录 (weclone 的父目录)
- project_root = os.path.dirname(tests_dir)
- sys.path.insert(0, project_root)
- from weclone.data.clean.strategies import LLMCleaningStrategy
- from weclone.data.models import QaPair
- from weclone.prompts.clean_data import CLEAN_PROMPT
-
-
-
-@pytest.fixture
-def sample_qa_pairs():
- """提供一些测试用的 QaPair 数据"""
- # now = datetime.now() # 不再需要 datetime
- return [
- QaPair(id=1, instruction="问题1", output="答案1", system="", history=[], time=pd.Timestamp.now(), score=0), # 使用 pd.Timestamp
- QaPair(id=2, instruction="问题2", output="答案2", system="", history=[], time=pd.Timestamp.now(), score=0), # 使用 pd.Timestamp
- ]
-
-@pytest.fixture
-def mock_make_dataset_config():
- """提供模拟的 make_dataset_config"""
- return {
- "model_name_or_path": "mock_model",
- "template": "mock_template",
- # 可以根据需要添加其他配置
- }
-
-@patch("weclone.data.clean.strategies.infer") # 模拟 infer 函数
-def test_llm_cleaning_strategy_clean(mock_infer, sample_qa_pairs, mock_make_dataset_config):
- """测试 LLMCleaningStrategy.clean 方法"""
- # 1. 准备
- print("--- 开始测试 test_llm_cleaning_strategy_clean ---")
- strategy = LLMCleaningStrategy(make_dataset_config=mock_make_dataset_config)
- prompt_template = PromptTemplate.from_template(CLEAN_PROMPT)
-
- # 预期 infer 函数的输入
- expected_inputs = []
- for qa in sample_qa_pairs:
- expected_inputs.append(prompt_template.invoke({"id": qa.id, "Q": qa.instruction, "A": qa.output}))
- print(f"预期 infer 输入: {expected_inputs}")
-
- # 设置模拟 infer 函数的返回值
- mock_cleaned_outputs = ["cleaned_output_1", "cleaned_output_2"]
- mock_infer.return_value = mock_cleaned_outputs
- print(f"设置 mock infer 返回值: {mock_cleaned_outputs}")
-
- # 2. 执行
- print("调用 strategy.clean...")
- # 注意：原始的 clean 方法没有 return 语句。如果需要测试返回值，
- # 需要在 weclone/data/clean/strategies.py 中取消注释 'return cleaned_data'
- cleaned_data = strategy.clean(sample_qa_pairs)
- # strategy.clean(sample_qa_pairs) # 暂时只调用，不获取返回值
- print(f"获取的 cleaned_data: {cleaned_data}") # 如果有返回值，取消注释此行
-
- # 3. 断言
- print("执行断言...")
- # 验证 infer 函数是否以正确的参数被调用
- try:
- mock_infer.assert_called_once_with(
- expected_inputs,
- mock_make_dataset_config["model_name_or_path"],
- template=mock_make_dataset_config["template"],
- temperature=0,
- )
- print("infer 函数调用断言成功！")
- except AssertionError as e:
- print(f"infer 函数调用断言失败: {e}")
- raise # 重新抛出异常，以便 pytest 能捕获
-
- # 验证 clean 方法的返回值（基于假设）
- # 如果原始 clean 方法确实没有 return，可以移除这个断言或者修改 clean 方法添加 return
- assert cleaned_data == mock_cleaned_outputs
- print("返回值断言成功！")
-
- print("--- 测试 test_llm_cleaning_strategy_clean 结束 ---")
-
-
-if __name__ == "__main__":
- print("直接运行测试脚本进行调试...")
-
- # 手动准备依赖项 (代替 pytest fixtures)
- # now_main = datetime.now() # 不再需要 datetime
- qa_pairs = [
- QaPair(id=101, instruction="调试问题1", output="调试答案1", system="", history=[], time=pd.Timestamp.now(), score=0), # 使用 pd.Timestamp
- QaPair(id=102, instruction="调试问题2", output="调试答案2", system="", history=[], time=pd.Timestamp.now(), score=0), # 使用 pd.Timestamp
- ]
- config = {
- "model_name_or_path": "debug_model",
- "template": "debug_template",
- }
-
- # 方案1：直接调用被测代码逻辑 (更简单)
- print("\n--- 方案1：直接调用被测代码逻辑 ---")
- try:
- from weclone.data.clean.strategies import infer # 需要导入 infer
- except ImportError:
- # 处理导入错误的代码已在文件顶部
- from weclone.data.clean.strategies import infer
-
- with patch("weclone.data.clean.strategies.infer") as mock_infer_main:
- strategy = LLMCleaningStrategy(make_dataset_config=config)
- prompt_template = PromptTemplate.from_template(CLEAN_PROMPT)
- inputs_main = []
- for qa in qa_pairs:
- inputs_main.append(prompt_template.invoke({"id": qa.id, "Q": qa.instruction, "A": qa.output}))
-
- mock_return = ["debug_cleaned_1", "debug_cleaned_2"]
- mock_infer_main.return_value = mock_return
- print(f"设置 main 中的 mock infer 返回值: {mock_return}")
-
- print("在 main 中调用 strategy.clean...")
- cleaned_result_main = strategy.clean(qa_pairs) # 如果 clean 有返回值
- # strategy.clean(qa_pairs) # 如果 clean 没有返回值
- print(f"Main 中获取的 cleaned_result: {cleaned_result_main}") # 如果有返回值
-
- print("在 main 中进行断言...")
- try:
- mock_infer_main.assert_called_once_with(
- inputs_main,
- config["model_name_or_path"],
- template=config["template"],
- temperature=0,
- )
- print("Main 中的 infer 函数调用断言成功！")
- if cleaned_result_main == mock_return: # 如果有返回值
- print("Main 中的返回值断言成功！")
- else:
- print(f"Main 中的返回值断言失败: 预期 {mock_return}, 得到 {cleaned_result_main}")
-
- except AssertionError as e:
- print(f"Main 中的 infer 函数调用断言失败: {e}")
-
-
- # # 方案2：手动调用测试函数（稍微复杂，需要手动创建 mock）
- # print("\\n--- 方案2：手动调用测试函数 ---")
- # # 创建一个 mock 对象手动传递
- # mock_infer_manual = MagicMock()
- # # 为手动创建的 mock 设置返回值 (如果需要)
- # mock_return_manual = ["debug_cleaned_1_manual", "debug_cleaned_2_manual"]
- # mock_infer_manual.return_value = mock_return_manual
- # print(f"设置 manual mock infer 返回值: {mock_return_manual}")
-
- # try:
- # print("手动调用 test_llm_cleaning_strategy_clean...")
- # # 注意：直接调用被 @patch 装饰的函数可能导致 TypeError
- # # 因为装饰器期望由测试运行器（如 pytest）注入 mock 对象
- # test_llm_cleaning_strategy_clean(mock_infer_manual, qa_pairs, config)
- # print("手动调用测试函数完成。请检查上面的打印输出。")
- # # 检查手动传入的 mock 是否被调用 (可能不会，因为 @patch 可能覆盖了它)
- # print("检查 manual mock 调用次数:", mock_infer_manual.call_count)
- # except TypeError as e:
- # print(f"\\n手动调用测试函数时捕获到 TypeError: {e}")
- # print("这通常发生在直接运行脚本时，@patch 装饰器未能正确处理 mock 注入。")
- # print("建议使用方案1（'with patch(...)' 上下文管理器）进行调试，因为它在 __main__ 块中更可靠。")
-
- print("\n调试脚本运行结束。")
\ No newline at end of file
diff --git a/tests/test_full_pipe.py b/tests/test_full_pipe.py
new file mode 100644
index 0000000..800126e
--- /dev/null
+++ b/tests/test_full_pipe.py
@@ -0,0 +1,154 @@
+import pytest
+from unittest import mock
+import sys
+import os
+import shutil
+import functools
+import subprocess
+import time
+from typing import Union, Optional, cast
+from weclone.utils.log import logger
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+PROJECT_ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+server_process: Optional[subprocess.Popen] = None
+
+test_logger = logger.bind()
+test_logger.remove()
+test_logger.add(
+ sys.stderr,
+ format="{message}",
+ colorize=True,
+ level="INFO",
+)
+
+def print_test_header(test_name: str):
+ line_length = 100
+ test_logger.info("\n" + "─" * line_length)
+ title = f" Testing Phase: {test_name} "
+ padding_total = line_length - len(title)
+ padding_left = padding_total // 2
+ padding_right = padding_total - padding_left
+ test_logger.info(" " * padding_left + title + " " * padding_right)
+ test_logger.info("─" * line_length)
+
+def setup_make_dataset_test_data():
+ PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+ DATASET_CSV_DIR = os.path.join(PROJECT_ROOT, "dataset", "csv")
+
+ TESTS_DIR = os.path.dirname(__file__)
+ TEST_DATA_PERSON_DIR = os.path.join(TESTS_DIR, "tests_data", "test_person")
+
+ os.makedirs(DATASET_CSV_DIR, exist_ok=True)
+
+ if os.path.exists(DATASET_CSV_DIR) and os.listdir(DATASET_CSV_DIR):
+ if all(f.startswith('.') or f.lower() == 'readme.md' for f in os.listdir(DATASET_CSV_DIR)):
+ for item_name in os.listdir(TEST_DATA_PERSON_DIR):
+ source_item_path = os.path.join(TEST_DATA_PERSON_DIR, item_name)
+ if os.path.isfile(source_item_path) and item_name.lower().endswith('.csv'):
+ destination_item_path = os.path.join(DATASET_CSV_DIR, item_name)
+ shutil.copy2(source_item_path, destination_item_path)
+
+
+def run_cli_command(command: list[str], timeout: int | None = None, background: bool = False) -> Union[subprocess.CompletedProcess, subprocess.Popen]:
+ """Execute a CLI command and return the result.
+
+ Args:
+ command: List of commands to execute.
+ timeout: Timeout in seconds.
+ background: Whether to run in the background.
+
+ Returns:
+ If background=True, returns a Popen object; otherwise, returns a CompletedProcess object.
+ """
+ env = os.environ.copy()
+ env["WECLONE_CONFIG_PATH"] = "tests/test.jsonc" # Set environment variable
+
+ if background:
+ process = subprocess.Popen(
+ [sys.executable, "-m", "weclone.cli"] + command,
+ stderr=subprocess.PIPE,
+ stdout=subprocess.PIPE,
+ text=True,
+ cwd=PROJECT_ROOT_DIR,
+ env=env
+ )
+ time.sleep(2)
+ return process
+ else:
+ process = subprocess.run(
+ [sys.executable, "-m", "weclone.cli"] + command,
+ stderr=None,
+ stdout=None,
+ text=True,
+ cwd=PROJECT_ROOT_DIR, # Execute in the project root directory
+ timeout=timeout,
+ env=env # Pass the modified environment variables
+ )
+ return process
+
+@pytest.mark.order(1)
+def test_cli_make_dataset():
+ """Test the make-dataset command."""
+ print_test_header("make-dataset")
+ setup_make_dataset_test_data()
+ result = run_cli_command(["make-dataset"])
+ assert result.returncode == 0, "make-dataset command execution failed"
+
+@pytest.mark.order(2)
+def test_cli_train_sft():
+ """Test the train-sft command."""
+ print_test_header("train-sft")
+ try:
+ result = run_cli_command(["train-sft"])
+ assert result.returncode == 0, "train-sft command failed or did not fail fast as expected"
+ except subprocess.TimeoutExpired:
+ test_logger.info("train-sft command terminated due to timeout, which is acceptable in testing, indicating the command has started execution.")
+ pass
+ except Exception as e:
+ pytest.fail(f"An unexpected error occurred during train-sft command execution: {e}")
+
+@pytest.mark.order(3)
+def test_cli_webchat_demo():
+ """Test the webchat-demo command."""
+ print_test_header("webchat-demo")
+
+ with mock.patch("weclone.eval.web_demo.main") as mock_main:
+ mock_main.return_value = None
+ try:
+ result = run_cli_command(["webchat-demo"], timeout=5)
+ assert result.returncode == 0, "webchat-demo command execution failed"
+ except subprocess.TimeoutExpired:
+ pass
+
+@pytest.mark.order(4)
+def test_cli_server():
+ """Test the server command.
+
+ Start the server in the background, without blocking subsequent tests.
+ """
+ print_test_header("server (background)")
+ global server_process
+ server_process = cast(subprocess.Popen, run_cli_command(["server"], background=True))
+ assert server_process.poll() is None, "Server startup failed"
+ test_logger.info("服务器已在后台启动")
+
+@pytest.mark.order(5)
+def test_cli_test_model():
+ """Test the test-model command.
+
+ Use the server for testing, and shut down the server after the test is complete.
+ """
+ print_test_header("test-model")
+ try:
+ result = run_cli_command(["test-model"])
+ assert result.returncode == 0, "test-model command execution failed"
+ finally:
+ global server_process
+ if server_process is not None and server_process.poll() is None:
+ test_logger.info("测试完成，正在关闭服务器...")
+ server_process.terminate()
+ server_process.wait(timeout=5)
+ if server_process.poll() is None:
+ server_process.kill() # Force kill if the process hasn't terminated
+ test_logger.info("服务器已关闭")
diff --git a/tests/test_full_pipeline.py b/tests/test_full_pipeline.py
deleted file mode 100644
index b33126b..0000000
--- a/tests/test_full_pipeline.py
+++ /dev/null
@@ -1,902 +0,0 @@
-import subprocess
-import sys
-import os
-import time
-import shutil
-import threading # 导入 threading
-from typing import Optional, Union, IO # 导入 IO
-import torch
-from loguru import logger
-from subprocess import Popen
-
-#TODO 放弃了改成测cli吧
-
-# 配置 Loguru
-logger.remove() # 移除默认处理器
-current_time = time.strftime('%Y%m%d_%H%M%S')
-log_file_path = os.path.join(os.path.dirname(__file__), f"pipeline_test_{current_time}.log") # 日志文件名包含执行时间
-logger.add(log_file_path, rotation="10 MB", encoding='utf-8', level="DEBUG", enqueue=True) # 文件记录 DEBUG 级别
-logger.add(sys.stdout, colorize=True, format="[test] {time:YYYY-MM-DD HH:mm:ss} | {level.name[0]} | {message}", level="INFO", enqueue=True) # 控制台保持 INFO 级别
-
-project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-logger.info(f"项目根目录: {project_root}")
-
-qa_script = "weclone/data/qa_generator.py"
-train_script = "weclone/train/train_sft.py"
-api_service_script = "weclone/server/api_service.py"
-eval_script = "weclone/eval/test_model.py"
-web_demo_script = "weclone/eval/web_demo.py"
-
-DEFAULT_TIMEOUT: Optional[Union[int, float]] = 45
-API_STARTUP_WAIT = 20
-API_TERMINATE_WAIT = 15
-WEB_DEMO_STARTUP_WAIT = 20
-WEB_DEMO_TERMINATE_WAIT = 15
-
-STEP_QA = "QA 数据生成"
-STEP_TRAIN = "SFT 训练"
-STEP_COPY_CKPT = "Checkpoint 复制"
-STEP_API_START = "API 服务启动"
-STEP_EVAL = "模型评估"
-STEP_WEB_DEMO = "Web Demo 启动"
-
-# Mapping from identifiers (script paths or custom keys) to step names
-step_identifiers = {
- qa_script: STEP_QA,
- train_script: STEP_TRAIN,
- "copy_checkpoint": STEP_COPY_CKPT, # Custom key for non-script step
- api_service_script: STEP_API_START, # Script associated with starting API
- eval_script: STEP_EVAL,
- web_demo_script: STEP_WEB_DEMO, # Script associated with starting Web Demo
-}
-# Order for fallback logic
-step_order = [STEP_QA, STEP_TRAIN, STEP_COPY_CKPT, STEP_API_START, STEP_EVAL, STEP_WEB_DEMO]
-
-#todo 需要测试前替换成测试的settings.jsonc 测试完再替换回来
-
-class PipelineStepError(Exception):
- """自定义异常类，用于表示 Pipeline 步骤执行失败。"""
- pass
-
-# --- 辅助函数：用于在线程中读取和记录流 ---
-def log_stream(stream: Optional[IO[str]], log_func):
- """读取流并使用指定的 log 函数记录每一行。"""
- if stream is None:
- return
- try:
- for line in iter(stream.readline, ''):
- if line:
- log_func(line.strip()) # 去除末尾换行符
- except ValueError:
- # 当 Popen 的 stream 在另一线程中被关闭时，readline 可能会抛出 ValueError
- logger.warning("日志流在读取时似乎已被关闭。")
- except Exception as e:
- # 捕获其他潜在的读取错误
- logger.warning(f"日志流读取时发生未预料的错误: {e}")
- finally:
- if stream:
- try:
- stream.close() # 确保流被关闭
- except Exception as close_e:
- logger.warning(f"关闭日志流时发生错误: {close_e}")
-
-# --- 新增：启动日志流线程的辅助函数 ---
-def _start_stream_logging_threads(process: Popen, stdout_log_func=logger.info, stderr_log_func=logger.error) -> tuple[threading.Thread, threading.Thread]:
- """为给定的进程启动 stdout 和 stderr 的日志记录线程。"""
- stdout_thread = threading.Thread(
- target=log_stream,
- args=(process.stdout, stdout_log_func),
- daemon=True
- )
- stderr_thread = threading.Thread(
- target=log_stream,
- args=(process.stderr, stderr_log_func),
- daemon=True
- )
- stdout_thread.start()
- stderr_thread.start()
- return stdout_thread, stderr_thread
-
-
-def run_script(script_relative_path: str, timeout: Optional[Union[int, float]] = DEFAULT_TIMEOUT, ignore_timeout_error: bool = False, env: Optional[dict] = None):
- """使用 Popen 执行脚本，通过线程实时记录 stdout/stderr 到 loguru。"""
- script_full_path = os.path.join(project_root, script_relative_path)
- timeout_str = '无限制' if timeout is None else f'{timeout}s'
- env_str = f" (环境变量: {env})" if env else ""
- logger.info(f"--- 开始执行 (流式): {script_relative_path} (超时: {timeout_str}){env_str} ---")
- if not os.path.exists(script_full_path):
- error_msg = f"脚本文件不存在 {script_full_path}"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
-
- process: Optional[Popen] = None
- stdout_thread: Optional[threading.Thread] = None
- stderr_thread: Optional[threading.Thread] = None
-
- # 准备环境变量
- run_env = os.environ.copy()
- if env:
- run_env.update(env)
-
- try:
- process = Popen(
- [sys.executable, script_full_path],
- cwd=project_root,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- encoding='utf-8',
- bufsize=1, # 行缓冲
- env=run_env # 传递环境变量
- )
-
- # 使用辅助函数启动日志线程
- stdout_thread, stderr_thread = _start_stream_logging_threads(process, logger.debug, logger.debug) # stdout/stderr 都用 debug
-
- # 等待子进程完成或超时
- try:
- return_code = process.wait(timeout=timeout)
- except subprocess.TimeoutExpired:
- warn_msg = f"{script_relative_path} 执行超时 ({timeout}s)。"
- logger.warning(warn_msg)
- # 尝试优雅地关闭流（可能已被 log_stream 关闭）
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill() # 强制终止超时进程
- logger.warning(f"已强制终止进程 {process.pid}")
- # 等待 I/O 线程完成（即使进程被 kill，也要尝试读取剩余输出）
- if stdout_thread: stdout_thread.join(timeout=5)
- if stderr_thread: stderr_thread.join(timeout=5)
- if not ignore_timeout_error:
- error_msg = f"{script_relative_path} 执行超时 ({timeout}s) 且未忽略。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- else:
- logger.info("--- 根据设置，超时不视为错误，继续执行后续步骤。 ---")
- return # 忽略超时，函数正常返回
-
- # 等待日志线程完成（确保所有输出都被记录）
- if stdout_thread: stdout_thread.join()
- if stderr_thread: stderr_thread.join()
-
- # 检查返回码
- if return_code != 0:
- error_msg = f"{script_relative_path} 执行失败，返回码 {return_code}"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- else:
- logger.success(f"--- {script_relative_path} 执行成功 ---")
-
- except FileNotFoundError:
- error_msg = f"Python 解释器 '{sys.executable}' 或脚本 '{script_full_path}' 未找到。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- except Exception as e:
- # 捕获其他潜在错误 (例如 Popen 本身失败)
- error_msg = f"执行 {script_relative_path} 时发生意外错误: {e}"
- logger.error(error_msg)
- # 尝试确保进程和线程被清理
- if process and process.poll() is None:
- try:
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill()
- logger.warning(f"因异常 {e}，强制终止进程 {process.pid}")
- except Exception as kill_e:
- logger.error(f"清理过程中强制终止进程失败: {kill_e}")
- if stdout_thread and stdout_thread.is_alive(): stdout_thread.join(timeout=1)
- if stderr_thread and stderr_thread.is_alive(): stderr_thread.join(timeout=1)
- raise PipelineStepError(error_msg)
-
-
-def start_api_service_background() -> Popen:
- """在后台启动 API 服务脚本，实时记录启动日志，失败时抛出 PipelineStepError。"""
- script_full_path = os.path.join(project_root, api_service_script)
- logger.info(f"--- 尝试在后台启动: {api_service_script} ---")
- if not os.path.exists(script_full_path):
- error_msg = f"脚本文件不存在 {script_full_path}"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
-
- process: Optional[Popen] = None
- stdout_thread: Optional[threading.Thread] = None
- stderr_thread: Optional[threading.Thread] = None
- try:
- logger.info(f"启动命令: {[sys.executable, script_full_path]}")
- process = Popen(
- [sys.executable, script_full_path],
- cwd=project_root,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- encoding='utf-8',
- bufsize=1 # 行缓冲
- )
-
- # 使用辅助函数启动日志线程
- stdout_thread, stderr_thread = _start_stream_logging_threads(process, logger.debug, logger.debug) # stdout/stderr 都用 debug
-
- logger.info(f"等待 {API_STARTUP_WAIT} 秒让服务初步启动 (日志将实时显示)...")
- time.sleep(API_STARTUP_WAIT)
-
- # 检查进程是否仍在运行
- if process.poll() is None:
- logger.success(f"--- {api_service_script} 似乎已在后台启动 (进程 PID: {process.pid}) ---")
- # 注意：不 join 日志线程，让它们继续运行
- return process
- else:
- # 进程过早退出
- logger.error(f"{api_service_script} 启动后在 {API_STARTUP_WAIT} 秒内过早退出，返回码 {process.returncode}")
- # 尝试等待日志线程结束以捕获最后输出
- if stdout_thread: stdout_thread.join(timeout=2)
- if stderr_thread: stderr_thread.join(timeout=2)
- # 读取 communicate 获取可能遗漏的最终输出 (虽然理论上线程应该读完了)
- try:
- # 设置短超时，因为进程已退出，communicate 应该立即返回
- stdout, stderr = process.communicate(timeout=1)
- except subprocess.TimeoutExpired:
- logger.warning("等待 communicate 超时，可能没有更多输出了。")
- stdout, stderr = "", "" # 假设没有更多输出
- except Exception as comm_e:
- logger.warning(f"调用 communicate 获取最后输出时出错: {comm_e}")
- stdout, stderr = "", ""
-
- error_message = f'''--- EARLY EXIT STDOUT ---
- {stdout}
- --- EARLY EXIT STDERR ---
- {stderr}'''
- logger.error(error_message)
- raise PipelineStepError(f"{api_service_script} 启动失败并过早退出。")
-
- except FileNotFoundError:
- error_msg = f"Python 解释器 '{sys.executable}' 或脚本 '{script_full_path}' 未找到。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- except Exception as e:
- # 捕获其他启动错误
- error_msg = f"启动 {api_service_script} 时发生意外错误: {e}"
- logger.error(error_msg)
- if process and process.poll() is None:
- logger.warning("捕获到异常，尝试强制终止进程...")
- try:
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill()
- except Exception as kill_e: logger.error(f"强制终止进程时出错: {kill_e}")
- # 尝试join线程
- if stdout_thread and stdout_thread.is_alive(): stdout_thread.join(timeout=1)
- if stderr_thread and stderr_thread.is_alive(): stderr_thread.join(timeout=1)
- raise PipelineStepError(error_msg)
-
-def stop_api_service(process: Optional[Popen]):
- """停止指定的 API 服务进程，采用更健壮的终止和清理逻辑。"""
- if process and process.poll() is None:
- pid = process.pid # Get PID for logging
- logger.info(f"--- 尝试停止 API 服务 (PID: {pid}) ---")
- try:
- logger.info(f"发送 SIGTERM 信号到进程 {pid}...")
- process.terminate()
- try:
- logger.info(f"等待最多 {API_TERMINATE_WAIT} 秒让进程 {pid} 优雅终止...")
- process.wait(timeout=API_TERMINATE_WAIT)
- logger.info(f"API 服务进程 {pid} 已优雅终止，返回码: {process.returncode}")
- # 进程已终止，尝试获取最终输出
- try:
- stdout, stderr = process.communicate(timeout=2)
- if stdout: logger.debug(f"进程 {pid} 最终 STDOUT:\n{stdout.strip()}")
- if stderr: logger.debug(f"进程 {pid} 最终 STDERR:\n{stderr.strip()}")
- except subprocess.TimeoutExpired:
- logger.warning(f"获取进程 {pid} 最终输出时超时。")
- except Exception as comm_e:
- logger.warning(f"获取进程 {pid} 最终输出时出错: {comm_e}")
-
- except subprocess.TimeoutExpired:
- logger.warning(f"进程 {pid} 优雅终止超时 ({API_TERMINATE_WAIT}s)，发送 SIGKILL 信号...")
- process.kill()
- logger.info(f"等待进程 {pid} 被强制终止...")
- # 在 kill 后等待，应该很快返回。增加安全超时。
- try:
- process.wait(timeout=5)
- logger.info(f"API 服务进程 {pid} 已被强制终止。")
- except subprocess.TimeoutExpired:
- logger.error(f"进程 {pid} 在发送 SIGKILL 后仍然没有终止！")
- except Exception as wait_kill_e:
- logger.error(f"等待强制终止进程 {pid} 时发生错误: {wait_kill_e}")
-
- # 尝试在 kill 后获取输出
- try:
- # 在 kill 后也使用 communicate，它隐式处理等待
- stdout, stderr = process.communicate(timeout=2)
- if stdout: logger.warning(f"来自进程 {pid} 的 Kill 后输出 (STDOUT):\n{stdout.strip()}")
- if stderr: logger.warning(f"来自进程 {pid} 的 Kill 后输出 (STDERR):\n{stderr.strip()}")
- except Exception as comm_e:
- logger.warning(f"获取进程 {pid} (强制终止后) 输出时出错: {comm_e}")
-
- except Exception as e:
- logger.error(f"停止 API 服务 (PID: {pid if process else '未知'}) 时发生意外错误: {e}")
- # 如果进程仍然存活，尝试最后一次强制 kill
- if process and process.poll() is None:
- logger.warning(f"最终尝试强制终止进程 {pid}...")
- try:
- process.kill()
- process.wait(timeout=5)
- except Exception as final_kill_e:
- logger.error(f"最终强制终止进程 {pid} 时出错: {final_kill_e}")
-
- elif process:
- logger.info(f"--- API 服务进程 (PID: {process.pid}) 在尝试停止前已经退出。 ---")
- else:
- logger.debug("--- 无需停止 API 服务 (进程不存在或已为 None) ---")
-
-def start_web_demo_background() -> Popen:
- """在后台启动 Web Demo 脚本，实时记录启动日志，失败时抛出 PipelineStepError。"""
- script_full_path = os.path.join(project_root, web_demo_script)
- logger.info(f"--- 尝试在后台启动: {web_demo_script} ---")
- if not os.path.exists(script_full_path):
- error_msg = f"脚本文件不存在 {script_full_path}"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
-
- process: Optional[Popen] = None
- stdout_thread: Optional[threading.Thread] = None
- stderr_thread: Optional[threading.Thread] = None
- try:
- logger.info(f"启动命令: {[sys.executable, script_full_path]}")
- process = Popen(
- [sys.executable, script_full_path],
- cwd=project_root,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- encoding='utf-8',
- bufsize=1 # 行缓冲
- )
-
- # 使用辅助函数启动日志线程 (stdout/stderr 都用 info)
- stdout_thread, stderr_thread = _start_stream_logging_threads(process, logger.debug, logger.debug) # stdout/stderr 都用 debug
-
-
- logger.info(f"等待 {WEB_DEMO_STARTUP_WAIT} 秒让 Web Demo 初步启动 (日志将实时显示)...")
- time.sleep(WEB_DEMO_STARTUP_WAIT)
-
- # 检查进程是否仍在运行
- if process.poll() is None:
- logger.success(f"--- {web_demo_script} 似乎已在后台启动 (进程 PID: {process.pid}) ---")
- # 注意：不 join 日志线程
- return process
- else:
- # 进程过早退出
- logger.error(f"{web_demo_script} 启动后在 {WEB_DEMO_STARTUP_WAIT} 秒内过早退出，返回码 {process.returncode}")
- if stdout_thread: stdout_thread.join(timeout=2)
- if stderr_thread: stderr_thread.join(timeout=2)
- try:
- stdout, stderr = process.communicate(timeout=1)
- except subprocess.TimeoutExpired:
- logger.warning("等待 communicate 超时，可能没有更多输出了。")
- stdout, stderr = "", ""
- except Exception as comm_e:
- logger.warning(f"调用 communicate 获取最后输出时出错: {comm_e}")
- stdout, stderr = "", ""
- error_message = f'''--- EARLY EXIT STDOUT ---
- {stdout}
- --- EARLY EXIT STDERR ---
- {stderr}'''
- logger.error(error_message)
- raise PipelineStepError(f"{web_demo_script} 启动失败并过早退出。")
-
- except FileNotFoundError:
- error_msg = f"Python 解释器 '{sys.executable}' 或脚本 '{script_full_path}' 未找到。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- except Exception as e:
- error_msg = f"启动 {web_demo_script} 时发生意外错误: {e}"
- logger.error(error_msg)
- if process and process.poll() is None:
- logger.warning("捕获到异常，尝试强制终止进程...")
- try:
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill()
- except Exception as kill_e: logger.error(f"强制终止进程时出错: {kill_e}")
- if stdout_thread and stdout_thread.is_alive(): stdout_thread.join(timeout=1)
- if stderr_thread and stderr_thread.is_alive(): stderr_thread.join(timeout=1)
- raise PipelineStepError(error_msg)
-
-def stop_web_demo(process: Optional[Popen]):
- """停止指定的 Web Demo 进程，采用更健壮的终止和清理逻辑。"""
- if process and process.poll() is None:
- pid = process.pid # Get PID for logging
- logger.info(f"--- 尝试停止 Web Demo 服务 (PID: {pid}) ---")
- try:
- logger.info(f"发送 SIGTERM 信号到进程 {pid}...")
- process.terminate()
- try:
- logger.info(f"等待最多 {WEB_DEMO_TERMINATE_WAIT} 秒让进程 {pid} 优雅终止...")
- process.wait(timeout=WEB_DEMO_TERMINATE_WAIT)
- logger.info(f"Web Demo 服务进程 {pid} 已优雅终止，返回码: {process.returncode}")
- # 进程已终止，尝试获取最终输出
- try:
- stdout, stderr = process.communicate(timeout=2)
- if stdout: logger.debug(f"进程 {pid} 最终 STDOUT:\n{stdout.strip()}")
- if stderr: logger.debug(f"进程 {pid} 最终 STDERR:\n{stderr.strip()}")
- except subprocess.TimeoutExpired:
- logger.warning(f"获取进程 {pid} 最终输出时超时。")
- except Exception as comm_e:
- logger.warning(f"获取进程 {pid} 最终输出时出错: {comm_e}")
-
- except subprocess.TimeoutExpired:
- logger.warning(f"进程 {pid} 优雅终止超时 ({WEB_DEMO_TERMINATE_WAIT}s)，发送 SIGKILL 信号...")
- process.kill()
- logger.info(f"等待进程 {pid} 被强制终止...")
- try:
- process.wait(timeout=5)
- logger.info(f"Web Demo 服务进程 {pid} 已被强制终止。")
- except subprocess.TimeoutExpired:
- logger.error(f"进程 {pid} 在发送 SIGKILL 后仍然没有终止！")
- except Exception as wait_kill_e:
- logger.error(f"等待强制终止进程 {pid} 时发生错误: {wait_kill_e}")
-
- # 尝试在 kill 后获取输出
- try:
- stdout, stderr = process.communicate(timeout=2)
- if stdout: logger.warning(f"来自进程 {pid} 的 Kill 后输出 (STDOUT):\n{stdout.strip()}")
- if stderr: logger.warning(f"来自进程 {pid} 的 Kill 后输出 (STDERR):\n{stderr.strip()}")
- except Exception as comm_e:
- logger.warning(f"获取进程 {pid} (强制终止后) 输出时出错: {comm_e}")
-
- except Exception as e:
- logger.error(f"停止 Web Demo 服务 (PID: {pid if process else '未知'}) 时发生意外错误: {e}")
- # 如果进程仍然存活，尝试最后一次强制 kill
- if process and process.poll() is None:
- logger.warning(f"最终尝试强制终止进程 {pid}...")
- try:
- process.kill()
- process.wait(timeout=5)
- except Exception as final_kill_e:
- logger.error(f"最终强制终止进程 {pid} 时出错: {final_kill_e}")
-
- elif process:
- logger.info(f"--- Web Demo 服务进程 (PID: {process.pid}) 在尝试停止前已经退出。 ---")
- else:
- logger.debug("--- 无需停止 Web Demo 服务 (进程不存在或已为 None) ---")
-
-# --- 新增：监控 Checkpoint 目录的函数 ---
-def monitor_checkpoints(process: Popen, model_output_dir: str, stop_event: threading.Event, check_interval: float = 5.0):
- """
- 在后台线程中监控指定目录，如果发现 checkpoint* 目录，则尝试终止目标进程。
- """
- logger.info(f"[Monitor] 开始监控目录 {model_output_dir} 的 checkpoint...")
- while not stop_event.is_set():
- if not os.path.isdir(model_output_dir):
- # 目录可能尚未创建，等待下一个间隔
- time.sleep(check_interval)
- continue
-
- try:
- found_checkpoint = False
- for item in os.listdir(model_output_dir):
- item_path = os.path.join(model_output_dir, item)
- if os.path.isdir(item_path) and item.startswith("checkpoint"):
- logger.warning(f"[Monitor] 检测到 Checkpoint 目录: {item_path}。尝试停止训练进程 (PID: {process.pid})...")
- found_checkpoint = True
- break # 找到一个就足够了
-
- if found_checkpoint:
- # 发送终止信号
- try:
- logger.info(f"[Monitor] 发送 SIGTERM 到进程 {process.pid}...")
- process.terminate()
- # 给进程一点时间响应 SIGTERM
- try:
- process.wait(timeout=5)
- logger.info(f"[Monitor] 进程 {process.pid} 已通过 SIGTERM 终止。")
- except subprocess.TimeoutExpired:
- logger.warning(f"[Monitor] 进程 {process.pid} 未在 5 秒内响应 SIGTERM，发送 SIGKILL...")
- process.kill()
- process.wait(timeout=5) # 等待 SIGKILL 生效
- logger.info(f"[Monitor] 进程 {process.pid} 已通过 SIGKILL 终止。")
- except Exception as term_err:
- logger.error(f"[Monitor] 尝试终止进程 {process.pid} 时出错: {term_err}")
- finally:
- stop_event.set() # 通知主线程停止等待
- logger.info("[Monitor] 已设置停止事件，监控结束。")
- return # 找到 checkpoint 并处理后，监控任务完成
-
- except FileNotFoundError:
- # 目录可能在检查时被删除，忽略
- pass
- except Exception as e:
- logger.error(f"[Monitor] 监控时发生错误: {e}")
- # 出现错误也设置停止信号，防止无限循环或未处理的异常
- stop_event.set()
- return
-
- # 如果没有找到，且进程仍在运行，则等待下一个检查周期
- if process.poll() is None:
- time.sleep(check_interval)
- else:
- # 如果进程已经结束（无论何种原因），监控也应结束
- logger.info(f"[Monitor] 训练进程 {process.pid} 似乎已结束，停止监控。")
- break # 进程已结束，退出循环
- logger.info("[Monitor] 监控循环正常结束。")
-
-
-# --- 新增：运行训练并进行监控的函数 ---
-def run_train_with_checkpoint_monitoring(
- script_relative_path: str,
- model_output_dir: str,
- timeout: Optional[Union[int, float]] = DEFAULT_TIMEOUT,
- ignore_timeout_error: bool = False,
- env: Optional[dict] = None
-) -> str:
- """
- 执行训练脚本，同时启动一个后台线程监控 checkpoint 目录。
- 如果检测到 checkpoint，会尝试停止训练进程。
- 返回执行状态: "success", "stopped_by_monitor", "timeout", "failed"
- """
- script_full_path = os.path.join(project_root, script_relative_path)
- timeout_str = '无限制' if timeout is None else f'{timeout}s'
- env_str = f" (环境变量: {env})" if env else ""
- logger.info(f"--- 开始执行 (带监控): {script_relative_path} (超时: {timeout_str}){env_str} ---")
- if not os.path.exists(script_full_path):
- error_msg = f"脚本文件不存在 {script_full_path}"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
-
- process: Optional[Popen] = None
- stdout_thread: Optional[threading.Thread] = None
- stderr_thread: Optional[threading.Thread] = None
- monitor_thread: Optional[threading.Thread] = None
- stop_event = threading.Event()
- status = "failed" # 默认状态
-
- # 准备环境变量
- run_env = os.environ.copy()
- if env:
- run_env.update(env)
-
- try:
- process = Popen(
- [sys.executable, script_full_path],
- cwd=project_root,
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- encoding='utf-8',
- bufsize=1, # 行缓冲
- env=run_env
- )
-
- # 启动日志线程
- stdout_thread, stderr_thread = _start_stream_logging_threads(process, logger.debug, logger.debug)
-
- # 启动监控线程
- monitor_thread = threading.Thread(
- target=monitor_checkpoints,
- args=(process, model_output_dir, stop_event),
- daemon=True
- )
- monitor_thread.start()
-
- # 等待进程完成、被监控停止或超时
- start_time = time.time()
- wait_interval = 1 # seconds to wait between checks
- while True:
- # 检查进程是否结束
- return_code = process.poll()
- if return_code is not None:
- # 进程已结束
- if stop_event.is_set(): # 如果是监控线程停止的
- logger.warning(f"{script_relative_path} 被监控线程停止。返回码可能为 {return_code}。")
- status = "stopped_by_monitor"
- elif return_code == 0:
- logger.success(f"{script_relative_path} 成功完成。")
- status = "success"
- else:
- logger.error(f"{script_relative_path} 执行失败，返回码 {return_code}")
- status = "failed"
- stop_event.set() # 确保监控线程也会退出
- break
-
- # 检查是否被监控线程要求停止
- if stop_event.is_set():
- logger.warning(f"{script_relative_path} 被监控线程标记为停止。")
- # 进程可能仍在运行，等待 monitor_checkpoints 中的终止逻辑生效
- # 但我们这里也应该退出等待循环
- status = "stopped_by_monitor"
- # 不需要再次 kill，monitor 线程会处理
- break
-
- # 检查是否超时
- if timeout is not None and (time.time() - start_time) > timeout:
- warn_msg = f"{script_relative_path} 执行超时 ({timeout}s)。"
- logger.warning(warn_msg)
- stop_event.set() # 通知监控线程停止
- # 尝试优雅地关闭流
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill() # 强制终止超时进程
- logger.warning(f"已强制终止进程 {process.pid}")
- status = "timeout"
- break
-
- # 短暂休眠后继续检查
- time.sleep(wait_interval)
-
- # --- 等待所有线程完成 ---
- logger.info("等待日志和监控线程完成...")
- if stdout_thread: stdout_thread.join(timeout=5)
- if stderr_thread: stderr_thread.join(timeout=5)
- if monitor_thread: monitor_thread.join(timeout=5) # 监控线程也需要 join
-
- # 处理最终状态
- if status == "failed":
- raise PipelineStepError(f"{script_relative_path} 执行失败。")
- elif status == "timeout":
- if not ignore_timeout_error:
- error_msg = f"{script_relative_path} 执行超时 ({timeout}s) 且未忽略。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- else:
- logger.info("--- 根据设置，超时不视为错误，继续执行后续步骤。 ---")
- # 即使忽略超时错误，状态仍然是 "timeout"
- return status # 返回 "timeout" 状态
-
- # 对于 success 和 stopped_by_monitor，直接返回状态
- logger.info(f"--- {script_relative_path} 执行结束，状态: {status} ---")
- return status
-
- except FileNotFoundError:
- error_msg = f"Python 解释器 '{sys.executable}' 或脚本 '{script_full_path}' 未找到。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- except Exception as e:
- # 捕获其他潜在错误 (例如 Popen 本身失败)
- error_msg = f"执行 {script_relative_path} (带监控) 时发生意外错误: {e}"
- logger.error(error_msg)
- stop_event.set() # 确保监控线程停止
- # 尝试确保进程和线程被清理
- if process and process.poll() is None:
- try:
- if process.stdout: process.stdout.close()
- if process.stderr: process.stderr.close()
- process.kill()
- logger.warning(f"因异常 {e}，强制终止进程 {process.pid}")
- except Exception as kill_e:
- logger.error(f"清理过程中强制终止进程失败: {kill_e}")
- if stdout_thread and stdout_thread.is_alive(): stdout_thread.join(timeout=1)
- if stderr_thread and stderr_thread.is_alive(): stderr_thread.join(timeout=1)
- if monitor_thread and monitor_thread.is_alive(): monitor_thread.join(timeout=1)
- raise PipelineStepError(error_msg)
-
-
-if __name__ == "__main__":
- logger.info("="*20 + " 开始执行 WeClone Pipeline 脚本 " + "="*20)
-
- is_cuda_available = torch.cuda.is_available()
- logger.info("--- CUDA 可用性检查 ---")
- if is_cuda_available:
- gpu_count = torch.cuda.device_count()
- logger.success(f"CUDA 可用 (找到 {gpu_count} 个 GPU)")
- for i in range(gpu_count):
- logger.info(f" - GPU {i}: {torch.cuda.get_device_name(i)}")
- else:
- logger.warning("CUDA 不可用，将使用 CPU (如果适用)。")
- logger.info("-" * 25)
-
- steps_completed = []
- api_process: Optional[Popen] = None
- web_demo_process: Optional[Popen] = None
-
- # 设置哪些步骤需要运行
- run_qa = True
- run_train = True
- run_copy_checkpoint = True # 依赖于 run_train
- run_api = True
- run_eval = True # 依赖于 run_api
- run_web_demo = True # 不依赖于 run_api
-
- try:
- # 步骤 1: QA Generator
- if run_qa:
- logger.info("-" * 10 + " 步骤 1: QA 数据生成 " + "-" * 10)
- run_script(qa_script)
- steps_completed.append(f"{STEP_QA}: 成功")
- else:
- logger.info(f"{STEP_QA}: 跳过 (配置)")
- steps_completed.append(f"{STEP_QA}: 跳过")
-
- # 步骤 2: Train SFT (with monitoring)
- if run_train:
- logger.info("-" * 10 + " 步骤 2: SFT 训练 (带 Checkpoint 监控) " + "-" * 10)
- model_output_dir = os.path.join(project_root, "model_output")
-
- # --- 删除 model_output 目录 ---
- if os.path.exists(model_output_dir):
- logger.info(f"删除现有的 model_output 目录: {model_output_dir}")
- try:
- shutil.rmtree(model_output_dir)
- logger.success("成功删除 model_output 目录")
- except Exception as e:
- logger.error(f"删除 model_output 目录时出错: {e}")
- # Treat failure to delete as a critical error before training
- raise PipelineStepError(f"删除 model_output 目录失败: {e} ###step_id:{train_script}###")
-
- # --- 执行训练脚本并进行监控 ---
- train_status = run_train_with_checkpoint_monitoring(
- train_script,
- model_output_dir, # Pass the directory to monitor
- timeout=2000,
- ignore_timeout_error=True,
- env={'TQDM_DISABLE': '1'}
- )
-
- # --- 根据训练状态更新完成列表 ---
- if train_status == "success":
- steps_completed.append(f"{STEP_TRAIN}: 成功")
- elif train_status == "stopped_by_monitor":
- steps_completed.append(f"{STEP_TRAIN}: 已停止 (检测到 Checkpoint)")
- elif train_status == "timeout":
- steps_completed.append(f"{STEP_TRAIN}: 超时 (已忽略)")
- else: # "failed" or other unexpected status handled by exception
- steps_completed.append(f"{STEP_TRAIN}: 失败") # Should be caught by exception, but added for completeness
-
- # 步骤 2.1: 复制 Checkpoint (只有在训练 *成功* 完成后才执行)
- if run_copy_checkpoint:
- if train_status == "success":
- logger.info("-" * 10 + " 步骤 2.1: 复制 Checkpoint 到 model_output " + "-" * 10)
- source_dir = os.path.join(project_root, "model_output", "checkpoint-2") # Note: Still assumes checkpoint-2 specifically.
- dest_dir = os.path.join(project_root, "model_output")
- if os.path.isdir(source_dir):
- try:
- logger.info(f"开始将 {source_dir} 的内容复制到 {dest_dir}...")
- shutil.copytree(source_dir, dest_dir, dirs_exist_ok=True)
- logger.success(f"--- {STEP_COPY_CKPT} 成功 ---")
- steps_completed.append(f"{STEP_COPY_CKPT}: 成功")
- except Exception as e:
- error_msg = f"{STEP_COPY_CKPT} 时发生错误: {e}"
- logger.error(error_msg)
- raise PipelineStepError(f"{error_msg} ###step_id:copy_checkpoint###")
- else:
- logger.warning(f"训练成功后，源 Checkpoint 目录 {source_dir} 不存在或不是目录，跳过复制。")
- steps_completed.append(f"{STEP_COPY_CKPT}: 跳过 (源不存在)")
- # Consider if missing checkpoint-2 after successful training is an error
- # raise PipelineStepError(f"训练成功但必需的源 Checkpoint 目录 {source_dir} 不存在")
- else:
- # If training didn't succeed (stopped, timeout, failed), skip copy
- logger.info(f"{STEP_COPY_CKPT}: 跳过 (训练未成功完成，状态: {train_status})")
- steps_completed.append(f"{STEP_COPY_CKPT}: 跳过 (训练未成功)")
- else:
- logger.info(f"{STEP_COPY_CKPT}: 跳过 (配置)")
- steps_completed.append(f"{STEP_COPY_CKPT}: 跳过 (配置)")
-
- else:
- # If run_train is false
- logger.info(f"{STEP_TRAIN}: 跳过 (配置)")
- steps_completed.append(f"{STEP_TRAIN}: 跳过 (配置)")
- logger.info(f"{STEP_COPY_CKPT}: 跳过 (训练未运行)")
- steps_completed.append(f"{STEP_COPY_CKPT}: 跳过 (训练未运行)")
-
- # 步骤 3: Start API Service
- if run_api:
- logger.info("-" * 10 + " 步骤 3: 启动 API 服务 " + "-" * 10)
- api_process = start_api_service_background()
- steps_completed.append(f"{STEP_API_START}: 成功")
- else:
- logger.info(f"{STEP_API_START}: 跳过 (配置)")
- steps_completed.append(f"{STEP_API_START}: 跳过 (配置)")
-
-
- # 步骤 4: Eval Model (依赖 API 服务)
- if run_eval:
- if not run_api:
- logger.info("-" * 10 + " 步骤 4: 模型评估 " + "-" * 10)
- logger.warning("--- 因 API 服务配置为不运行，跳过执行: weclone/eval/test_model.py ---")
- steps_completed.append(f"{STEP_EVAL}: 跳过 (API未配置运行)")
- elif api_process is None or api_process.poll() is not None: # 检查进程是否已退出
- error_msg = "尝试运行评估，但 API 服务进程不存在或已退出。"
- logger.error(error_msg)
- raise PipelineStepError(error_msg)
- else:
- logger.info("-" * 10 + " 步骤 4: 模型评估 " + "-" * 10)
- # 在调用评估脚本时禁用 tqdm
- run_script(eval_script, timeout=9999, env={'TQDM_DISABLE': '1'})
- steps_completed.append(f"{STEP_EVAL}: 成功")
- stop_api_service(api_process) # 评估完成后停止API
- api_process = None # 标记为已停止
- else:
- logger.info(f"{STEP_EVAL}: 跳过 (配置)")
- steps_completed.append(f"{STEP_EVAL}: 跳过 (配置)")
- if api_process: # 如果API在运行但评估被跳过，也停止API
- logger.info("评估被跳过，停止 API 服务...")
- stop_api_service(api_process)
- api_process = None
-
-
- # 步骤 5: Start Web Demo (不依赖 API 服务)
- if run_web_demo:
- logger.info("-" * 10 + " 步骤 5: 启动 Web Demo " + "-" * 10)
- web_demo_process = start_web_demo_background()
- steps_completed.append(f"{STEP_WEB_DEMO}: 成功")
- logger.info("--- Web Demo 已启动，测试流程继续... ---")
- else:
- logger.info(f"{STEP_WEB_DEMO}: 跳过 (配置)")
- steps_completed.append(f"{STEP_WEB_DEMO}: 跳过 (配置)")
-
- # Pipeline 成功完成所有请求的步骤
- logger.info("="*20 + " Pipeline 执行摘要 " + "="*20)
- for step in steps_completed:
- logger.info(f"- {step}")
- logger.success("✅ 所有请求执行的 Pipeline 步骤均成功完成！")
-
- skipped_steps = [s for s in steps_completed if "跳过" in s]
- if skipped_steps:
- logger.warning("注意: 以下步骤被设置为跳过，如需执行请修改脚本顶部的 run_xxx 变量：")
- for skipped in skipped_steps:
- logger.warning(f" - {skipped.split(':')[0]}")
-
-
- except PipelineStepError as e:
- logger.error("="*20 + " Pipeline 执行失败 " + "="*20)
-
- failing_step = "未知步骤"
- error_details = str(e)
- cleaned_error_details = error_details # Store original/cleaned details for logging
-
- # Attempt 1: Check for explicit marker (e.g., from copy_checkpoint)
- marker_prefix = "###step_id:"
- marker_suffix = "###"
- marker_start = error_details.find(marker_prefix)
- if marker_start != -1:
- marker_end = error_details.find(marker_suffix, marker_start + len(marker_prefix))
- if marker_end != -1:
- step_id = error_details[marker_start + len(marker_prefix):marker_end]
- failing_step = step_identifiers.get(step_id, f"未知标记 ({step_id})")
- # Clean the marker from the displayed error details
- cleaned_error_details = error_details[:marker_start].strip()
-
- # Attempt 2: Check for known script paths in the error message if marker not found
- if failing_step == "未知步骤":
- found_script = False
- # Iterate through potential script paths stored as keys in step_identifiers
- for identifier, step_name in step_identifiers.items():
- # Check if the identifier looks like a path and is in the error message
- if isinstance(identifier, str) and ('/' in identifier or '\\\\' in identifier) and identifier in error_details:
- failing_step = step_name
- found_script = True
- break # Found the most likely script
-
- # Attempt 3: Fallback based on last completed step (if still unknown)
- if failing_step == "未知步骤":
- if steps_completed:
- last_completed = steps_completed[-1].split(':')[0]
- try:
- last_completed_index = step_order.index(last_completed)
- if last_completed_index + 1 < len(step_order):
- # Assume the next step in the defined order failed
- failing_step = step_order[last_completed_index + 1] + " (推断)"
- else:
- failing_step = "Pipeline末尾或未知 (推断)" # Error after the last known step
- except ValueError:
- # Last completed step name wasn't found in our defined order
- failing_step = f"未知 (最后完成: {last_completed})"
- else:
- failing_step = "初始化期间" # No steps completed
-
- logger.error(f"错误发生在步骤: {failing_step}")
- logger.error(f"错误详情: {cleaned_error_details}") # Log the cleaned error details
- logger.info("--- 已完成步骤 ---")
- for step in steps_completed:
- logger.info(f"- {step}")
- logger.error("="*50)
- sys.exit(1) # 测试失败时退出码为 1
-
- finally:
- logger.info("--- Pipeline 结束，开始清理后台服务 ---")
- # 确保在 finally 块中总是尝试停止服务
- stop_web_demo(web_demo_process)
- stop_api_service(api_process) # 即使评估步骤停止了它，这里也尝试停止，无害
- logger.info("--- 后台服务清理完成 ---")
-
- # 如果 Pipeline 成功，确保退出码为 0
- sys.exit(0)
\ No newline at end of file
diff --git a/tests/test_get_sample_audio.py b/tests/test_get_sample_audio.py
deleted file mode 100644
index 75e3f5a..0000000
--- a/tests/test_get_sample_audio.py
+++ /dev/null
@@ -1,151 +0,0 @@
-import os
-import subprocess
-import sys
-import pytest
-
-# 获取 weclone-audio/src 目录的绝对路径
-# 这假设 tests 目录和 weclone-audio 在同一个父目录下
-SRC_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'weclone-audio', 'src'))
-SCRIPT_PATH = os.path.join(SRC_DIR, 'get_sample_audio.py')
-
-# --- 测试配置 ---
-# 请将下面的路径替换为你的测试数据库文件的实际路径
-# 最好放在 tests/data 目录下，并使用相对路径
-TEST_DB_PATH = r"D:\projects\python projects\WeClone-data\wxdump_work\wxid_d6wwiru2zsmo22\merge_all.db"# <--- 修改这里
-# 请将下面的 ID 替换为测试数据库中一个有效的音频消息的 MsgSvrID
-TEST_MSG_SVR_ID = "3269716813078873653" # <--- 修改这里
-# ----------------
-
-@pytest.fixture(scope="module")
-def setup_test_environment():
- """确保测试所需的文件和目录存在"""
- if not os.path.exists(TEST_DB_PATH):
- pytest.fail(f"测试数据库文件未找到: {TEST_DB_PATH}。请提供一个有效的测试数据库。")
- if not os.path.exists(SCRIPT_PATH):
- pytest.fail(f"待测试的脚本未找到: {SCRIPT_PATH}")
- # 可以添加其他设置，例如创建测试数据目录
-
-def test_audio_extraction(tmp_path, setup_test_environment):
- """
- 测试 get_sample_audio.py 是否能成功提取音频并保存为 wav 文件。
- """
- output_filename = "test_output.wav"
- output_path = tmp_path / output_filename # 使用 pytest 的 tmp_path fixture 创建临时输出路径
-
- # 构建命令行参数
- cmd = [
- sys.executable, # 使用当前的 Python 解释器
- SCRIPT_PATH,
- "--db-path", TEST_DB_PATH,
- "--MsgSvrID", TEST_MSG_SVR_ID,
- "--save-path", str(output_path),
- "--rate", "24000" # 可以根据需要调整
- ]
-
- # 运行脚本
- # 注意：脚本中的 'key' 可能需要根据实际情况调整，或者修改脚本以允许通过参数传递 key
- # 目前脚本中硬编码了 key="test1"
- result = subprocess.run(cmd, capture_output=True, text=True, check=False) # check=False 允许我们检查返回码
-
- # 打印输出以便调试 (如果测试失败)
- print("STDOUT:", result.stdout)
- print("STDERR:", result.stderr)
-
- # 断言脚本成功运行
- assert result.returncode == 0, f"脚本执行失败，错误信息: {result.stderr}"
-
- # 断言输出文件已创建
- assert output_path.exists(), f"输出文件 {output_path} 未被创建"
-
- # (可选) 断言文件大小大于 0
- assert output_path.stat().st_size > 0, f"输出文件 {output_path} 为空"
-
- # (可选) 更复杂的检查，例如使用 wave 库检查文件头或内容
- # import wave
- # try:
- # with wave.open(str(output_path), 'rb') as wf:
- # assert wf.getnchannels() == 1 # 假设是单声道
- # assert wf.getframerate() == 24000 # 检查采样率
- # except wave.Error as e:
- # pytest.fail(f"无法读取输出的 WAV 文件: {e}")
-
-def main_debug():
- """用于直接运行和调试的主要函数"""
- print("--- 开始调试运行 ---")
-
- # 检查基本环境
- if not os.path.exists(TEST_DB_PATH):
- print(f"错误: 测试数据库文件未找到: {TEST_DB_PATH}")
- return
- if not os.path.exists(SCRIPT_PATH):
- print(f"错误: 待测试的脚本未找到: {SCRIPT_PATH}")
- return
- if TEST_MSG_SVR_ID == "YOUR_TEST_MSG_SVR_ID":
- print(f"警告: TEST_MSG_SVR_ID 似乎未配置 ({TEST_MSG_SVR_ID})")
- # 可以选择在这里 return 或继续执行
-
- # 定义调试输出路径
- debug_output_dir = os.path.join(os.path.dirname(__file__), "debug_output")
- os.makedirs(debug_output_dir, exist_ok=True) # 创建输出目录（如果不存在）
- debug_output_path = os.path.join(debug_output_dir, "debug_sample.wav")
-
- print(f"脚本路径: {SCRIPT_PATH}")
- print(f"数据库路径: {TEST_DB_PATH}")
- print(f"消息 ID: {TEST_MSG_SVR_ID}")
- print(f"输出路径: {debug_output_path}")
-
- # 构建命令行参数
- cmd = [
- sys.executable,
- SCRIPT_PATH,
- "--db-path", TEST_DB_PATH,
- "--MsgSvrID", TEST_MSG_SVR_ID,
- "--save-path", debug_output_path,
- "--rate", "24000"
- ]
-
- print(f"执行命令: {' '.join(cmd)}")
-
- # 运行脚本
- try:
- result = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=30) # 添加超时
- print("\\n--- 脚本执行结果 ---")
- print("返回码:", result.returncode)
- print("STDOUT:")
- print(result.stdout)
- print("STDERR:")
- print(result.stderr)
-
- # 检查结果
- if result.returncode == 0:
- print("\\n--- 结果检查 ---")
- if os.path.exists(debug_output_path):
- print(f"[成功] 输出文件已创建: {debug_output_path}")
- if os.path.getsize(debug_output_path) > 0:
- print(f"[成功] 输出文件大小 > 0 ({os.path.getsize(debug_output_path)} bytes)")
- else:
- print(f"[失败] 输出文件为空: {debug_output_path}")
- else:
- print(f"[失败] 输出文件未找到: {debug_output_path}")
- else:
- print("\\n[失败] 脚本执行失败。")
-
- except subprocess.TimeoutExpired:
- print("\\n[失败] 脚本执行超时。")
- except Exception as e:
- print(f"\\n[失败] 执行命令时发生异常: {e}")
-
- print("\\n--- 调试运行结束 ---")
-
-
-if __name__ == "__main__":
- # 确保在直接运行时正确设置了测试数据路径
- # 注意：这里仍然使用文件顶部的 TEST_DB_PATH 和 TEST_MSG_SVR_ID
- # 请确保它们已经被修改为有效值！
- if TEST_DB_PATH == "tests/data/your_test_db.sqlite" or TEST_MSG_SVR_ID == "YOUR_TEST_MSG_SVR_ID":
- print("*"*40)
- print("警告：请先在脚本顶部修改 TEST_DB_PATH 和 TEST_MSG_SVR_ID 为有效的测试值！")
- print("*"*40)
- # sys.exit(1) # 可以取消注释以强制退出，如果未配置
-
- main_debug()
\ No newline at end of file
diff --git a/tests/test_get_score.py b/tests/test_get_score.py
deleted file mode 100644
index 8a2eb27..0000000
--- a/tests/test_get_score.py
+++ /dev/null
@@ -1,72 +0,0 @@
-import pytest
-from weclone.data.clean.get_score import adjust_score_tiered
-
-# 定义通用的参数
-THRESHOLDS = [0.6, 0.3] # 置信度阈值：>=0.6 高, >=0.3 中, <0.3 低
-DOWNGRADE_LEVELS = [0, 1, 2] # 对应降级幅度：高->0级, 中->1级, 低->2级
-
-THRESHOLDS_FINE = [0.7, 0.5, 0.3]
-DOWNGRADE_LEVELS_FINE = [0, 1, 2, 3] # 对应 >=0.7, >=0.5, >=0.3, <0.3
-
-test_cases = [
- # 案例 1: 高置信度
- (5, [0.05, 0.05, 0.1, 0.1, 0.7], THRESHOLDS, DOWNGRADE_LEVELS, 5, "高置信度"),
- # 案例 2: 中等置信度
- (4, [0.1, 0.15, 0.2, 0.45, 0.1], THRESHOLDS, DOWNGRADE_LEVELS, 3, "中等置信度"),
- # 案例 3: 低置信度
- (4, [0.15, 0.2, 0.25, 0.25, 0.15], THRESHOLDS, DOWNGRADE_LEVELS, 2, "低置信度"),
- # 案例 4: 低置信度，但原始分较低
- (2, [0.3, 0.2, 0.2, 0.15, 0.15], THRESHOLDS, DOWNGRADE_LEVELS, 1, "低置信度，原始分较低"),
- # 案例 5: 边界情况 - 刚好等于高阈值
- (3, [0.1, 0.1, 0.6, 0.1, 0.1], THRESHOLDS, DOWNGRADE_LEVELS, 3, "边界情况 - 等于高阈值"),
- # 案例 6: 边界情况 - 刚好等于中阈值
- (3, [0.2, 0.2, 0.3, 0.15, 0.15], THRESHOLDS, DOWNGRADE_LEVELS, 2, "边界情况 - 等于中阈值"),
- # 案例 7: 细分阈值 - 中高置信度
- (4, [0.1, 0.1, 0.2, 0.55, 0.05], THRESHOLDS_FINE, DOWNGRADE_LEVELS_FINE, 3, "细分阈值 - 中高置信度"),
- # 案例 8: 细分阈值 - 中低置信度
- (4, [0.15, 0.15, 0.2, 0.35, 0.15], THRESHOLDS_FINE, DOWNGRADE_LEVELS_FINE, 2, "细分阈值 - 中低置信度"),
- # 案例 9: 概率和异常 (预期行为是打印警告并继续计算)
- (3, [0.1, 0.1, 0.5, 0.1, 0.1], THRESHOLDS, DOWNGRADE_LEVELS, 3, "概率和异常"),
-]
-
-@pytest.mark.parametrize("initial_score, probabilities, thresholds, downgrade_levels, expected_score, description", test_cases)
-def test_adjust_score_tiered(initial_score, probabilities, thresholds, downgrade_levels, expected_score, description):
- """ 测试 adjust_score_tiered 函数在各种情况下的表现 """
- print(f"测试案例: {description}")
- print(f" 输入: score={initial_score}, probs={probabilities}, thresholds={thresholds}, levels={downgrade_levels}")
- adjusted_score = adjust_score_tiered(initial_score, probabilities, thresholds, downgrade_levels)
- print(f" 输出: adjusted_score={adjusted_score}, 预期: {expected_score}")
- assert adjusted_score == expected_score
-
-# 测试非法输入
-def test_adjust_score_invalid_input():
- """ 测试非法输入是否按预期引发 ValueError """
- # initial_score 无效
- with pytest.raises(ValueError, match="initial_score 必须在 1 到 5 之间"):
- adjust_score_tiered(0, [0.2]*5, THRESHOLDS, DOWNGRADE_LEVELS)
- with pytest.raises(ValueError, match="initial_score 必须在 1 到 5 之间"):
- adjust_score_tiered(6, [0.2]*5, THRESHOLDS, DOWNGRADE_LEVELS)
-
- # probabilities 长度无效
- with pytest.raises(ValueError, match="probabilities 列表必须包含 5 个元素"):
- adjust_score_tiered(3, [0.2]*4, THRESHOLDS, DOWNGRADE_LEVELS)
- with pytest.raises(ValueError, match="probabilities 列表必须包含 5 个元素"):
- adjust_score_tiered(3, [0.1]*6, THRESHOLDS, DOWNGRADE_LEVELS) # 总和也不为1
-
- # # probabilities 和不为 1 (现在是警告，不抛异常)
- # with pytest.raises(ValueError, match="probabilities 中元素的和必须接近 1.0"):
- # adjust_score_tiered(3, [0.1]*5, THRESHOLDS, DOWNGRADE_LEVELS)
-
- # downgrade_levels 长度无效
- with pytest.raises(ValueError, match="downgrade_levels 的长度必须比 thresholds 的长度多 1"):
- adjust_score_tiered(3, [0.2]*5, THRESHOLDS, [0, 1])
- with pytest.raises(ValueError, match="downgrade_levels 的长度必须比 thresholds 的长度多 1"):
- adjust_score_tiered(3, [0.2]*5, THRESHOLDS, [0, 1, 2, 3])
-
- # thresholds 不是降序
- with pytest.raises(ValueError, match="thresholds 列表必须是降序排列的"):
- adjust_score_tiered(3, [0.2]*5, [0.3, 0.6], DOWNGRADE_LEVELS)
-
- # downgrade_levels 包含负数
- with pytest.raises(ValueError, match="downgrade_levels 中的降级幅度不能为负数"):
- adjust_score_tiered(3, [0.2]*5, THRESHOLDS, [0, -1, 2])
\ No newline at end of file
diff --git a/tests/test_old_csv_to_json.py b/tests/test_old_csv_to_json.py
deleted file mode 100644
index e386e51..0000000
--- a/tests/test_old_csv_to_json.py
+++ /dev/null
@@ -1,312 +0,0 @@
-import csv
-import json
-import os
-import re
-import sys
-
-import pandas as pd
-from collections import deque
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-root_dir = os.path.dirname(current_dir)
-sys.path.append(root_dir)
-
-from make_dataset.qa_generator import DataProcessor
-
-csv_folder = "./data/csv"
-# csv_folder = './data/test'
-os.chdir(root_dir)
-
-print(f"当前处理目录{csv_folder}")
-
-
-def handle_pt_csv(csvfile):
- chat_df = pd.read_csv(csvfile)
- # 选择type_name为文本的行、is_sender为1的行
- chat_df = chat_df[chat_df["type_name"] == "文本"]
- chat_df = chat_df[chat_df["is_sender"] == 1]
- # 对每一行的content进行处理 转为dict 再取'msg'字段
- chat_df["content"] = chat_df["content"].apply(lambda x: json.loads(x)["msg"])
- # 如果content 包含 手机号、身份证号、邮箱、网址则删除这行
- chat_df = chat_df[~chat_df["content"].str.contains("1\d{10}")]
- chat_df = chat_df[~chat_df["content"].str.contains("\d{18}")]
- chat_df = chat_df[~chat_df["content"].str.contains("\w+@\w+")]
- chat_df = chat_df[~chat_df["content"].str.contains("http")]
- chat_df = chat_df[~chat_df["content"].str.contains(r"\\xa0")]
- chat_df = chat_df[~chat_df["content"].str.contains(r"\\u")]
-
- # 纯content
- chat_df = chat_df["content"]
- chat_df = chat_df.dropna()
-
- return chat_df
-
-
-def make_pt_dataset():
- csv_res = []
- # csv文件夹里全是不同聊天对象文件夹 每个文件夹里是csv文件 先遍历不同聊天对象文件夹 再遍历聊天对象的csv文件
- for chat_obj_folder in os.listdir(csv_folder):
- chat_obj_folder_path = os.path.join(csv_folder, chat_obj_folder)
- for csvfile in os.listdir(chat_obj_folder_path):
- if not csvfile.endswith(".csv"):
- continue
- csvfile_path = os.path.join(chat_obj_folder_path, csvfile)
- chat_df = handle_pt_csv(csvfile_path)
- csv_res.append(chat_df)
-
- csv_res = pd.concat(csv_res)
- csv_res = csv_res.apply(lambda x: {"c": x}) # 设置数据集prompt键为c
-
- csv_res.to_json("./data/res_csv/pt-my.json", orient="records", force_ascii=False)
-
-
-def handle_sft_csv(csvfile):
- chat_df = pd.read_csv(csvfile)
- blocked_words = json.load(
- open("./make_dataset/blocked_words.json", encoding="utf-8")
- )["blocked_words"]
- # 选择type_name为文本的行、is_sender为1的行
- # 需要保留的type_name字段名
- type_list = [
- "文本",
- "图片",
- "视频",
- "合并转发的聊天记录",
- "语音",
- "(分享)音乐",
- "(分享)卡片式链接",
- "(分享)笔记",
- "(分享)小程序",
- "(分享)收藏夹",
- "(分享)小说(猜)",
- "(分享)视频号名片",
- "(分享)视频号视频",
- "粘贴的文本", # 无法解析的分享链接
- ]
- chat_df = chat_df[chat_df["type_name"].isin(values=type_list)]
-
- # chat_df['content'] = chat_df['content'].apply(func=lambda x: json.loads(x)['msg'])
- chat_df["content"] = chat_df["msg"]
-
- # 如果type_name为文本 并且content 包含 手机号、身份证号、邮箱、网址则删除这行
- for i in chat_df.index:
- if chat_df.loc[i, "type_name"] == "文本":
- if (
- re.search(r"1\d{10}", chat_df.loc[i, "content"])
- or re.search(r"\d{18}", chat_df.loc[i, "content"])
- or re.search(r"\w+@\w+", chat_df.loc[i, "content"])
- or "http" in chat_df.loc[i, "content"]
- or r"\\xa0" in chat_df.loc[i, "content"]
- or r"\\u" in chat_df.loc[i, "content"]
- ):
- chat_df = chat_df.drop(index=i)
- continue
- for blocked_word in blocked_words:
- if blocked_word in chat_df.loc[i, "content"]:
- chat_df = chat_df.drop(index=i)
- break
- else:
- chat_df.loc[i, "content"] = ""
-
- chat_df = chat_df[["is_sender", "type_name", "content", "CreateTime"]]
- chat_df = chat_df.dropna()
- # 时间格式 2021-07-07 10:27:23
- # 遍历行 相同is_sender的行合并content（）遇到不同is_sender就重新开始
- # CreateTime字段保留最后的CreateTime
- chat_df["CreateTime"] = pd.to_datetime(chat_df["CreateTime"])
-
- # 改到这了
-
- type_list.remove("文本")
- skip_list = type_list
- res_df = []
- last_is_sender = chat_df.iloc[0]["is_sender"]
- last_content: str = chat_df.iloc[0]["content"]
- last_CreateTime = chat_df.iloc[0]["CreateTime"]
- # 超时处理 半天没说话就重新开始
- # 注意这里只是处理了组装成一个句子 最后封装对话、配对在make_sft_dataset
- # 遇到图片 连接 直接封装成一个句子
- for i, row in chat_df.iterrows():
- if row["type_name"] in skip_list:
- if last_content != "":
- if last_content[-1] == "，":
- last_content = last_content[:-1]
- elif last_content[-1] not in ["。", "！", "？", "…", "."]:
- last_content += ""
- res_df.append(
- {
- "is_sender": last_is_sender,
- "content": last_content,
- "CreateTime": last_CreateTime,
- }
- )
- last_CreateTime = row["CreateTime"]
- last_content = ""
- # cut表示被skip字段截断
- res_df.append(
- {
- "is_sender": row["is_sender"],
- "content": "cut",
- "CreateTime": row["CreateTime"],
- }
- )
- continue
- if last_content == "": # 重新开始
- last_content = row["content"]
- last_is_sender = row["is_sender"]
- last_CreateTime = row["CreateTime"]
- continue
- if row["is_sender"] == last_is_sender:
- if row["CreateTime"] - last_CreateTime > pd.Timedelta(value="2m"):
- # 如果超时 前面的添加到res_df 并重新开始
- if last_content[-1] == "，":
- last_content = last_content[:-1]
- elif last_content[-1] not in ["。", "！", "？", "…", "."]:
- last_content += ""
- res_df.append(
- {
- "is_sender": last_is_sender,
- "content": last_content,
- "CreateTime": last_CreateTime,
- }
- )
- last_content = row["content"]
- last_CreateTime = row["CreateTime"]
- continue
- # 如果content的结尾没有标点符号则添加逗号，最后结尾是句号
- if last_content[-1] not in ["。", "！", "？", "…", "，"]:
- last_content += "，"
- last_content = last_content + row["content"]
- last_CreateTime = row["CreateTime"]
- else:
- if last_content[-1] == "，":
- last_content = last_content[:-1]
- elif last_content[-1] not in ["。", "！", "？", "…", "."]:
- last_content += ""
- res_df.append(
- {
- "is_sender": last_is_sender,
- "content": last_content,
- "CreateTime": last_CreateTime,
- }
- )
- last_is_sender = row["is_sender"]
- last_content = row["content"]
- last_CreateTime = row["CreateTime"]
- res_df = pd.DataFrame(res_df)
- return res_df
-
-
-def make_sft_dataset():
- processor = DataProcessor()
- csv_files = processor.get_csv_files()
-
- csv_concat = []
- csv_res = []
-
- for csvfile_path in csv_files:
- chat_df = handle_sft_csv(csvfile_path)
- csv_concat.append(chat_df)
-
- # 后续代码保持不变
- csv_concat = pd.concat(csv_concat)
-
- # 更全面地处理cut标记
- # 1. 将连续的cut标记合并为一个
- # 2. 标记数据区块的开始和结束
- processed_rows = []
- skip_row = False
- last_row_was_cut = False
-
- for i in range(len(csv_concat)):
- if skip_row:
- skip_row = False
- continue
-
- current_row = csv_concat.iloc[i].copy()
-
- # 处理当前行是cut的情况
- if current_row["content"] == "cut":
- # 如果上一行已经是cut，则跳过当前行
- if last_row_was_cut:
- continue
-
- # 查找连续的cut
- j = i + 1
- while j < len(csv_concat) and csv_concat.iloc[j]["content"] == "cut":
- j += 1
-
- # 如果有连续的cut，只保留最后一个
- if j > i + 1:
- current_row = csv_concat.iloc[j - 1].copy()
- skip_row = True
-
- last_row_was_cut = True
- else:
- last_row_was_cut = False
-
- processed_rows.append(current_row)
-
- # 创建新的DataFrame
- csv_concat = pd.DataFrame(processed_rows)
-
- # csv_res里is_sender必须是01 01 01 的顺序 csv_concat里不一定是01 01
- # 相差超过1小时的时间戳分为不同的对话
- # temp_res为一个长度为2的队列
- # 将合并后的数据保存到CSV文件中
- output_dir = "./test_output"
-
- # 生成带时间戳的文件名
- import datetime
-
- now = datetime.datetime.now()
- output_file = os.path.join(output_dir, f"csv_old_.csv")
-
- # 保存合并后的数据
- # csv_concat.to_csv(output_file, index=False, encoding="utf-8-sig")
- # print(f"已将合并后的数据保存到: {output_file}")
- # print(f"合并后数据总量: {len(csv_concat)} 条记录")
-
- temp_res = deque(maxlen=2)
- # 6种情况
- # temp_res 为空 遇到 0入队 遇到1不处理 遇到cut不处理
- # temp_res 有0 遇到0清空队列再入队 遇到1相差超过1小时清空队列 没有相差一小时入队再全部出队 遇到cut清空队列
-
- for i, row in csv_concat.iterrows():
- if len(temp_res) == 0:
- if row["content"] == "cut":
- continue
- if row["is_sender"] == 0:
- temp_res.append(row["content"])
- last_CreateTime = row["CreateTime"]
- else:
- continue
- elif len(temp_res) == 1:
- if row["content"] == "cut":
- temp_res.clear()
- last_CreateTime = row["CreateTime"]
- elif row["is_sender"] == 0:
- # 遇到0 清空队列再入队
- temp_res.clear()
- temp_res.append(row["content"])
- last_CreateTime = row["CreateTime"]
- else:
- if row["CreateTime"] - last_CreateTime > pd.Timedelta("5m"):
- # 相差超过1小时清空队列
- temp_res.clear()
- last_CreateTime = row["CreateTime"]
- else:
- # 没有相差一小时入队再全部出队
- temp_res.append(row["content"])
- csv_res.append({"instruction": temp_res[0], "output": temp_res[1]})
- temp_res.clear()
- last_CreateTime = row["CreateTime"]
-
- csv_res_df = pd.DataFrame(csv_res)
- print(f"处理后数据量：{csv_res_df.shape[0]}")
- csv_res_df.to_json('./data/res_csv/sft/sft-old-my.json', orient='records', force_ascii=False)
-
-
-if __name__ == "__main__":
- # make_pt_dataset()
- make_sft_dataset()
diff --git a/tests/test_qa_generator.py b/tests/test_qa_generator.py
deleted file mode 100644
index 3bc50d2..0000000
--- a/tests/test_qa_generator.py
+++ /dev/null
@@ -1,353 +0,0 @@
-import sys
-import os
-import pytest
-from datetime import datetime, timedelta
-import pandas as pd
-
-# 添加项目根目录到sys.path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-root_dir = os.path.dirname(current_dir)
-sys.path.append(root_dir)
-
-from make_dataset.models import ChatMessage, CutMessage
-from make_dataset.qa_generator import DataProcessor
-
-# 将当前工作目录更改为项目根目录
-os.chdir(root_dir)
-
-# # 测试数据处理器类的初始化和配置加载
-# def test_data_processor_init():
-# """测试DataProcessor初始化"""
-# processor = DataProcessor()
-# assert processor.csv_folder == "./data/csv"
-# assert "文本" not in processor.skip_type_list
-# assert len(processor.type_list) == 8
-
-
-class MockDataProcessor(DataProcessor):
- def __init__(self):
- super().__init__()
-
-
-@pytest.fixture
-def processor():
- """创建一个测试用的处理器实例"""
- return MockDataProcessor()
-
-
-def test_empty_messages(processor):
- """测试空消息列表的情况"""
- messages = []
- result = processor.group_consecutive_messages(messages)
- assert result == []
-
-
-def test_single_message(processor):
- """测试单条消息的情况"""
- now = datetime.now()
- message = ChatMessage(
- id=1,
- MsgSvrID=1001,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="你好",
- src="",
- CreateTime=now,
- )
-
- result = processor.group_consecutive_messages([message])
- assert len([msg for msg in result if isinstance(msg, ChatMessage)]) == 1
- assert len([msg for msg in result if isinstance(msg, CutMessage)]) == 0
- assert result[0].msg == "你好"
-
-
-def test_consecutive_messages_same_sender(processor):
- """测试同一发送者的连续消息"""
- now = datetime.now()
- messages = [
- ChatMessage(
- id=1,
- MsgSvrID=1001,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="你好",
- src="",
- CreateTime=now,
- ),
- ChatMessage(
- id=2,
- MsgSvrID=1002,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="最近怎么样",
- src="",
- CreateTime=now + timedelta(minutes=5),
- ),
- ChatMessage(
- id=3,
- MsgSvrID=1003,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="我想问个问题",
- src="",
- CreateTime=now + timedelta(minutes=10),
- ),
- ]
-
- result = processor.group_consecutive_messages(messages)
- assert len([msg for msg in result if isinstance(msg, ChatMessage)]) == 1
- assert len([msg for msg in result if isinstance(msg, CutMessage)]) == 0
- assert result[0].msg == "你好，最近怎么样，我想问个问题"
-
-
-def test_messages_different_senders(processor):
- """测试不同发送者的消息"""
- now = datetime.now()
- messages = [
- ChatMessage(
- id=1,
- MsgSvrID=1001,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="你好",
- src="",
- CreateTime=now,
- ),
- ChatMessage(
- id=2,
- MsgSvrID=1002,
- type_name="文本",
- is_sender=1,
- talker="user2",
- room_name="testroom",
- msg="你好，有什么可以帮你的",
- src="",
- CreateTime=now + timedelta(minutes=5),
- ),
- ChatMessage(
- id=3,
- MsgSvrID=1003,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="我想问个问题",
- src="",
- CreateTime=now + timedelta(minutes=10),
- ),
- ]
-
- result = processor.group_consecutive_messages(messages)
- assert len([msg for msg in result if isinstance(msg, ChatMessage)]) == 3
- assert len([msg for msg in result if isinstance(msg, CutMessage)]) == 0
- assert result[0].msg == "你好"
- assert result[1].msg == "你好，有什么可以帮你的"
- assert result[2].msg == "我想问个问题"
-
-
-def test_skip_non_text_messages(processor):
- """测试跳过非文本消息"""
- now = datetime.now()
- messages = [
- ChatMessage(
- id=1,
- MsgSvrID=1001,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="你好",
- src="",
- CreateTime=now,
- ),
- ChatMessage(
- id=2,
- MsgSvrID=1002,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="先生",
- src="image.jpg",
- CreateTime=now + timedelta(minutes=9.9),
- ),
- ChatMessage(
- id=2,
- MsgSvrID=1002,
- type_name="图片",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="",
- src="image.jpg",
- CreateTime=now + timedelta(minutes=1+9.9),
- ),
- ChatMessage(
- id=3,
- MsgSvrID=1003,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="看到图片了吗",
- src="",
- CreateTime=now + timedelta(minutes=1+9.9),
- ),
- ]
-
- result = processor.group_consecutive_messages(messages)
- chat_messages = [msg for msg in result if isinstance(msg, ChatMessage)]
- cut_messages = [msg for msg in result if isinstance(msg, CutMessage)]
- assert len(chat_messages) == 2
- assert len(cut_messages) == 1
- assert chat_messages[0].msg == "你好，先生"
-
-
-def test_time_window_limit(processor):
- """测试时间窗口限制（超过1小时的消息不会合并）"""
- now = datetime.now()
- messages = [
- ChatMessage(
- id=1,
- MsgSvrID=1001,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="你好",
- src="",
- CreateTime=now,
- ),
- ChatMessage(
- id=2,
- MsgSvrID=1002,
- type_name="文本",
- is_sender=0,
- talker="user1",
- room_name="testroom",
- msg="晚上好",
- src="",
- CreateTime=now + timedelta(hours=2), # 超过1小时
- ),
- ]
-
- result = processor.group_consecutive_messages(messages)
- assert len([msg for msg in result if isinstance(msg, ChatMessage)]) == 2
- assert len([msg for msg in result if isinstance(msg, CutMessage)]) == 0
- assert result[0].msg == "你好"
- assert result[1].msg == "晚上好"
-
-
-def test_consecutive_messages_to_csv():
- """
- 测试使用DataProcessor的main函数从CSV文件中读取数据，
- 应用group_consecutive_messages函数，并将结果保存为CSV
- """
- processor = MockDataProcessor()
-
- # 获取CSV文件列表
- csv_files = processor.get_csv_files()
-
- # 如果没有找到CSV文件，创建一个模拟的CSV文件供测试使用
- if not csv_files:
- print("警告：未找到CSV文件，请确保数据目录中有CSV文件")
- return "无法找到CSV文件"
-
- # 存储所有处理后的消息
- all_grouped_messages = []
-
- # 处理每个CSV文件
- for csv_file in csv_files:
- print(f"处理文件: {csv_file}")
- # 加载CSV文件中的消息
- chat_messages = processor.load_csv(csv_file)
- print(f"加载了 {len(chat_messages)} 条消息")
-
- # 应用group_consecutive_messages函数
- grouped_messages = processor.group_consecutive_messages(messages=chat_messages)
- print(f"分组后得到 {len(grouped_messages)} 条消息")
-
- # 添加到结果列表
- all_grouped_messages.extend(grouped_messages)
-
- # 如果没有处理到任何消息，提前返回
- if not all_grouped_messages:
- print("警告：未处理到任何消息")
- return "未处理到任何消息"
-
- # 将结果转换为DataFrame
- messages_dict = []
- for msg in all_grouped_messages:
- if isinstance(msg, ChatMessage):
- messages_dict.append(
- {
- "id": msg.id,
- "MsgSvrID": msg.MsgSvrID,
- "type_name": msg.type_name,
- "is_sender": msg.is_sender,
- "talker": msg.talker,
- "room_name": msg.room_name,
- "msg": msg.msg,
- "src": msg.src,
- "CreateTime": msg.CreateTime,
- }
- )
- elif hasattr(msg, "cut_type"): # 处理CutMessage对象
- messages_dict.append(
- {
- "id": None,
- "MsgSvrID": None,
- "type_name": msg.cut_type,
- "is_sender": msg.is_sender,
- "talker": None,
- "room_name": None,
- "msg": f"cut",
- "src": None,
- "CreateTime": msg.CreateTime,
- }
- )
-
- # 创建DataFrame
- df = pd.DataFrame(messages_dict)
-
- # 确保输出目录存在
- output_dir = "./test_output"
- os.makedirs(output_dir, exist_ok=True)
-
- # 保存为CSV文件
- import datetime
-
- now = datetime.datetime.now()
- output_file = os.path.join(output_dir, f"grouped_messages_.csv")
- # 使用utf-8-sig编码保存，添加BOM标记以解决中文乱码问题
- df.to_csv(output_file, index=False, encoding="utf-8-sig")
-
- # 验证结果
- assert os.path.exists(output_file)
- print(f"已成功保存分组消息到: {output_file}")
- print(f"共保存了 {len(messages_dict)} 条消息")
-
- # 显示前5条消息示例
- if len(messages_dict) > 0:
- print("\n消息示例:")
- for i, msg in enumerate(messages_dict[:5]):
- print(
- f"{i+1}. {'用户' if msg['is_sender'] == 0 else '对方'}: {msg['msg'][:50]}..."
- )
-
- return output_file
-
-
-if __name__ == "__main__":
- output_file = test_consecutive_messages_to_csv()
- print(f"测试完成，消息已保存到 {output_file}")
diff --git a/tests/test_weclone_pipeline_mock.py b/tests/test_weclone_pipeline_mock.py
deleted file mode 100644
index 62c7761..0000000
--- a/tests/test_weclone_pipeline_mock.py
+++ /dev/null
@@ -1,306 +0,0 @@
-import os
-import sys
-import json
-import shutil
-import unittest
-import tempfile
-from unittest.mock import patch, MagicMock
-import pandas as pd
-
-# 添加项目根目录到系统路径
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-# 导入需要测试的模块
-from weclone.data.qa_generator import DataProcessor
-from weclone.utils.config import load_config
-
-
-class TestWeclonePipeline(unittest.TestCase):
- @classmethod
- def setUpClass(cls):
- """设置测试环境"""
- # 创建临时目录用于测试
- cls.test_dir = tempfile.mkdtemp()
- cls.test_data_dir = os.path.join(cls.test_dir, "data")
- cls.test_model_dir = os.path.join(cls.test_dir, "model_output")
- cls.test_eval_dir = os.path.join(cls.test_dir, "eval_output")
-
- # 创建必要的目录
- os.makedirs(cls.test_data_dir, exist_ok=True)
- os.makedirs(cls.test_model_dir, exist_ok=True)
- os.makedirs(cls.test_eval_dir, exist_ok=True)
-
- # 创建测试数据集结构
- cls.csv_folder = os.path.join(cls.test_data_dir, "csv")
- os.makedirs(cls.csv_folder, exist_ok=True)
-
- # 创建示例聊天文件夹和CSV文件
- chat_folder = os.path.join(cls.csv_folder, "test_chat")
- os.makedirs(chat_folder, exist_ok=True)
-
- # 创建简单的测试CSV数据
- cls._create_test_csv(os.path.join(chat_folder, "test_chat.csv"))
-
- # 创建测试用的settings.jsonc
- cls._create_test_settings()
-
- # 创建测试用的test_data.json用于模型评估
- cls._create_test_eval_data()
-
- @classmethod
- def tearDownClass(cls):
- """清理测试环境"""
- # 删除临时目录
- shutil.rmtree(cls.test_dir, ignore_errors=True)
-
- @classmethod
- def _create_test_csv(cls, file_path):
- """创建测试用CSV文件"""
- import pandas as pd
-
- # 创建简单的聊天记录数据
- data = {
- "id": list(range(1, 5)),
- "MsgSvrID": list(range(1001, 1005)),
- "type": ["1", "1", "1", "1"], # 文本类型
- "is_sender": [0, 1, 0, 1], # 0=对方发送，1=自己发送
- "talker": ["test_user", "me", "test_user", "me"],
- "room_name": ["", "", "", ""],
- "content": ["你好，请问你是谁？", "我是你的微信助手", "你能帮我做什么？", "我可以回答问题，提供信息和帮助你完成各种任务"],
- "src": ["", "", "", ""],
- "CreateTime": [1609459200, 1609459220, 1609459240, 1609459260] # 时间戳
- }
-
- # 创建DataFrame并保存为CSV
- df = pd.DataFrame(data)
- df.to_csv(file_path, index=False)
-
- @classmethod
- def _create_test_settings(cls):
- """创建测试用的settings.jsonc"""
- # 简化版的设置文件，只包含测试所需的最小配置
- settings = {
- "train_sft_args": {
- "stage": "sft",
- "dataset": "wechat-sft",
- "dataset_dir": cls.test_data_dir + "/res_csv/sft",
- "lora_target": "query_key_value",
- "lora_rank": 4,
- "lora_dropout": 0.5,
- "overwrite_cache": True,
- "per_device_train_batch_size": 1,
- "gradient_accumulation_steps": 1,
- "lr_scheduler_type": "cosine",
- "logging_steps": 1,
- "save_steps": 1,
- "learning_rate": 0.0001,
- "num_train_epochs": 1,
- "plot_loss": False,
- "fp16": False
- },
- "infer_args": {
- "repetition_penalty": 1.2,
- "temperature": 0.5,
- "max_length": 50,
- "top_p": 0.65
- },
- "make_dataset_args": {
- "single_combine_strategy": "time_window",
- "qa_match_strategy": "time_window",
- "single_combine_time_window": 2,
- "qa_match_time_window": 5,
- "prompt_with_history": False
- },
- "common_args": {
- "model_name_or_path": "./chatglm3-6b", # 假设已有模型
- "adapter_name_or_path": cls.test_model_dir,
- "template": "chatglm3-weclone",
- "finetuning_type": "lora",
- "trust_remote_code": True
- }
- }
-
- # 保存到临时目录
- with open(os.path.join(cls.test_dir, "settings.jsonc"), "w", encoding="utf-8") as f:
- json.dump(settings, f, indent=4)
-
- @classmethod
- def _create_test_eval_data(cls):
- """创建测试用的评估数据"""
- test_data = {
- "questions": [
- ["你好", "你是谁"],
- ["你能做什么"]
- ]
- }
-
- # 确保目录存在
- data_dir = os.path.join(cls.test_dir, "data")
- os.makedirs(data_dir, exist_ok=True)
-
- # 保存测试数据
- with open(os.path.join(data_dir, "test_data.json"), "w", encoding="utf-8") as f:
- json.dump(test_data, f, ensure_ascii=False, indent=4)
-
- @patch('weclone.data.qa_generator.DataProcessor.get_csv_files')
- @patch('weclone.data.qa_generator.DataProcessor.load_csv')
- @patch('weclone.data.qa_generator.DataProcessor.save_result')
- def test_qa_generator(self, mock_save_result, mock_load_csv, mock_get_csv_files):
- """测试QA生成器"""
- print("\n测试QA生成器...")
-
- # 准备模拟数据
- from weclone.data.models import ChatMessage
- mock_get_csv_files.return_value = ["test_csv_file.csv"]
-
- # 模拟从CSV加载的消息
- mock_messages = [
- ChatMessage(id=1, MsgSvrID=1001, type_name="文本", is_sender=0,
- talker="test_user", room_name="", msg="你好，请问你是谁？",
- src="", CreateTime=pd.Timestamp(1609459200, unit='s')),
- ChatMessage(id=2, MsgSvrID=1002, type_name="文本", is_sender=1,
- talker="me", room_name="", msg="我是你的微信助手",
- src="", CreateTime=pd.Timestamp(1609459220, unit='s')),
- ChatMessage(id=3, MsgSvrID=1003, type_name="文本", is_sender=0,
- talker="test_user", room_name="", msg="你能帮我做什么？",
- src="", CreateTime=pd.Timestamp(1609459240, unit='s')),
- ChatMessage(id=4, MsgSvrID=1004, type_name="文本", is_sender=1,
- talker="me", room_name="", msg="我可以回答问题，提供信息和帮助你完成各种任务",
- src="", CreateTime=pd.Timestamp(1609459260, unit='s'))
- ]
- mock_load_csv.return_value = mock_messages
-
- # 创建DataProcessor实例
- with patch('weclone.utils.config.load_config') as mock_load_config:
- # 模拟配置
- mock_config = {
- "single_combine_strategy": "time_window",
- "qa_match_strategy": "time_window",
- "single_combine_time_window": 2,
- "qa_match_time_window": 5,
- "prompt_with_history": False
- }
- mock_load_config.return_value = mock_config
-
- # 执行QA生成
- processor = DataProcessor()
- processor.csv_folder = self.csv_folder # 设置为测试目录
- processor.main()
-
- # 验证是否调用了预期的方法
- mock_get_csv_files.assert_called_once()
- mock_load_csv.assert_called_once()
- mock_save_result.assert_called_once()
-
- # 验证结果格式
- # 获取保存的结果
- call_args = mock_save_result.call_args[0][0]
- self.assertTrue(isinstance(call_args, list))
- self.assertEqual(len(call_args), 2) # 应该有两个QA对
-
- # 验证QA对的结构
- for qa in call_args:
- self.assertTrue("instruction" in qa)
- self.assertTrue("output" in qa)
-
- print("QA生成器测试成功")
-
- def test_train_sft(self):
- """测试SFT训练过程"""
- print("\n测试SFT训练过程...")
- # 由于训练需要实际的模型和数据，这里我们只模拟调用
-
- with patch('llamafactory.train.tuner.run_exp') as mock_run_exp:
- # 导入训练模块并运行
- from weclone.train.train_sft import run_exp
-
-
- # 验证是否正确调用了训练函数
- self.assertTrue(mock_run_exp.called)
- print("SFT训练过程测试成功")
-
- def test_api_service(self):
- """测试API服务"""
- print("\n测试API服务...")
-
- # 模拟服务器进程
- with patch('uvicorn.run') as mock_run:
- # 导入API服务模块
- from weclone.server.api_service import main, create_app, ChatModel
-
- # 模拟配置和模型
- with patch('weclone.utils.config.load_config') as mock_load_config:
- mock_config = {"model_path": "test_model_path"}
- mock_load_config.return_value = mock_config
-
- # 模拟ChatModel
- with patch('llamafactory.chat.ChatModel') as MockChatModel:
- mock_chat_model = MagicMock()
- MockChatModel.return_value = mock_chat_model
-
- # 运行API服务
- main()
-
- # 验证服务是否正确启动
- mock_run.assert_called_once()
- call_args = mock_run.call_args[1]
- self.assertEqual(call_args["host"], "0.0.0.0")
- self.assertEqual(call_args["port"], 8005) # 默认端口
- self.assertEqual(call_args["workers"], 1)
-
- print("API服务测试成功")
-
- def test_model_evaluation(self):
- """测试模型评估"""
- print("\n测试模型评估...")
-
- # 模拟OpenAI API调用
- with patch('openai.ChatCompletion.create') as mock_create:
- # 设置模拟返回值
- mock_response = MagicMock()
- mock_response.choices = [MagicMock()]
- mock_response.choices[0].message.content = "这是模型的测试回复"
- mock_create.return_value = mock_response
-
- # 运行评估脚本
- with patch('builtins.open', create=True) as mock_open:
- # 模拟打开测试数据文件
- test_data_content = '{"questions": [["你好", "你是谁"], ["你能做什么"]]}'
- mock_file = MagicMock()
- mock_file.read.return_value = test_data_content
- mock_open.return_value.__enter__.return_value = mock_file
-
- # 导入并运行评估模块
- from weclone.eval.test_model import main
-
- # 执行评估
- main()
-
- # 验证API调用次数（应该是测试问题的数量）
- self.assertEqual(mock_create.call_count, 3) # 3个测试问题
-
- print("模型评估测试成功")
-
- def test_full_pipeline(self):
- """测试完整流程"""
- print("\n测试完整流程...")
-
- # 这个测试方法会依次调用上面的各个测试方法，模拟完整的流程
-
- # 1. 测试QA生成器
- self.test_qa_generator()
-
- # 2. 测试SFT训练
- self.test_train_sft()
-
- # 3. 测试API服务
- self.test_api_service()
-
- # 4. 测试模型评估
- self.test_model_evaluation()
-
- print("完整流程测试完成")
-
-
-if __name__ == "__main__":
- unittest.main()
\ No newline at end of file
diff --git a/weclone/data/qa_generator.py b/weclone/data/qa_generator.py
index 98a5673..0233c79 100644
--- a/weclone/data/qa_generator.py
+++ b/weclone/data/qa_generator.py
@@ -74,7 +74,7 @@ class DataProcessor:
if self.config.get("prompt_with_history", False):
logger.warning("开启 prompt_with_history 不支持 clean_dataset 功能")
exit()
-
+
if not is_vllm_available():
logger.warning("vLLM 不可用，暂不清洗数据集。")
clean_dataset_config["enable_clean"] = False
@@ -164,7 +164,8 @@ class DataProcessor:
csvfile_path = os.path.join(chat_obj_folder_path, csvfile)
csv_files.append(csvfile_path)
# 提取文件名中的起始数字，比如 wxid_..._0_5000.csv → 0
- pattern = re.compile(r'_(\d+)_\d+\.csv$')
+ pattern = re.compile(r"_(\d+)_\d+\.csv$")
+
def extract_start(fp: str) -> int:
name = os.path.basename(fp)
m = pattern.search(name)
@@ -339,7 +340,9 @@ class DataProcessor:
combined_content += content
if len(combined_content) > self.c["combine_msg_max_length"]:
- logger.warning(f"组合后消息长度超过{self.c['combine_msg_max_length']}将截断：\n {combined_content[: 50]}")
+ logger.warning(
+ f"组合后消息长度超过{self.c['combine_msg_max_length']}将截断：\n {combined_content[:50]}"
+ )
combined_content = combined_content[: self.c["combine_msg_max_length"]]
combined_message = ChatMessage(