mirror of
https://github.com/xming521/WeClone.git
synced 2026-08-28 18:07:28 +08:00
feat(data): 支持图像识别并优化数据清洗流程
- 新增 VisionApiConfig 类用于配置视觉 API - 在数据处理中集成图像识别功能，支持并行处理 - 重构数据清洗策略，支持在线和离线两种方式- 优化数据清洗流程，提高可扩展性和可维护性
This commit is contained in:
BAIKEMARK
@@ -13,5 +13,20 @@
|
||||
"user_tag": "user",
|
||||
"assistant_tag": "assistant"
|
||||
}
|
||||
},
|
||||
"chat-sft-cleaned": {
|
||||
"file_name": "./sft-my-cleaned.json",
|
||||
"formatting": "sharegpt",
|
||||
"columns": {
|
||||
"messages": "messages",
|
||||
"system": "system",
|
||||
"images": "images"
|
||||
},
|
||||
"tags": {
|
||||
"role_tag": "role",
|
||||
"content_tag": "content",
|
||||
"user_tag": "user",
|
||||
"assistant_tag": "assistant"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,63 +1,121 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from tqdm import tqdm
|
||||
|
||||
from weclone.data.models import QaPairScore, QaPairV2
|
||||
from weclone.prompts.clean_data import CLEAN_PROMPT
|
||||
from weclone.core.inference.online_infer import OnlineLLM
|
||||
from weclone.data.models import QaPair, QaPairScore, QaPairV2
|
||||
from weclone.prompts.clean_data import CLEAN_PROMPT, ONLINE_LLM_CLEAN_PROMPT
|
||||
from weclone.utils.config_models import WCMakeDatasetConfig
|
||||
from weclone.utils.log import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class CleaningStrategy(ABC):
|
||||
"""数据清洗策略的抽象基类"""
|
||||
"""数据清洗策略的抽象基类，但提供通用的清洗方法"""
|
||||
|
||||
make_dataset_config: Dict
|
||||
make_dataset_config: WCMakeDatasetConfig
|
||||
|
||||
@abstractmethod
|
||||
def clean(self, data: Any) -> Any:
|
||||
def judge(self, data: List[QaPair] | List[QaPairV2]) -> None:
|
||||
"""
|
||||
执行数据清洗操作。
|
||||
|
||||
Args:
|
||||
data: 需要清洗的数据。
|
||||
|
||||
Returns:
|
||||
清洗后的数据。
|
||||
打分方法是抽象的，强制每个子类根据自己的方式去实现。
|
||||
"""
|
||||
pass
|
||||
|
||||
def clean(self) -> str:
|
||||
"""
|
||||
通用策略
|
||||
根据score筛选SFT数据，并返回最终应使用的dataset名称。
|
||||
"""
|
||||
config = self.make_dataset_config
|
||||
original_dataset_name = config.dataset
|
||||
cleaned_dataset_name = "chat-sft-cleaned"
|
||||
|
||||
if not config.clean_dataset.enable_clean or "image" in config.include_type:
|
||||
logger.info("数据清洗未启用或包含图像，将使用原始数据集。")
|
||||
return original_dataset_name
|
||||
|
||||
dataset_dir = config.dataset_dir
|
||||
dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
|
||||
|
||||
# 获取文件名称
|
||||
try:
|
||||
with open(dataset_info_path, "r", encoding="utf-8") as f:
|
||||
info = json.load(f)
|
||||
paths = {
|
||||
name: os.path.join(dataset_dir, info.get(name, {}).get("file_name"))
|
||||
for name in [original_dataset_name, cleaned_dataset_name]
|
||||
}
|
||||
original_data_path, cleaned_data_path = paths.values()
|
||||
if not all(paths.values()):
|
||||
raise ValueError(f"缺失 '{original_dataset_name}' 或 '{cleaned_dataset_name}' 文件配置。")
|
||||
except Exception as e:
|
||||
logger.error(f"加载 dataset_info.json 出错: {e}，将使用原始数据集。")
|
||||
return original_dataset_name
|
||||
|
||||
# 执行清洗流程
|
||||
logger.info(f"数据清洗已启用，将从 '{original_data_path}' 读取数据...")
|
||||
try:
|
||||
if not os.path.exists(original_data_path):
|
||||
logger.error(f"原始数据文件 '{original_data_path}' 不存在，清洗中止。")
|
||||
return original_dataset_name
|
||||
with open(original_data_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
accept_score = config.clean_dataset.llm.accept_score
|
||||
filtered_data = [item for item in data if item.get("score", 0) >= accept_score]
|
||||
|
||||
if not filtered_data:
|
||||
logger.warning("清洗后无数据保留，将使用原始数据集。")
|
||||
return original_dataset_name
|
||||
|
||||
with open(cleaned_data_path, "w", encoding="utf-8") as f:
|
||||
json.dump(filtered_data, f, ensure_ascii=False, indent=2)
|
||||
logger.success(
|
||||
f"已筛出低于 {accept_score} 分的数据，保留 {len(filtered_data)} 条，保存至 {cleaned_data_path}"
|
||||
)
|
||||
return cleaned_dataset_name
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"数据清洗过程中发生错误，将使用原始数据集: {e}")
|
||||
return original_dataset_name
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMCleaningStrategy(CleaningStrategy):
|
||||
"""使用大模型进行数据清洗的策略"""
|
||||
|
||||
def judge(self, data: List[QaPairV2]) -> None:
|
||||
make_dataset_config: WCMakeDatasetConfig
|
||||
|
||||
def judge(self, data: List[QaPair] | List[QaPairV2]) -> None:
|
||||
"""
|
||||
调用llm打分，并将分数直接赋值给传入的QaPair。
|
||||
"""
|
||||
from weclone.core.inference.offline_infer import vllm_infer
|
||||
|
||||
config_dict = self.make_dataset_config.model_dump()
|
||||
|
||||
logger.info("开始使用llm对数据打分")
|
||||
inputs = []
|
||||
prompt_template = PromptTemplate.from_template(CLEAN_PROMPT)
|
||||
for qa in data:
|
||||
messages_str = ""
|
||||
for msg in qa.messages:
|
||||
if msg.role == "user":
|
||||
messages_str += f"Q: {msg.content}\n"
|
||||
elif msg.role == "assistant":
|
||||
messages_str += f"A: {msg.content}\n"
|
||||
prompt_value = prompt_template.invoke({"id": qa.id, "messages": messages_str.strip()})
|
||||
inputs.append(prompt_value.to_string())
|
||||
qa_info_list = [
|
||||
{
|
||||
"id": qa.id,
|
||||
"Q": next((msg.content for msg in qa.messages if msg.role == "user"), ""),
|
||||
"A": next((msg.content for msg in qa.messages if msg.role == "assistant"), ""),
|
||||
}
|
||||
for qa in data
|
||||
]
|
||||
inputs = [prompt_template.invoke(info).text for info in qa_info_list]
|
||||
outputs = vllm_infer(
|
||||
inputs,
|
||||
self.make_dataset_config["model_name_or_path"],
|
||||
template=self.make_dataset_config["template"],
|
||||
config_dict["model_name_or_path"], # 使用转换后的字典
|
||||
template=config_dict["template"],
|
||||
temperature=0,
|
||||
guided_decoding_class=QaPairScore,
|
||||
repetition_penalty=1.2,
|
||||
@@ -97,61 +155,86 @@ class LLMCleaningStrategy(CleaningStrategy):
|
||||
printable_df_str = distribution_df.reset_index().to_string(index=False)
|
||||
logger.success(f"llm打分分数分布情况:\n{printable_df_str}")
|
||||
|
||||
def clean(self) -> str:
|
||||
"""
|
||||
清洗 SFT 数据并返回清洗后的文件路径。
|
||||
如果未启用清洗，则返回原始路径。
|
||||
"""
|
||||
|
||||
@dataclass
|
||||
class OlineLLMCleaningStrategy(CleaningStrategy):
|
||||
"""使用大模型进行数据清洗的策略"""
|
||||
|
||||
def judge(self, data: List[QaPair] | List[QaPairV2]) -> None:
|
||||
config = self.make_dataset_config
|
||||
dataset_dir = config["dataset_dir"]
|
||||
dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
|
||||
logger.info("开始使用在线模型对数据打分")
|
||||
logger.info(f"使用模型 {config.model_name}")
|
||||
|
||||
sft_json_path = os.path.join(dataset_dir, "sft-my.json")
|
||||
output_json_path = os.path.join(dataset_dir, "sft-my-l.json")
|
||||
accept_score = config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1)
|
||||
client = OnlineLLM(
|
||||
api_key=config.llm_api_key,
|
||||
base_url=config.base_url,
|
||||
model_name=config.model_name,
|
||||
default_system=config.default_system,
|
||||
)
|
||||
prompt_template = PromptTemplate.from_template(ONLINE_LLM_CLEAN_PROMPT)
|
||||
|
||||
if not config.get("clean_dataset", {}).get("enable_clean") or "image" in config.get(
|
||||
"include_type", ""
|
||||
):
|
||||
logger.info("不启用数据清洗功能")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json")
|
||||
return sft_json_path
|
||||
parsed_scores = []
|
||||
clean_batch_size = config.clean_batch_size
|
||||
|
||||
try:
|
||||
with open(sft_json_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
filtered_data = [item for item in data if item.get("score", 0) >= accept_score]
|
||||
for i in tqdm(range(0, len(data), clean_batch_size), desc="在线模型评分进度"):
|
||||
batch = data[i : i + clean_batch_size]
|
||||
# 构造当前批次的 qa_list
|
||||
# qa_list = [{"id": qa.id, "Q": qa.instruction, "A": qa.output} for qa in batch]
|
||||
qa_list = [
|
||||
{
|
||||
"id": qa.id,
|
||||
"Q": next((msg.content for msg in qa.messages if msg.role == "user"), ""),
|
||||
"A": next((msg.content for msg in qa.messages if msg.role == "assistant"), ""),
|
||||
}
|
||||
for qa in batch
|
||||
]
|
||||
qa_list_json = json.dumps(qa_list, ensure_ascii=False)
|
||||
# 填充模板
|
||||
prompt_text = prompt_template.invoke({"qa_list": qa_list_json}).text
|
||||
try:
|
||||
response = client.chat(prompt_text)
|
||||
result_text = response.choices[0].message.content
|
||||
# print("大模型返回：",result_text)
|
||||
# 如果有 <think> … </think>，只保留 </think> 之后的内容
|
||||
if "</think>" in result_text:
|
||||
result_text = result_text.split("</think>", 1)[1]
|
||||
# 去掉开头和结尾的 ```json 或 ``` 等代码块标记
|
||||
result_text = re.sub(r"^```json\s*|```$", "", result_text.strip(), flags=re.MULTILINE)
|
||||
# 如果偶尔的几次解析失败就跳过
|
||||
try:
|
||||
score_list = json.loads(result_text)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"JSON 解析失败，跳过本批次: {e}\n内容：{result_text}")
|
||||
continue
|
||||
|
||||
with open(output_json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(filtered_data, f, ensure_ascii=False, indent=4)
|
||||
for item in score_list:
|
||||
parsed_scores.append(QaPairScore(**item))
|
||||
except Exception as e:
|
||||
ids_in_batch = [qa["id"] for qa in qa_list]
|
||||
logger.error(
|
||||
f"调用在线模型或解析结果失败，当前 batch QA ID 列表: {ids_in_batch}，错误信息: {str(e)}"
|
||||
)
|
||||
|
||||
logger.success(f"已筛出低于{accept_score}分的数据，共保留 {len(filtered_data)} 条数据")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my-l.json")
|
||||
return output_json_path
|
||||
score_map = {score.id: score.score for score in parsed_scores}
|
||||
for qa in data:
|
||||
if qa.id in score_map:
|
||||
qa.score = score_map[qa.id]
|
||||
else:
|
||||
logger.warning(f"未获取到QA ID {qa.id}的分数，默认赋值0")
|
||||
qa.score = 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"清洗数据失败，使用原始数据: {str(e)}")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json")
|
||||
return sft_json_path
|
||||
|
||||
def _update_dataset_info_file(self, dataset_info_path: str, new_file_name: str):
|
||||
"""
|
||||
修改 dataset_info.json 文件中的 file_name 字段
|
||||
"""
|
||||
try:
|
||||
with open(dataset_info_path, "r", encoding="utf-8") as f:
|
||||
dataset_info = json.load(f)
|
||||
|
||||
# 更新所有支持的数据集的 file_name
|
||||
for key in ["wechat-sft", "wechat-sft-with-history"]:
|
||||
if key in dataset_info:
|
||||
dataset_info[key]["file_name"] = new_file_name
|
||||
|
||||
# 写回文件
|
||||
with open(dataset_info_path, "w", encoding="utf-8") as f:
|
||||
json.dump(dataset_info, f, indent=4, ensure_ascii=False)
|
||||
|
||||
logger.info(f"已更新 dataset_info.json 中的 file_name 为 {new_file_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"无法更新 dataset_info.json: {e}")
|
||||
# 统计分数分布，打印日志（和本地版本保持一致）
|
||||
scores = [qa.score for qa in data if qa.score is not None]
|
||||
score_series = pd.Series(scores)
|
||||
score_counts = score_series.value_counts().sort_index()
|
||||
score_percentages = score_series.value_counts(normalize=True).sort_index() * 100
|
||||
pd.set_option("display.unicode.east_asian_width", True)
|
||||
distribution_df = pd.DataFrame(
|
||||
{
|
||||
"数量": score_counts,
|
||||
"占比(%)": score_percentages.round(2),
|
||||
}
|
||||
)
|
||||
distribution_df.index.name = "分数"
|
||||
printable_df_str = distribution_df.reset_index().to_string(index=False)
|
||||
logger.success(f"在线模型打分分数分布情况:\n{printable_df_str}")
|
||||
|
||||
@@ -1,158 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pandas as pd
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from tqdm import tqdm
|
||||
|
||||
from weclone.core.inference.online_infer import OnlineLLM
|
||||
from weclone.data.models import QaPair, QaPairScore, QaPairV2
|
||||
from weclone.prompts.clean_data import ONLINE_LLM_CLEAN_PROMPT
|
||||
from weclone.utils.log import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class CleaningStrategy(ABC):
|
||||
"""数据清洗策略的抽象基类"""
|
||||
|
||||
make_dataset_config: Dict
|
||||
|
||||
@abstractmethod
|
||||
def clean(self, data: Any) -> Any:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class OlineLLMCleaningStrategy(CleaningStrategy):
|
||||
"""使用大模型进行数据清洗的策略"""
|
||||
|
||||
def judge(self, data: List[QaPair] | List[QaPairV2]) -> None:
|
||||
logger.info("开始使用在线模型对数据打分")
|
||||
|
||||
logger.info(f"使用模型 {self.make_dataset_config.get('model_name', '')}")
|
||||
|
||||
client = OnlineLLM(
|
||||
api_key=self.make_dataset_config.get("llm_api_key"),
|
||||
base_url=self.make_dataset_config.get("base_url"),
|
||||
model_name=self.make_dataset_config.get("model_name"),
|
||||
default_system=self.make_dataset_config.get("default_system"),
|
||||
)
|
||||
prompt_template = PromptTemplate.from_template(ONLINE_LLM_CLEAN_PROMPT)
|
||||
|
||||
parsed_scores = []
|
||||
clean_batch_size = int(self.make_dataset_config.get("clean_batch_size", 10))
|
||||
for i in tqdm(range(0, len(data), clean_batch_size), desc="在线模型评分进度"):
|
||||
batch = data[i : i + clean_batch_size]
|
||||
# 构造当前批次的 qa_list
|
||||
qa_list = [{"id": qa.id, "Q": qa.instruction, "A": qa.output} for qa in batch]
|
||||
qa_list_json = json.dumps(qa_list, ensure_ascii=False)
|
||||
# 填充模板
|
||||
prompt_text = prompt_template.invoke({"qa_list": qa_list_json}).text
|
||||
try:
|
||||
response = client.chat(prompt_text)
|
||||
result_text = response.choices[0].message.content
|
||||
# print("大模型返回：",result_text)
|
||||
# 如果有 <think> … </think>，只保留 </think> 之后的内容
|
||||
if "</think>" in result_text:
|
||||
result_text = result_text.split("</think>", 1)[1]
|
||||
# 去掉开头和结尾的 ```json 或 ``` 等代码块标记
|
||||
result_text = re.sub(r"^```json\s*|```$", "", result_text.strip(), flags=re.MULTILINE)
|
||||
# 如果偶尔的几次解析失败就跳过
|
||||
try:
|
||||
score_list = json.loads(result_text)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"JSON 解析失败，跳过本批次: {e}\n内容：{result_text}")
|
||||
continue
|
||||
|
||||
for item in score_list:
|
||||
parsed_scores.append(QaPairScore(**item))
|
||||
except Exception as e:
|
||||
ids_in_batch = [qa["id"] for qa in qa_list]
|
||||
logger.error(
|
||||
f"调用在线模型或解析结果失败，当前 batch QA ID 列表: {ids_in_batch}，错误信息: {str(e)}"
|
||||
)
|
||||
|
||||
score_map = {score.id: score.score for score in parsed_scores}
|
||||
for qa in data:
|
||||
if qa.id in score_map:
|
||||
qa.score = score_map[qa.id]
|
||||
else:
|
||||
logger.warning(f"未获取到QA ID {qa.id}的分数，默认赋值0")
|
||||
qa.score = 0
|
||||
|
||||
# 统计分数分布，打印日志（和本地版本保持一致）
|
||||
scores = [qa.score for qa in data if qa.score is not None]
|
||||
score_series = pd.Series(scores)
|
||||
score_counts = score_series.value_counts().sort_index()
|
||||
score_percentages = score_series.value_counts(normalize=True).sort_index() * 100
|
||||
pd.set_option("display.unicode.east_asian_width", True)
|
||||
distribution_df = pd.DataFrame(
|
||||
{
|
||||
"数量": score_counts,
|
||||
"占比(%)": score_percentages.round(2),
|
||||
}
|
||||
)
|
||||
distribution_df.index.name = "分数"
|
||||
printable_df_str = distribution_df.reset_index().to_string(index=False)
|
||||
logger.success(f"在线模型打分分数分布情况:\n{printable_df_str}")
|
||||
|
||||
def clean(self) -> str:
|
||||
"""
|
||||
清洗 SFT 数据并返回清洗后的文件路径。
|
||||
如果未启用清洗，则返回原始路径。
|
||||
"""
|
||||
config = self.make_dataset_config
|
||||
dataset_dir = config["dataset_dir"]
|
||||
dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
|
||||
|
||||
sft_json_path = os.path.join(dataset_dir, "sft-my.json")
|
||||
output_json_path = os.path.join(dataset_dir, "sft-my-l.json")
|
||||
accept_score = config.get("clean_dataset", {}).get("llm", {}).get("accept_score", 1)
|
||||
|
||||
if not config.get("clean_dataset", {}).get("enable_clean"):
|
||||
logger.info("未启用清洗功能")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json")
|
||||
return sft_json_path
|
||||
|
||||
try:
|
||||
with open(sft_json_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
filtered_data = [item for item in data if item.get("score", 0) >= accept_score]
|
||||
|
||||
with open(output_json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(filtered_data, f, ensure_ascii=False, indent=4)
|
||||
|
||||
logger.success(f"已筛出低于{accept_score}分的数据，共保留 {len(filtered_data)} 条数据")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my-l.json")
|
||||
return output_json_path
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"清洗数据失败，使用原始数据: {str(e)}")
|
||||
self._update_dataset_info_file(dataset_info_path, new_file_name="sft-my.json")
|
||||
return sft_json_path
|
||||
|
||||
def _update_dataset_info_file(self, dataset_info_path: str, new_file_name: str):
|
||||
"""
|
||||
修改 dataset_info.json 文件中的 file_name 字段
|
||||
"""
|
||||
try:
|
||||
with open(dataset_info_path, "r", encoding="utf-8") as f:
|
||||
dataset_info = json.load(f)
|
||||
|
||||
# 更新所有支持的数据集的 file_name
|
||||
for key in ["wechat-sft", "wechat-sft-with-history"]:
|
||||
if key in dataset_info:
|
||||
dataset_info[key]["file_name"] = new_file_name
|
||||
|
||||
# 写回文件
|
||||
with open(dataset_info_path, "w", encoding="utf-8") as f:
|
||||
json.dump(dataset_info, f, indent=4, ensure_ascii=False)
|
||||
|
||||
logger.info(f"已更新 dataset_info.json 中的 file_name 为 {new_file_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"无法更新 dataset_info.json: {e}")
|
||||
@@ -1,508 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import List, Union
|
||||
|
||||
import pandas as pd
|
||||
from llamafactory.extras.packages import is_vllm_available
|
||||
from pandas import Timestamp
|
||||
|
||||
from weclone.data.clean.strategies import LLMCleaningStrategy
|
||||
from weclone.data.clean.strategies_online import OlineLLMCleaningStrategy
|
||||
from weclone.data.models import ChatMessage, CutMessage, QaPair, skip_type_list
|
||||
from weclone.data.strategies import LLMStrategy, TimeWindowStrategy
|
||||
from weclone.utils.config import load_config
|
||||
from weclone.utils.log import logger
|
||||
|
||||
|
||||
class DataProcessor:
|
||||
def __init__(self):
|
||||
self.config = load_config(arg_type="make_dataset")
|
||||
self.csv_folder = "./dataset/csv"
|
||||
self.system_prompt = self.config["default_system"]
|
||||
self.cut_type_list = [
|
||||
"图片",
|
||||
"视频",
|
||||
"合并转发的聊天记录",
|
||||
"语音",
|
||||
"(分享)音乐",
|
||||
"(分享)卡片式链接",
|
||||
"(分享)笔记",
|
||||
"(分享)小程序",
|
||||
"(分享)收藏夹",
|
||||
"(分享)小说(猜)",
|
||||
"(分享)视频号名片",
|
||||
"(分享)视频号视频",
|
||||
"粘贴的文本", # 无法解析的分享链接
|
||||
]
|
||||
|
||||
# blocked_words
|
||||
config_blocked_words = self.config.get("blocked_words", [])
|
||||
file_blocked_words = []
|
||||
try:
|
||||
with open("./dataset/blocked_words.json", encoding="utf-8") as f:
|
||||
file_blocked_words = json.load(f).get("blocked_words", [])
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
pass
|
||||
|
||||
self.blocked_words = list(set(config_blocked_words + file_blocked_words))
|
||||
# logger.info(f"聊天记录禁用词: {self.blocked_words}")
|
||||
|
||||
if self.config["single_combine_strategy"] == "time_window":
|
||||
self.single_combine_strategy = TimeWindowStrategy(
|
||||
time_window=self.config["single_combine_time_window"] * 60,
|
||||
is_single_chat=True,
|
||||
)
|
||||
elif self.config["single_combine_strategy"] == "llm":
|
||||
self.single_combine_strategy = LLMStrategy(
|
||||
is_single_chat=True,
|
||||
)
|
||||
|
||||
if self.config["qa_match_strategy"] == "time_window":
|
||||
self.qa_match_strategy = TimeWindowStrategy(
|
||||
time_window=self.config["qa_match_time_window"] * 60,
|
||||
is_single_chat=False,
|
||||
)
|
||||
elif self.config["qa_match_strategy"] == "llm":
|
||||
self.qa_match_strategy = LLMStrategy(is_single_chat=False)
|
||||
|
||||
clean_dataset_config = self.config.get("clean_dataset", {})
|
||||
enable_clean = clean_dataset_config.get("enable_clean", False)
|
||||
|
||||
if enable_clean:
|
||||
if self.config.get("prompt_with_history", False):
|
||||
logger.warning("开启 prompt_with_history 不支持 clean_dataset 功能")
|
||||
exit()
|
||||
|
||||
if not is_vllm_available() and not self.config.get("online_llm_clear"):
|
||||
logger.warning("vLLM 不可用，暂不清洗数据集。")
|
||||
clean_dataset_config["enable_clean"] = False
|
||||
|
||||
if self.config.get("clean_dataset", {}).get("enable_clean", False):
|
||||
if self.config.get("clean_dataset", {}).get("clean_strategy", "llm") == "llm":
|
||||
if self.config.get("online_llm_clear"):
|
||||
self.clean_strategy = OlineLLMCleaningStrategy(make_dataset_config=self.config)
|
||||
else:
|
||||
self.clean_strategy = LLMCleaningStrategy(make_dataset_config=self.config)
|
||||
self.c = self.config
|
||||
|
||||
def main(self):
|
||||
if not os.path.exists(self.csv_folder) or not os.listdir(self.csv_folder):
|
||||
logger.error(
|
||||
f"错误：目录 '{self.csv_folder}' 不存在或为空，请检查路径并确保其中包含 CSV 聊天数据文件。"
|
||||
)
|
||||
return
|
||||
|
||||
csv_files = self.get_csv_files()
|
||||
logger.info(f"共发现 {len(csv_files)} 个 CSV 文件，开始处理")
|
||||
message_list: List[ChatMessage] = []
|
||||
for csv_file in csv_files:
|
||||
logger.debug(f"开始处理 CSV 文件: {csv_file}")
|
||||
chat_messages = self.load_csv(csv_file)
|
||||
message_list.extend(self.group_consecutive_messages(messages=chat_messages))
|
||||
# self.process_by_msgtype(chat_message)
|
||||
logger.debug(f"处理完成: {csv_file}，共加载 {len(chat_messages)} 条消息")
|
||||
qa_res = self.match_qa(message_list)
|
||||
if self.c["prompt_with_history"]:
|
||||
qa_res = self.add_history_to_qa(qa_res)
|
||||
else:
|
||||
qa_res = [item for item in qa_res if isinstance(item, QaPair)]
|
||||
|
||||
if self.c.get("clean_dataset", {}).get("enable_clean", False):
|
||||
self.clean_strategy.judge(qa_res)
|
||||
# qa_res = self.clean_strategy.clean(qa_res)
|
||||
self.save_result(qa_res)
|
||||
self._execute_length_cdf_script()
|
||||
|
||||
logger.success(f"聊天记录处理成功，共{len(qa_res)}条，保存到 ./dataset/res_csv/sft/sft-my.json")
|
||||
|
||||
def _execute_length_cdf_script(self):
|
||||
"""执行 length_cdf.py 脚本来计算cutoff_len。"""
|
||||
try:
|
||||
python_executable = sys.executable
|
||||
# 脚本路径是相对于项目根目录的
|
||||
script_path = os.path.join("weclone", "utils", "length_cdf.py")
|
||||
|
||||
command_parts = [
|
||||
python_executable,
|
||||
script_path,
|
||||
f'--model_name_or_path="{self.c["model_name_or_path"]}"',
|
||||
f'--dataset="{self.c["dataset"]}"',
|
||||
f'--dataset_dir="{self.c["dataset_dir"]}"',
|
||||
f'--template="{self.c["template"]}"',
|
||||
f"--interval={self.c['cutoff_len']}",
|
||||
]
|
||||
|
||||
child_env = os.environ.copy()
|
||||
child_env["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
child_env["LLAMAFACTORY_VERBOSITY"] = "ERROR"
|
||||
|
||||
process = subprocess.Popen(
|
||||
command_parts,
|
||||
env=child_env,
|
||||
stdout=None, # 使用 None 表示使用父进程的标准输出（即终端）
|
||||
stderr=None, # 使用 None 表示使用父进程的标准错误（即终端）
|
||||
text=True,
|
||||
bufsize=1, # 行缓冲
|
||||
)
|
||||
return_code = process.wait()
|
||||
if return_code != 0:
|
||||
logger.error(f"命令 '{' '.join(command_parts)}' 执行失败，返回码 {return_code}")
|
||||
except FileNotFoundError:
|
||||
# command_parts[0] 是 python_executable, command_parts[1] 是 script_path
|
||||
logger.error(f"命令执行失败: 找不到可执行文件 '{command_parts[0]}' 或脚本 '{command_parts[1]}'")
|
||||
except KeyError as e:
|
||||
logger.error(f"执行 length_cdf.py 脚本失败：配置项缺失 {str(e)}")
|
||||
except Exception as e:
|
||||
logger.error(f"执行 length_cdf.py 脚本时发生未知错误: {str(e)}")
|
||||
|
||||
def get_csv_files(self):
|
||||
"""遍历文件夹获取所有CSV文件路径，并按文件名中的起始序号排序"""
|
||||
|
||||
csv_files = []
|
||||
for chat_obj_folder in os.listdir(self.csv_folder):
|
||||
chat_obj_folder_path = os.path.join(self.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)
|
||||
csv_files.append(csvfile_path)
|
||||
# 提取文件名中的起始数字，比如 wxid_..._0_5000.csv → 0
|
||||
pattern = re.compile(r"_(\d+)_\d+\.csv$")
|
||||
|
||||
def extract_start(fp: str) -> int:
|
||||
name = os.path.basename(fp)
|
||||
m = pattern.search(name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
# 按起始数字升序排序
|
||||
csv_files.sort(key=extract_start)
|
||||
return csv_files
|
||||
|
||||
def match_qa(self, messages: List[ChatMessage]) -> List[Union[QaPair, CutMessage]]:
|
||||
"""
|
||||
匹配问答对
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
|
||||
Returns:
|
||||
List[Union[QaPair, CutMessage]]: 包含指令和输出的问答对列表
|
||||
"""
|
||||
# 状态定义
|
||||
WAITING_INSTRUCTION = "waiting_instruction" # 等待指令
|
||||
WAITING_RESPONSE = "waiting_response" # 等待回复
|
||||
|
||||
current_state = WAITING_INSTRUCTION
|
||||
qa_res: List[Union[QaPair, CutMessage]] = []
|
||||
last_message = None
|
||||
current_instruction = None
|
||||
qa_id_counter = 0
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, CutMessage):
|
||||
current_state = WAITING_INSTRUCTION
|
||||
current_instruction = None
|
||||
last_message = None
|
||||
if self.c["prompt_with_history"]:
|
||||
qa_res.append(msg)
|
||||
continue
|
||||
|
||||
if current_state == WAITING_INSTRUCTION:
|
||||
if msg.is_sender == 0: # 收到对方消息
|
||||
current_instruction = msg.msg
|
||||
last_message = msg
|
||||
current_state = WAITING_RESPONSE
|
||||
|
||||
elif current_state == WAITING_RESPONSE:
|
||||
if msg.is_sender == 0: # 收到对方消息
|
||||
current_instruction = msg.msg
|
||||
last_message = msg
|
||||
# 状态保持不变
|
||||
else: # 自己的回复 使用策略判断是否属于同一对话
|
||||
if last_message and self.qa_match_strategy.is_same_conversation([last_message], msg):
|
||||
assert current_instruction is not None, (
|
||||
"current_instruction should not be None when creating a QA pair"
|
||||
)
|
||||
qa_pair = QaPair(
|
||||
id=qa_id_counter,
|
||||
system=self.system_prompt,
|
||||
instruction=current_instruction,
|
||||
output=msg.msg,
|
||||
history=[], # No history in this context yet
|
||||
time=msg.CreateTime, # Use the response message time
|
||||
score=0, # Default score
|
||||
)
|
||||
qa_res.append(qa_pair)
|
||||
qa_id_counter += 1 # 增加计数器
|
||||
else:
|
||||
if self.c["prompt_with_history"]:
|
||||
qa_res.append(
|
||||
CutMessage(
|
||||
is_sender=msg.is_sender,
|
||||
cut_type=msg.type_name,
|
||||
CreateTime=msg.CreateTime,
|
||||
)
|
||||
)
|
||||
# 无论是否匹配，都重置状态
|
||||
current_state = WAITING_INSTRUCTION
|
||||
current_instruction = None
|
||||
last_message = None
|
||||
|
||||
return qa_res
|
||||
|
||||
# TODO: need review
|
||||
def add_history_to_qa(self, qa_res: List[Union[QaPair, CutMessage]]) -> List[QaPair]:
|
||||
"""
|
||||
Adds conversation history to QaPair objects.
|
||||
|
||||
Args:
|
||||
qa_res: A list containing QaPair and CutMessage objects.
|
||||
|
||||
Returns:
|
||||
A list of QaPair objects with history populated.
|
||||
"""
|
||||
qa_res_with_history: List[QaPair] = []
|
||||
current_history: List[List[str]] = []
|
||||
last_timestamp: Timestamp = None # type: ignore
|
||||
|
||||
for item in qa_res:
|
||||
if isinstance(item, CutMessage):
|
||||
if current_history:
|
||||
instruction = current_history[-1][0]
|
||||
output = current_history[-1][1]
|
||||
history = current_history[:-1]
|
||||
qa_pair_with_history = QaPair(
|
||||
id=-1,
|
||||
system=self.system_prompt,
|
||||
instruction=instruction,
|
||||
output=output,
|
||||
history=history,
|
||||
time=last_timestamp,
|
||||
score=0,
|
||||
)
|
||||
qa_res_with_history.append(qa_pair_with_history)
|
||||
current_history = []
|
||||
last_timestamp = None # type: ignore
|
||||
elif isinstance(item, QaPair):
|
||||
current_history.append([item.instruction, item.output])
|
||||
last_timestamp = item.time
|
||||
|
||||
if current_history:
|
||||
instruction = current_history[-1][0]
|
||||
output = current_history[-1][1]
|
||||
history = current_history[:-1]
|
||||
# Ensure last_timestamp is not None before assignment
|
||||
final_timestamp_end = last_timestamp
|
||||
assert final_timestamp_end is not None, "Timestamp cannot be None for the final QaPair"
|
||||
qa_pair_with_history = QaPair(
|
||||
id=-1,
|
||||
system=self.system_prompt,
|
||||
instruction=instruction,
|
||||
output=output,
|
||||
history=history,
|
||||
time=final_timestamp_end,
|
||||
score=0,
|
||||
)
|
||||
qa_res_with_history.append(qa_pair_with_history)
|
||||
|
||||
return qa_res_with_history
|
||||
|
||||
def group_consecutive_messages(self, messages: List[ChatMessage]) -> List[ChatMessage]:
|
||||
"""
|
||||
将同一个人连续发送的多条消息组合成一条消息，遇到cut_type添加cut
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
|
||||
Returns:
|
||||
List[ChatMessage]: 组合后的消息列表
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
def _combine_text(messages: List[ChatMessage]) -> ChatMessage:
|
||||
"""
|
||||
合并多条消息为一条
|
||||
|
||||
Args:
|
||||
messages: 要合并的消息列表
|
||||
|
||||
Returns:
|
||||
ChatMessage: 合并后的消息
|
||||
"""
|
||||
base_msg = messages[0]
|
||||
combined_content = messages[0].msg
|
||||
|
||||
for i in messages[1:]:
|
||||
content = i.msg
|
||||
if not content:
|
||||
continue
|
||||
|
||||
if combined_content and combined_content[-1] not in ["。", "！", "？", "…", "，", "."]:
|
||||
combined_content += "，"
|
||||
|
||||
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]}"
|
||||
)
|
||||
combined_content = combined_content[: self.c["combine_msg_max_length"]]
|
||||
|
||||
combined_message = ChatMessage(
|
||||
id=base_msg.id,
|
||||
MsgSvrID=base_msg.MsgSvrID,
|
||||
type_name=base_msg.type_name,
|
||||
is_sender=base_msg.is_sender,
|
||||
talker=base_msg.talker,
|
||||
room_name=base_msg.room_name,
|
||||
msg=combined_content,
|
||||
src=base_msg.src,
|
||||
CreateTime=messages[-1].CreateTime, # 使用最后一条消息的时间
|
||||
)
|
||||
|
||||
return combined_message
|
||||
|
||||
def _create_cut_message(message: ChatMessage) -> CutMessage:
|
||||
return CutMessage(
|
||||
is_sender=message.is_sender,
|
||||
cut_type=message.type_name,
|
||||
CreateTime=message.CreateTime,
|
||||
)
|
||||
|
||||
def _combine_current_group(group):
|
||||
"""
|
||||
处理当前消息组并添加到grouped_messages
|
||||
|
||||
Args:
|
||||
group: 当前消息组
|
||||
"""
|
||||
if len(group) > 1:
|
||||
combined_msg = _combine_text(group)
|
||||
grouped_messages.append(combined_msg)
|
||||
else:
|
||||
grouped_messages.append(group[0])
|
||||
|
||||
grouped_messages = []
|
||||
current_group = []
|
||||
|
||||
for _, current_msg in enumerate(messages):
|
||||
if current_msg.type_name in self.cut_type_list:
|
||||
if current_group:
|
||||
# 当前组有消息，合并当前组，并添加一条cut
|
||||
_combine_current_group(current_group)
|
||||
current_group = []
|
||||
|
||||
cut_msg = _create_cut_message(current_msg)
|
||||
grouped_messages.append(cut_msg)
|
||||
else:
|
||||
# 当前组没消息，检查上一个组
|
||||
if grouped_messages:
|
||||
if not isinstance(grouped_messages[-1], CutMessage):
|
||||
cut_msg = _create_cut_message(current_msg)
|
||||
grouped_messages.append(cut_msg)
|
||||
# 如果上一个组没消息或最后一条是CutMessage，直接continue
|
||||
continue
|
||||
|
||||
if not current_group:
|
||||
current_group = [current_msg]
|
||||
continue
|
||||
|
||||
last_msg = current_group[-1]
|
||||
|
||||
# 判断是否是同一个人的连续消息
|
||||
if (
|
||||
current_msg.is_sender == last_msg.is_sender
|
||||
and current_msg.talker == last_msg.talker
|
||||
and self.single_combine_strategy.is_same_conversation([last_msg], current_msg)
|
||||
):
|
||||
current_group.append(current_msg)
|
||||
else:
|
||||
# 不是同一个人的消息，处理当前组并开始新组
|
||||
_combine_current_group(current_group)
|
||||
# 开始新组
|
||||
current_group = [current_msg]
|
||||
|
||||
# 处理最后一组消息
|
||||
if current_group:
|
||||
_combine_current_group(current_group)
|
||||
|
||||
return grouped_messages
|
||||
|
||||
def process_by_msgtype(self, chat_message: ChatMessage):
|
||||
if chat_message.type_name == "文本":
|
||||
self.process_text(chat_message)
|
||||
# elif chat_message.type_name == "图片":
|
||||
# self.process_image(chat_message)
|
||||
|
||||
def load_csv(self, file_path) -> List[ChatMessage]:
|
||||
"""
|
||||
做整体第一次预处理，过滤不符合条件的行
|
||||
"""
|
||||
df = pd.read_csv(file_path, encoding="utf-8", dtype={"msg": str})
|
||||
|
||||
df = df[~df["type_name"].isin(values=skip_type_list)]
|
||||
|
||||
# 如果type_name为文本 并且msg 包含 手机号、身份证号、邮箱、网址则删除这行
|
||||
for i in df.index:
|
||||
if df.loc[i, "type_name"] == "文本":
|
||||
msg_str = str(df.loc[i, "msg"])
|
||||
if (
|
||||
re.search(r"1\d{10}", msg_str)
|
||||
or re.search(r"\d{18}", msg_str)
|
||||
or re.search(r"\w+@\w+", msg_str)
|
||||
or "http" in msg_str
|
||||
or r"\\xa0" in msg_str
|
||||
or r"\\u" in msg_str
|
||||
):
|
||||
df = df.drop(index=i)
|
||||
continue
|
||||
for blocked_word in self.blocked_words:
|
||||
if blocked_word in msg_str:
|
||||
df = df.drop(index=i)
|
||||
break
|
||||
else:
|
||||
df.loc[i, "msg"] = ""
|
||||
|
||||
df = df.dropna(how="all")
|
||||
# 时间格式 2021-07-07 10:27:23
|
||||
# 遍历行 相同is_sender的行合并msg（）遇到不同is_sender就重新开始
|
||||
df["CreateTime"] = pd.to_datetime(df["CreateTime"])
|
||||
|
||||
return [ChatMessage(*row) for row in df.values]
|
||||
|
||||
def process_text(self, chat_message: ChatMessage):
|
||||
pass
|
||||
|
||||
def save_result(self, qa_res: List[QaPair]):
|
||||
"""
|
||||
Saves the list of QaPair objects to a JSON file after converting them to dictionaries.
|
||||
|
||||
Args:
|
||||
qa_res: A list of QaPair objects.
|
||||
"""
|
||||
processed_qa_res = []
|
||||
for idx, item in enumerate(qa_res):
|
||||
item_dict = {
|
||||
"id": idx,
|
||||
"system": item.system,
|
||||
"instruction": item.instruction,
|
||||
"output": item.output,
|
||||
"history": item.history,
|
||||
"time": item.time.isoformat() if item.time else None,
|
||||
"score": item.score,
|
||||
}
|
||||
processed_qa_res.append(item_dict)
|
||||
|
||||
output_path = "./dataset/res_csv/sft/sft-my.json"
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
json.dump(processed_qa_res, f, ensure_ascii=False, indent=4)
|
||||
logger.success(f"聊天记录处理成功，共{len(qa_res)}条，保存到 {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
processor = DataProcessor()
|
||||
processor.main()
|
||||
@@ -1,3 +1,4 @@
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
@@ -8,8 +9,9 @@ from typing import List, Union, cast
|
||||
import pandas as pd
|
||||
from pandas import Timestamp
|
||||
|
||||
from weclone.data.clean.strategies import LLMCleaningStrategy
|
||||
from weclone.data.clean.strategies_online import OlineLLMCleaningStrategy
|
||||
from weclone.data.clean.strategies import LLMCleaningStrategy, OlineLLMCleaningStrategy
|
||||
|
||||
# from weclone.data.clean.strategies_online import OlineLLMCleaningStrategy
|
||||
from weclone.data.models import (
|
||||
ChatMessage,
|
||||
CutMessage,
|
||||
@@ -19,7 +21,7 @@ from weclone.data.models import (
|
||||
skip_type_list,
|
||||
)
|
||||
from weclone.data.strategies import LLMStrategy, TimeWindowStrategy
|
||||
from weclone.data.utils import check_image_file_exists
|
||||
from weclone.data.utils import ImageToTextProcessor, check_image_file_exists
|
||||
from weclone.utils.config_models import DataModality, PlatformType, WCMakeDatasetConfig
|
||||
from weclone.utils.configV2 import load_config
|
||||
from weclone.utils.log import logger
|
||||
@@ -86,9 +88,7 @@ class DataProcessor:
|
||||
|
||||
if clean_dataset_config.clean_strategy == "llm":
|
||||
if self.config.online_llm_clear:
|
||||
self.clean_strategy = OlineLLMCleaningStrategy(
|
||||
make_dataset_config=self.config.model_dump(mode="json")
|
||||
)
|
||||
self.clean_strategy = OlineLLMCleaningStrategy(make_dataset_config=self.config)
|
||||
else:
|
||||
from llamafactory.extras.packages import is_vllm_available
|
||||
|
||||
@@ -97,12 +97,71 @@ class DataProcessor:
|
||||
# 注意：这里我们不能直接修改config对象的属性，因为它是不可变的
|
||||
self.enable_clean = False
|
||||
else:
|
||||
self.clean_strategy = LLMCleaningStrategy(
|
||||
make_dataset_config=self.config.model_dump(mode="json")
|
||||
)
|
||||
self.clean_strategy = LLMCleaningStrategy(make_dataset_config=self.config)
|
||||
|
||||
# 基于配置初始化图片识别处理器
|
||||
vision_config = self.config.vision_api
|
||||
if vision_config.enable and vision_config.api_key:
|
||||
self.image_processor = ImageToTextProcessor(
|
||||
api_url=vision_config.api_url,
|
||||
api_key=vision_config.api_key,
|
||||
model_name=vision_config.model_name,
|
||||
)
|
||||
logger.info(f"已启用图片识别功能, 模型: {self.image_processor.model_name}")
|
||||
else:
|
||||
self.image_processor = None
|
||||
|
||||
self.c = self.config
|
||||
|
||||
def _process_images_in_parallel(self, qa_list: List[QaPairV2]) -> List[QaPairV2]:
|
||||
"""并行处理所有对话中的图片，并将描述替换回对话文本。"""
|
||||
all_image_paths = []
|
||||
media_dir = self.c.media_dir
|
||||
|
||||
# 遍历所有对话，收集并构造完整的图片路径
|
||||
for qa_pair in qa_list:
|
||||
if qa_pair.images:
|
||||
image_list = qa_pair.images if isinstance(qa_pair.images, list) else [qa_pair.images]
|
||||
for relative_path in image_list:
|
||||
full_path = os.path.join(media_dir, relative_path)
|
||||
all_image_paths.append(full_path)
|
||||
|
||||
if not all_image_paths:
|
||||
logger.info("未在对话中找到任何图片，跳过识别。")
|
||||
return qa_list
|
||||
|
||||
logger.info(f"共找到 {len(all_image_paths)} 张有效图片需要识别。")
|
||||
max_workers = self.c.vision_api.max_workers
|
||||
|
||||
# 使用线程池并行调用API，executor.map 会保持结果顺序与输入一致
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# 现在传递给 image_processor 的是完整的路径
|
||||
image_descriptions = list(executor.map(self.image_processor.describe_image, all_image_paths))
|
||||
|
||||
desc_iterator = iter(image_descriptions)
|
||||
for qa_pair in qa_list:
|
||||
if not qa_pair.images:
|
||||
continue
|
||||
|
||||
for message in qa_pair.messages:
|
||||
# 替换消息内容中的 <image> 占位符
|
||||
num_images_in_message = message.content.count("<image>")
|
||||
for _ in range(num_images_in_message):
|
||||
try:
|
||||
description = next(desc_iterator)
|
||||
# 使用 count=1 确保每次只替换一个占位符，并添加换行符以增强可读性
|
||||
message.content = message.content.replace(
|
||||
"<image>", f"\n[图片描述: {description}]\n", 1
|
||||
)
|
||||
except StopIteration:
|
||||
logger.error("图片数量与描述数量不匹配，可能存在逻辑错误。")
|
||||
message.content = message.content.replace("<image>", "\n[图片描述缺失]\n", 1)
|
||||
|
||||
# 清空图片列表，因为它们已被转换为文本
|
||||
qa_pair.images.clear()
|
||||
|
||||
return qa_list
|
||||
|
||||
def main(self):
|
||||
if not os.path.exists(self.csv_folder) or not os.listdir(self.csv_folder):
|
||||
logger.error(
|
||||
@@ -122,6 +181,12 @@ class DataProcessor:
|
||||
qa_res = self.match_qa(message_list)
|
||||
qa_res = [item for item in qa_res if isinstance(item, QaPairV2)]
|
||||
|
||||
# 如果启用图片识别，则执行并行处理
|
||||
if self.image_processor:
|
||||
logger.info("开始执行图片识别流程...")
|
||||
qa_res = self._process_images_in_parallel(qa_res)
|
||||
logger.info("图片识别流程完成。")
|
||||
|
||||
if self.enable_clean:
|
||||
self.clean_strategy.judge(qa_res) # type: ignore
|
||||
# qa_res = self.clean_strategy.clean(qa_res) #改到sft.py中
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
from weclone.utils.log import logger
|
||||
|
||||
|
||||
@@ -34,6 +38,119 @@ def check_image_file_exists(file_path: str) -> str | bool:
|
||||
return False
|
||||
|
||||
|
||||
class ImageToTextProcessor:
|
||||
"""通过兼容OpenAI API的多模态LLM将图片转换为文本。"""
|
||||
|
||||
def __init__(self, api_url: str, api_key: str, model_name: str):
|
||||
self.api_url = api_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model_name = model_name
|
||||
self.prompt = """
|
||||
请描述这张图片的内容，重点关注：
|
||||
1. 如果是截图，描述界面内容和操作
|
||||
2. 如果是表格，描述表格结构和数据
|
||||
3. 如果是文档，提取关键文字信息
|
||||
4. 如果是生活照片，简要描述场景和内容。
|
||||
请用简洁明了的语言描述，不超过100字。"""
|
||||
|
||||
def _encode_image_to_base64(self, image_path: str) -> str:
|
||||
"""将图片编码为base64"""
|
||||
try:
|
||||
with open(image_path, "rb") as image_file:
|
||||
return base64.b64encode(image_file.read()).decode("utf-8")
|
||||
except Exception as e:
|
||||
logger.error(f"编码图片失败 {image_path}: {e}")
|
||||
return None
|
||||
|
||||
def _get_image_format(self, image_path: str) -> str:
|
||||
"""获取图片格式"""
|
||||
suffix = Path(image_path).suffix.lower().replace(".", "")
|
||||
if suffix == "jpg":
|
||||
return "jpeg"
|
||||
return suffix
|
||||
|
||||
def _call_vision_api(self, image_path: str) -> str:
|
||||
"""调用Vision API（增加了重试机制）"""
|
||||
base64_image = self._encode_image_to_base64(image_path)
|
||||
if not base64_image:
|
||||
return "[图片处理失败：无法编码]"
|
||||
|
||||
image_format = self._get_image_format(image_path)
|
||||
|
||||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
payload = {
|
||||
"model": self.model_name,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": self.prompt},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/{image_format};base64,{base64_image}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 1000,
|
||||
"temperature": 0.1,
|
||||
}
|
||||
|
||||
# --- 重试逻辑 ---
|
||||
max_retries = 5 # 最大重试次数
|
||||
base_delay = 5 # 基础等待时间（秒）
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.api_url}/chat/completions", headers=headers, json=payload, timeout=60
|
||||
)
|
||||
if response.status_code == 200:
|
||||
pass
|
||||
elif response.status_code in [429, 500, 502, 503, 504]:
|
||||
response.raise_for_status()
|
||||
else:
|
||||
logger.error(f"API请求失败，状态码: {response.status_code}，原因: {response.reason}")
|
||||
return "[图片描述获取失败]"
|
||||
|
||||
result = response.json()
|
||||
|
||||
if "choices" in result and len(result["choices"]) > 0:
|
||||
content = result["choices"][0]["message"]["content"]
|
||||
return content.strip()
|
||||
else:
|
||||
logger.warning(f"API响应格式异常: {result}")
|
||||
return "[图片描述获取失败：API格式错误]"
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.warning(f"API请求失败 (尝试 {attempt + 1}/{max_retries}): {e}")
|
||||
if attempt < max_retries - 1:
|
||||
# 指数退避等待
|
||||
wait_time = base_delay * (2**attempt)
|
||||
logger.info(f"将在 {wait_time} 秒后重试...")
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
logger.error(f"API请求在 {max_retries} 次尝试后最终失败: {image_path}")
|
||||
return "[图片描述获取失败：请求异常]"
|
||||
except Exception as e:
|
||||
logger.error(f"处理API响应时出现未知错误 {image_path}: {e}")
|
||||
# 对于未知错误，可以选择不重试，直接返回
|
||||
return "[图片描述获取失败：未知错误]"
|
||||
|
||||
# 如果循环结束仍未成功（理论上不会执行到这里，因为上面已有返回）
|
||||
return "[图片描述获取失败：所有重试均失败]"
|
||||
|
||||
def describe_image(self, image_path: str) -> str:
|
||||
"""公开方法，用于描述单张图片内容"""
|
||||
if not os.path.exists(image_path):
|
||||
logger.warning(f"图片文件不存在: {image_path}")
|
||||
return "[图片文件不存在]"
|
||||
|
||||
logger.debug(f"正在识别图片: {os.path.basename(image_path)}")
|
||||
return self._call_vision_api(image_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
path = "Storage\\Image\2021-08\6ce3f785b4230246639c3dd0d4a8848c.dat"
|
||||
print(check_image_file_exists(path))
|
||||
|
||||
@@ -20,11 +20,29 @@ def main():
|
||||
if device == "cpu":
|
||||
logger.warning("请注意你正在使用CPU训练，非Mac设备可能会出现问题")
|
||||
|
||||
cleaner = LLMCleaningStrategy(make_dataset_config=dataset_config.model_dump(mode="json"))
|
||||
cleaned_data_path = cleaner.clean()
|
||||
cleaner = LLMCleaningStrategy(make_dataset_config=dataset_config)
|
||||
final_dataset_name = cleaner.clean()
|
||||
|
||||
if not os.path.exists(cleaned_data_path):
|
||||
logger.error(f"错误：文件 '{cleaned_data_path}' 不存在，请确保数据处理步骤已正确生成该文件。")
|
||||
if train_config.dataset != final_dataset_name:
|
||||
logger.info(
|
||||
f"根据清洗结果，将训练数据集从 '{train_config.dataset}' 动态更新为 '{final_dataset_name}'。"
|
||||
)
|
||||
train_config.dataset = final_dataset_name
|
||||
|
||||
dataset_info_path = os.path.join(train_config.dataset_dir, "dataset_info.json")
|
||||
try:
|
||||
with open(dataset_info_path, "r", encoding="utf-8") as f:
|
||||
dataset_info = json.load(f)
|
||||
target_file_name = dataset_info.get(final_dataset_name, {}).get("file_name")
|
||||
if not target_file_name:
|
||||
raise FileNotFoundError(
|
||||
f"在 dataset_info.json 中未找到数据集 '{final_dataset_name}' 的 file_name 配置。"
|
||||
)
|
||||
final_data_path = os.path.join(train_config.dataset_dir, target_file_name)
|
||||
if not os.path.exists(final_data_path):
|
||||
raise FileNotFoundError(f"最终要使用的SFT数据文件 '{final_data_path}' 不存在。")
|
||||
except (FileNotFoundError, json.JSONDecodeError) as e:
|
||||
logger.error(f"校验最终数据集时出错: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
formatted_config = json.dumps(train_config.model_dump(mode="json"), indent=4, ensure_ascii=False)
|
||||
|
||||
@@ -38,7 +38,10 @@ def load_config(arg_type: str):
|
||||
)
|
||||
config["dataset"] = "wechat-sft-with-history"
|
||||
if "image" in s_config["make_dataset_args"]["include_type"]:
|
||||
config["dataset"] = "wechat-mllm-sft"
|
||||
if config["vision_api"].get("enable", False):
|
||||
config["dataset"] = "wechat-img-rec-sft" # 图像识别类模型使用的数据集
|
||||
else:
|
||||
config["dataset"] = "wechat-mllm-sft" # 多模态模型使用的数据集
|
||||
|
||||
elif arg_type == "make_dataset":
|
||||
config = {**s_config["make_dataset_args"], **s_config["common_args"]}
|
||||
@@ -46,7 +49,10 @@ def load_config(arg_type: str):
|
||||
config["dataset_dir"] = s_config["train_sft_args"]["dataset_dir"]
|
||||
config["cutoff_len"] = s_config["train_sft_args"]["cutoff_len"]
|
||||
if "image" in config["include_type"]:
|
||||
config["dataset"] = "wechat-mllm-sft"
|
||||
if config["vision_api"].get("enable", False):
|
||||
config["dataset"] = "wechat-img-rec-sft" # 图像识别类模型使用的数据集
|
||||
else:
|
||||
config["dataset"] = "wechat-mllm-sft" # 多模态模型使用的数据集
|
||||
|
||||
else:
|
||||
raise ValueError("暂不支持的参数类型")
|
||||
|
||||
@@ -61,6 +61,7 @@ def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel:
|
||||
elif arg_type == "train_sft":
|
||||
config_dict = {**common_config, **wc_config.train_sft_args.model_dump()}
|
||||
config_dict["include_type"] = wc_config.make_dataset_args.include_type
|
||||
config_dict["vision_api"] = wc_config.make_dataset_args.vision_api
|
||||
return WCTrainSftConfig(**config_dict)
|
||||
|
||||
elif arg_type == "make_dataset":
|
||||
|
||||
@@ -93,6 +93,16 @@ class CleanDatasetConfig(BaseModel):
|
||||
llm: LLMCleanConfig = LLMCleanConfig(accept_score=2)
|
||||
|
||||
|
||||
class VisionApiConfig(BaseModel):
|
||||
"""Vision API specific configuration"""
|
||||
|
||||
enable: bool = Field(False, description="是否启用Vision API进行图像识别")
|
||||
api_key: Optional[str] = None
|
||||
api_url: Optional[str] = None
|
||||
model_name: Optional[str] = None
|
||||
max_workers: Optional[int] = None
|
||||
|
||||
|
||||
class MakeDatasetArgs(BaseModel):
|
||||
platform: PlatformType = Field(..., description="Data source platform")
|
||||
include_type: List[DataModality] = Field([DataModality.TEXT], description="包含的数据类型")
|
||||
@@ -113,6 +123,7 @@ class MakeDatasetArgs(BaseModel):
|
||||
llm_api_key: Optional[str] = Field(None, description="在线LLM的api_key")
|
||||
model_name: Optional[str] = Field(None, description="在线LLM的模型名称, 建议使用参数较大的模型")
|
||||
clean_batch_size: int = Field(10, description="数据清洗批次大小")
|
||||
vision_api: VisionApiConfig = Field(default_factory=VisionApiConfig)
|
||||
|
||||
|
||||
class TrainSftArgs(BaseModel):
|
||||
|
||||
Reference in New Issue
Block a user