opensource/WeClone-new
3
0
Fork 0
mirror of https://github.com/xming521/WeClone.git synced 2026-08-28 18:07:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add full-flow test pipeline

Browse Source
This commit is contained in:
xming521
2025-05-22 20:05:22 +08:00
parent bb0e4c4ece
commit bbe16fda73
12 changed files with 169 additions and 2316 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+9 -2
View File
@@ -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"
tests/README.md
-39
View File
@@ -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
```
tests/__init__.py
View File
tests/data/clean/test_strategies.py
-176
View File
@@ -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调试脚本运行结束。")
tests/test_full_pipe.py
+154
View File
@@ -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="<yellow><b>{message}</b></yellow>",
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("服务器已关闭")
tests/test_full_pipeline.py
-902
View File
@@ -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] <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level.name[0]}</level> | <level>{message}</level>", 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)
tests/test_get_sample_audio.py
-151
View File
@@ -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()
tests/test_get_score.py
-72
View File
@@ -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])
tests/test_old_csv_to_json.py
-312
View File
@@ -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()
tests/test_qa_generator.py
-353
View File
@@ -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}")
tests/test_weclone_pipeline_mock.py
-306
View File
@@ -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()
weclone/data/qa_generator.py
+6 -3
View File
@@ -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(