From 055fb01792d573008db8daa1978b9916e7de8d9f Mon Sep 17 00:00:00 2001 From: xming521 <1223398803@qq.com> Date: Sun, 13 Apr 2025 15:27:00 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0pyproject.toml=E4=BB=A5?= =?UTF-8?q?=E6=B7=BB=E5=8A=A0pyright=E7=B1=BB=E5=9E=8B=E6=A3=80=E6=9F=A5?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=EF=BC=8C=E4=BF=AE=E6=94=B9settings.json?= =?UTF-8?q?=E4=BB=A5=E8=B0=83=E6=95=B4=E6=95=B0=E6=8D=AE=E9=9B=86=E5=8F=82?= =?UTF-8?q?=E6=95=B0=EF=BC=8C=E9=87=8D=E6=9E=84qa=5Fgenerator.py=E4=BB=A5?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=96=B0=E7=9A=84=E6=B6=88=E6=81=AF=E5=A4=84?= =?UTF-8?q?=E7=90=86=E7=AD=96=E7=95=A5=EF=BC=8C=E4=BC=98=E5=8C=96=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=A4=84=E7=90=86=E9=80=BB=E8=BE=91=EF=BC=8C=E6=9B=B4?= =?UTF-8?q?=E6=96=B0=E6=B5=8B=E8=AF=95=E7=94=A8=E4=BE=8B=E4=BB=A5=E9=80=82?= =?UTF-8?q?=E5=BA=94=E6=96=B0=E5=8A=9F=E8=83=BD=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- make_dataset/qa_generator.py | 210 +++++++++++-------- make_dataset/strategies.py | 51 +++-- pyproject.toml | 13 ++ settings.json | 15 +- tests/test_old_csv_to_json copy.py | 311 +++++++++++++++++++++++++++++ tests/test_old_csv_to_json.py | 262 ++++++++++++++---------- tests/test_qa_generator.py | 115 ++++++----- 7 files changed, 717 insertions(+), 260 deletions(-) create mode 100644 tests/test_old_csv_to_json copy.py diff --git a/make_dataset/qa_generator.py b/make_dataset/qa_generator.py index 1e577f9..f1a9606 100644 --- a/make_dataset/qa_generator.py +++ b/make_dataset/qa_generator.py @@ -1,9 +1,9 @@ -from dataclasses import dataclass import sys import os -from typing import List +from typing import Dict, List +from collections import deque +import re -from pandas import Timestamp import pandas as pd import json @@ -12,13 +12,12 @@ root_dir = os.path.dirname(current_dir) sys.path.append(root_dir) from src.utils.config import load_config from make_dataset.models import ChatMessage, CutMessage, skip_type_list -from make_dataset.strategies import TimeWindowStrategy, LagerModelStrategy +from make_dataset.strategies import TimeWindowStrategy, LLMStrategy class DataProcessor: def __init__(self): self.config = load_config(arg_type="make_dataset") - self.data = None self.csv_folder = "./data/csv" self.cut_type_list = [ "图片", @@ -36,13 +35,25 @@ class DataProcessor: "粘贴的文本", # 无法解析的分享链接 ] - # 根据self.config.make_dataset_args.conversation_strategy 判断初始化哪一个策略类 - if self.config["conversation_strategy"] == "time_window": - self.conversation_strategy = TimeWindowStrategy( - time_window=self.config["time_window"] + 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["conversation_strategy"] == "lager_model": - self.conversation_strategy = LagerModelStrategy() + 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) + + self.c = self.config def get_csv_files(self): """遍历文件夹获取所有CSV文件路径""" @@ -61,9 +72,94 @@ class DataProcessor: message_list: List[ChatMessage] = [] for csv_file in csv_files: chat_messages = self.load_csv(csv_file) - # 第一次预处理后 将chat_message 加入rcsv_df_list - message_list.append(self.group_consecutive_messages(messages=chat_messages)) + message_list.extend(self.group_consecutive_messages(messages=chat_messages)) # self.process_by_msgtype(chat_message) + qa_res = self.match_qa(message_list) + if self.c["prompt_with_history"]: + qa_res = self.add_history_to_qa(qa_res) + self.save_result(qa_res) + + def match_qa(self, messages: List[ChatMessage]) -> List[Dict]: + """ + 匹配问答对 + + Args: + messages: 消息列表 + + Returns: + List[Dict]: 包含指令和输出的问答对列表 + """ + # 状态定义 + WAITING_INSTRUCTION = "waiting_instruction" # 等待指令 + WAITING_RESPONSE = "waiting_response" # 等待回复 + + current_state = WAITING_INSTRUCTION + qa_res = [] + last_message = None + current_instruction = None + + for msg in messages: + # 检查是否为CutMessage + 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 + ): + qa_res.append( + {"instruction": current_instruction, "output": msg.msg} + ) + # 无论是否匹配，都重置状态 + current_state = WAITING_INSTRUCTION + current_instruction = None + last_message = None + + return qa_res + + def add_history_to_qa(self, qa_res: List[Dict]): + qa_res_with_history = [] + last_res = {"instruction": "", "output": "", "history": []} + + for _, qa in enumerate(qa_res): + if isinstance(qa, CutMessage): + if len(last_res["history"]) == 0: + continue + else: + if len(last_res["history"]) == 1: + last_res = { + "instruction": last_res["history"][0][0], + "output": last_res["history"][0][1], + "history": [], + } + else: + last_res = { + "instruction": last_res["history"][-1][0], + "output": last_res["history"][-1][1], + "history": last_res["history"][:-1], + } + qa_res_with_history.append(last_res) + last_res = {"instruction": "", "output": "", "history": []} + else: + last_res["history"].append([qa["instruction"], qa["output"]]) + + return qa_res_with_history def group_consecutive_messages( self, messages: List[ChatMessage] @@ -125,15 +221,7 @@ class DataProcessor: return combined_message def _create_cut_message(message: ChatMessage) -> CutMessage: - """ - 创建一个CutMessage实例 - Args: - message: 当前处理的消息，用于获取属性 - - Returns: - CutMessage: 创建的CutMessage实例 - """ return CutMessage( is_sender=message.is_sender, cut_type=message.type_name, @@ -185,8 +273,9 @@ class DataProcessor: if ( current_msg.is_sender == last_msg.is_sender and current_msg.talker == last_msg.talker - and (current_msg.CreateTime - last_msg.CreateTime).total_seconds() - < 3600 + and self.single_combine_strategy.is_same_conversation( + [last_msg], current_msg + ) ): current_group.append(current_msg) else: @@ -201,48 +290,6 @@ class DataProcessor: return grouped_messages - def create_conversation_data(self, messages: List[ChatMessage]) -> dict: - """ - 将一组消息组成一条对话数据 - - Args: - messages: 属于同一对话的消息列表 - - Returns: - dict: 包含对话历史的数据字典 - """ - conversation = [] - for msg in messages: - if msg.type_name in self.config["include_type"]: - conversation.append( - { - "role": "user" if msg.is_sender == 0 else "assistant", - "content": msg.msg, - "timestamp": msg.CreateTime, - } - ) - - # 确保对话是按时间顺序排列的 - conversation.sort(key=lambda x: x["timestamp"]) - - # 限制历史长度 - if len(conversation) > self.config["history_length"]: - conversation = conversation[-self.config["history_length"] :] - - return { - "conversation": conversation, - "metadata": { - "conversation_id": str(messages[0].id) if messages else None, - "timestamp": messages[-1].CreateTime if messages else None, - }, - } - - 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 process_by_msgtype(self, chat_message: ChatMessage): if chat_message.type_name == "文本": self.process_text(chat_message) @@ -264,18 +311,19 @@ class DataProcessor: # 如果type_name为文本 并且msg 包含 手机号、身份证号、邮箱、网址则删除这行 for i in df.index: if df.loc[i, "type_name"] == "文本": + msg_str = str(df.loc[i, "msg"]) if ( - "1\d{10}" in df.loc[i, "msg"] - or "\d{18}" in df.loc[i, "msg"] - or "\w+@\w+" in df.loc[i, "msg"] - or "http" in df.loc[i, "msg"] - or r"\\xa0" in df.loc[i, "msg"] - or r"\\u" in df.loc[i, "msg"] + 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 blocked_words: - if blocked_word in df.loc[i, "msg"]: + if blocked_word in msg_str: df = df.drop(index=i) break else: @@ -292,17 +340,15 @@ class DataProcessor: pass - def process_image(self): - # 处理方法1 - pass - - def process_method2(self): - # 处理方法2 - pass - - def save_result(self): + def save_result(self, qa_res: List[Dict]): # 保存结果 - pass + with open( + f"./data/res_csv/sft/sft-{self.c['single_combine_strategy']}-{self.c['qa_match_strategy']}-my.json", + "w", + encoding="utf-8", + ) as f: + json.dump(qa_res, f, ensure_ascii=False) + print(f"聊天记录处理成功，共{len(qa_res)}条，保存到{f.name}") if __name__ == "__main__": diff --git a/make_dataset/strategies.py b/make_dataset/strategies.py index 9106669..4a457f1 100644 --- a/make_dataset/strategies.py +++ b/make_dataset/strategies.py @@ -1,39 +1,60 @@ from dataclasses import dataclass -from typing import Protocol, List +from typing import List from .models import ChatMessage +from abc import ABC, abstractmethod -class ConversationStrategy(Protocol): - def is_same_conversation(self, msg1: ChatMessage, msg2: ChatMessage) -> bool: + +@dataclass +class ConversationStrategy(ABC): + """对话策略的抽象基类""" + + is_single_chat: bool + + @abstractmethod + def is_same_conversation( + self, history_msg: List[ChatMessage], current_msg: ChatMessage + ) -> bool: """判断两条消息是否属于同一个对话""" - pass - @dataclass class TimeWindowStrategy(ConversationStrategy): """基于时间窗口的判断策略""" - time_window: int # 时间窗口（分钟） - - def is_same_conversation(self, msg1: ChatMessage, msg2: ChatMessage) -> bool: - time_diff = abs((msg2.timestamp - msg1.timestamp)) + time_window: int # 时间窗口（分钟） + + def is_same_conversation( + self, history_msg: List[ChatMessage], current_msg: ChatMessage + ) -> bool: + time_diff = abs( + (current_msg.CreateTime - history_msg[-1].CreateTime) + ).total_seconds() return time_diff <= self.time_window @dataclass -class LagerModelStrategy(ConversationStrategy): +class LLMStrategy(ConversationStrategy): """基于大模型判断策略""" - def is_same_conversation(self, msg1: ChatMessage, msg2: ChatMessage) -> bool: - return msg1.user_id == msg2.user_id + + def is_same_conversation( + self, history_msg: List[ChatMessage], current_msg: ChatMessage + ) -> bool: + # 修复user_id错误，使用talker字段代替user_id + return current_msg.talker == history_msg[-1].talker if history_msg else False @dataclass class CompositeStrategy(ConversationStrategy): """组合多个策略的复合策略""" + strategies: List[ConversationStrategy] require_all: bool = True # True表示所有策略都满足，False表示任一策略满足即可 - def is_same_conversation(self, msg1: ChatMessage, msg2: ChatMessage) -> bool: - results = [s.is_same_conversation(msg1, msg2) for s in self.strategies] - return all(results) if self.require_all else any(results) \ No newline at end of file + def is_same_conversation( + self, history_msg: List[ChatMessage], current_msg: ChatMessage + ) -> bool: + results = [ + s.is_same_conversation(history_msg, current_msg) for s in self.strategies + ] + return all(results) if self.require_all else any(results) diff --git a/pyproject.toml b/pyproject.toml index 78bce99..d6bfc14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,3 +45,16 @@ conflicts = [ [[tool.uv.index]] url = "https://pypi.tuna.tsinghua.edu.cn/simple/" default = true + +[tool.pyright] +typeCheckingMode = "basic" +include = ["src"] +exclude = ["**/archive","**/tests"] +ignore = ["**/archive"] + +reportMissingImports = "error" +reportMissingTypeStubs = false + +pythonVersion = "3.9" +pythonPlatform = "Linux" + diff --git a/settings.json b/settings.json index ab419f6..bd07b35 100644 --- a/settings.json +++ b/settings.json @@ -43,14 +43,14 @@ "top_p": 0.65 }, "make_dataset_args": { - "enable_vision_model": false, + // "enable_vision_model": false,//后续实现 "include_type": [ - "文本", - "图片" + "文本" ], - "history_length": 10, - "conversation_strategy": "time_window", // 基于时间窗口的判断策略 - "time_window": 10, // 时间窗口（分钟）, + "single_combine_strategy": "time_window", // 单人组成单句策略 + "qa_match_strategy": "time_window", // 多人组成qa策略 + "single_combine_time_window": 10, // 单人组成单句时间窗口（分钟）, + "qa_match_time_window": 60, // 多人组成qa时间窗口（分钟）, "prompt_with_history": false // 是否在prompt中包含历史对话 }, "common_args": { @@ -59,6 +59,5 @@ "template": "chatglm3-weclone", "finetuning_type": "lora", "trust_remote_code": true - }, - "_comment": "adapter_name_or_path同时做为train_sft_args的output_dir " + } } \ No newline at end of file diff --git a/tests/test_old_csv_to_json copy.py b/tests/test_old_csv_to_json copy.py new file mode 100644 index 0000000..c2dff02 --- /dev/null +++ b/tests/test_old_csv_to_json copy.py @@ -0,0 +1,311 @@ +import csv +import json +import os +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 ( + "1\d{10}" in chat_df.loc[i, "content"] + or "\d{18}" in chat_df.loc[i, "content"] + or "\w+@\w+" in 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="10m"): + # 如果超时 前面的添加到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("10m"): + # 相差超过1小时清空队列 + temp_res.clear() + last_CreateTime = row["CreateTime"] + else: + # 没有相差一小时入队再全部出队 + temp_res.append(row["content"]) + csv_res.append({"instruction": temp_res[0], "output": temp_res[1]}) + temp_res.clear() + last_CreateTime = row["CreateTime"] + + csv_res_df = pd.DataFrame(csv_res) + print(f"处理后数据量：{csv_res_df.shape[0]}") + csv_res_df.to_json('./data/res_csv/sft/sft-old-my.json', orient='records', force_ascii=False) + + +if __name__ == "__main__": + # make_pt_dataset() + make_sft_dataset() diff --git a/tests/test_old_csv_to_json.py b/tests/test_old_csv_to_json.py index f0f6ab5..ddb368a 100644 --- a/tests/test_old_csv_to_json.py +++ b/tests/test_old_csv_to_json.py @@ -12,30 +12,30 @@ sys.path.append(root_dir) from make_dataset.qa_generator import DataProcessor -csv_folder = './data/csv' +csv_folder = "./data/csv" # csv_folder = './data/test' os.chdir(root_dir) -print(f'当前处理目录{csv_folder}') +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] + 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']) + 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')] + 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["content"] chat_df = chat_df.dropna() return chat_df @@ -47,109 +47,151 @@ def make_pt_dataset(): 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'): + 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 = csv_res.apply(lambda x: {"c": x}) # 设置数据集prompt键为c - csv_res.to_json('./data/res_csv/pt-my.json', orient='records', force_ascii=False) + 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'] + 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)] + 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'] + chat_df["content"] = chat_df["msg"] # 如果type_name为文本 并且content 包含 手机号、身份证号、邮箱、网址则删除这行 for i in chat_df.index: - if chat_df.loc[i, 'type_name'] == '文本': - if ('1\d{10}' in chat_df.loc[i, 'content'] or - '\d{18}' in chat_df.loc[i, 'content'] or - '\w+@\w+' in 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']): + if chat_df.loc[i, "type_name"] == "文本": + 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 + ): chat_df = chat_df.drop(index=i) continue for blocked_word in blocked_words: - if blocked_word in chat_df.loc[i, 'content']: + 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.loc[i, "content"] = "" - chat_df = chat_df[['is_sender', 'type_name', 'content', 'CreateTime']] + 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']) + chat_df["CreateTime"] = pd.to_datetime(chat_df["CreateTime"]) - -# 改到这了 + # 改到这了 - type_list.remove('文本') + 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'] + 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 = '' + 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']}) + 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'] + 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='1h'): + if row["is_sender"] == last_is_sender: + if row["CreateTime"] - last_CreateTime > pd.Timedelta(value="10m"): # 如果超时 前面的添加到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'] + 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'] + 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'] + 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 @@ -157,71 +199,72 @@ def handle_sft_csv(csvfile): 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) - + 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': + 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': + 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() + 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)} 条记录") + # 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种情况 @@ -230,40 +273,39 @@ def make_sft_dataset(): for i, row in csv_concat.iterrows(): if len(temp_res) == 0: - if row['content'] == 'cut': + if row["content"] == "cut": continue - if row['is_sender'] == 0: - temp_res.append(row['content']) - last_CreateTime = row['CreateTime'] + 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': + if row["content"] == "cut": temp_res.clear() - last_CreateTime = row['CreateTime'] - elif row['is_sender'] == 0: + last_CreateTime = row["CreateTime"] + elif row["is_sender"] == 0: # 遇到0 清空队列再入队 temp_res.clear() - temp_res.append(row['content']) - last_CreateTime = row['CreateTime'] + temp_res.append(row["content"]) + last_CreateTime = row["CreateTime"] else: - if row['CreateTime'] - last_CreateTime > pd.Timedelta('1h'): + if row["CreateTime"] - last_CreateTime > pd.Timedelta("1h"): # 相差超过1小时清空队列 temp_res.clear() - last_CreateTime = row['CreateTime'] + last_CreateTime = row["CreateTime"] else: # 没有相差一小时入队再全部出队 - temp_res.append(row['content']) - csv_res.append({'instruction': temp_res[0], 'output': temp_res[1]}) + temp_res.append(row["content"]) + csv_res.append({"instruction": temp_res[0], "output": temp_res[1]}) temp_res.clear() - last_CreateTime = row['CreateTime'] - + 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-my.json', orient='records', force_ascii=False) + 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__': +if __name__ == "__main__": # make_pt_dataset() make_sft_dataset() diff --git a/tests/test_qa_generator.py b/tests/test_qa_generator.py index 5739986..3bc50d2 100644 --- a/tests/test_qa_generator.py +++ b/tests/test_qa_generator.py @@ -58,7 +58,8 @@ def test_single_message(processor): ) result = processor.group_consecutive_messages([message]) - assert len(result) == 1 + 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 == "你好" @@ -102,7 +103,8 @@ def test_consecutive_messages_same_sender(processor): ] result = processor.group_consecutive_messages(messages) - assert len(result) == 1 + 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 == "你好，最近怎么样，我想问个问题" @@ -146,7 +148,8 @@ def test_messages_different_senders(processor): ] result = processor.group_consecutive_messages(messages) - assert len(result) == 3 + 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 == "我想问个问题" @@ -167,6 +170,17 @@ def test_skip_non_text_messages(processor): 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, @@ -176,7 +190,7 @@ def test_skip_non_text_messages(processor): room_name="testroom", msg="", src="image.jpg", - CreateTime=now + timedelta(minutes=1), + CreateTime=now + timedelta(minutes=1+9.9), ), ChatMessage( id=3, @@ -187,13 +201,16 @@ def test_skip_non_text_messages(processor): room_name="testroom", msg="看到图片了吗", src="", - CreateTime=now + timedelta(minutes=2), + CreateTime=now + timedelta(minutes=1+9.9), ), ] result = processor.group_consecutive_messages(messages) - assert len(result) == 1 - assert result[0].msg == "你好，看到图片了吗" + 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): @@ -225,7 +242,8 @@ def test_time_window_limit(processor): ] result = processor.group_consecutive_messages(messages) - assert len(result) == 2 + 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 == "晚上好" @@ -236,90 +254,97 @@ def test_consecutive_messages_to_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, - }) + 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, - }) - + 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]}...") - + print( + f"{i+1}. {'用户' if msg['is_sender'] == 0 else '对方'}: {msg['msg'][:50]}..." + ) + return output_file