mirror of
https://github.com/xming521/WeClone.git
synced 2026-08-28 18:07:28 +08:00
pre-commit-ci[bot]
509 lines
20 KiB
Python
509 lines
20 KiB
Python
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()
|