更新pyproject.toml以添加pyright类型检查配置，修改settings.json以调整数据集参数，重构qa_generator.py以支持新的消息处理策略，优化数据处理逻辑，更新测试用例以适应新功能。

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__":

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) 
+    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)

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"
+

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 "
+    }
 }

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()

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()

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