使用Python将WeChat聊天记录转为LangChain消息格式

本文链接：https://blog.csdn.net/qahaj/article/details/146484167

在AI模型的微调过程中，往往需要输入一些带有上下文的对话数据。然而，对于WeChat这样的流行聊天工具，直接导出聊天信息并不方便。本文将向您展示如何通过手动复制和粘贴聊天记录到文本文件，再利用自定义的加载器，将这些记录转换成LangChain所需的消息格式用于模型训练。此方法基于LangChain的Discord聊天加载器库，经过简单的修改可以适用于WeChat。

技术背景介绍

在进行大语言模型的微调和少样本学习时，拥有来自不同来源的对话数据可以提升模型的表现。在本案例中，目标是将WeChat聊天记录转换为LangChain框架所需的消息格式，进一步应用于AI模型的训练与推理。

核心原理解析

首先，通过WeChat桌面应用程序的复制功能，将所需聊天记录复制到一个文本文件中。接着，通过一个自定义的WeChat聊天加载器，解析文本文件内容，并将格式转换为LangChain消息对象。最后，将这些对象用作微调模型的输入。

代码实现演示

下面是具体的代码实现步骤：

1. 创建聊天记录文本文件

首先，从WeChat桌面应用程序中复制所需的消息，并粘贴到本地的文本文件中。例如：

女朋友 2023/09/16 2:51 PM
天气有点凉

男朋友 2023/09/16 2:51 PM
珍簟凉风著，瑶琴寄恨生。嵇君懒书札，底物慰秋情。

2. 定义WeChat聊天加载器

import logging
import re
from typing import Iterator, List
from langchain_community.chat_loaders import base as chat_loaders
from langchain_core.messages import BaseMessage, HumanMessage

logger = logging.getLogger()

class WeChatChatLoader(chat_loaders.BaseChatLoader):
    def __init__(self, path: str):
        self.path = path
        self._message_line_regex = re.compile(
            r"(?P<sender>.+?) (?P<timestamp>\d{4}/\d{2}/\d{2} \d{1,2}:\d{2} (?:AM|PM))"
        )

    def _append_message_to_results(self, results: List, current_sender: str, current_timestamp: str, current_content: List[str]):
        content = "\n".join(current_content).strip()
        if not re.match(r"\[.*\]", content):
            results.append(
                HumanMessage(content=content, additional_kwargs={"sender": current_sender, "events": [{"message_time": current_timestamp}]})
            )
        return results

    def _load_single_chat_session_from_txt(self, file_path: str) -> chat_loaders.ChatSession:
        with open(file_path, "r", encoding="utf-8") as file:
            lines = file.readlines()

        results: List[BaseMessage] = []
        current_sender = None
        current_timestamp = None
        current_content = []
        for line in lines:
            if re.match(self._message_line_regex, line):
                if current_sender and current_content:
                    results = self._append_message_to_results(results, current_sender, current_timestamp, current_content)
                current_sender, current_timestamp = re.match(self._message_line_regex, line).groups()
                current_content = []
            else:
                current_content.append(line.strip())

        if current_sender and current_content:
            results = self._append_message_to_results(results, current_sender, current_timestamp, current_content)

        return chat_loaders.ChatSession(messages=results)

    def lazy_load(self) -> Iterator[chat_loaders.ChatSession]:
        yield self._load_single_chat_session_from_txt(self.path)

3. 加载消息

from typing import List
from langchain_community.chat_loaders.utils import map_ai_messages, merge_chat_runs
from langchain_core.chat_sessions import ChatSession

loader = WeChatChatLoader(path="./wechat_chats.txt")
raw_messages = loader.lazy_load()
merged_messages = merge_chat_runs(raw_messages)
messages: List[ChatSession] = list(map_ai_messages(merged_messages, sender="男朋友"))

from langchain_openai import ChatOpenAI

llm = ChatOpenAI()

for chunk in llm.stream(messages[0]["messages"]):
    print(chunk.content, end="", flush=True)