420 lines
15 KiB
Python
420 lines
15 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
预处理模块
|
||
"""
|
||
|
||
|
||
from os import utime
|
||
from pathlib import Path
|
||
from re import sub, MULTILINE
|
||
from hashlib import blake2s
|
||
from frontmatter import load
|
||
from langchain_text_splitters import (
|
||
MarkdownHeaderTextSplitter,
|
||
RecursiveCharacterTextSplitter,
|
||
)
|
||
from paramiko import AutoAddPolicy, SSHClient
|
||
from pgvector.psycopg import register_vector
|
||
from psycopg import connect
|
||
from pydantic import BaseModel, Field
|
||
from transformers import AutoTokenizer
|
||
|
||
|
||
class Chunk(BaseModel):
|
||
"""
|
||
分块数据模型
|
||
"""
|
||
|
||
index: int = Field(..., description="分块索引")
|
||
document_id: str = Field(..., description="文档唯一标识")
|
||
embedding_text: str = Field(..., description="分块嵌入文本")
|
||
return_text: str = Field(..., description="分块返回文本")
|
||
return_text_fingerprint: str = Field(..., description="分块返回文本指纹")
|
||
|
||
|
||
LOCAL_DIRECTORY = Path(__file__).parent / "vault" # 本地库文件夹
|
||
REMOTE_DIRECTORY = "/root/obsidian/vaults/MyNewVault/RawsSources" # 云端库文件夹
|
||
|
||
CHUNK_SIZE = 8192 # 分块大小，BAAI/BGE-M3 标值
|
||
|
||
|
||
def sync_vault() -> None:
|
||
"""
|
||
同步本地库和云端库
|
||
"""
|
||
# 创建 SSHClient
|
||
ssh = SSHClient()
|
||
ssh.set_missing_host_key_policy(AutoAddPolicy())
|
||
ssh.connect(
|
||
hostname="www.liubiren.cloud",
|
||
port=22,
|
||
username="root",
|
||
password="Te198752.",
|
||
)
|
||
# 创建 SFTP
|
||
sftp = ssh.open_sftp()
|
||
|
||
def get_remote_vault_metadata(folder: str) -> list[dict]:
|
||
"""
|
||
获取云端库指定文件夹下所有元数据（文档即MD文件，每个文档包含云端路径和修改时间戳）
|
||
:param folder: 文件夹
|
||
:return: [ {"remote_path":str, "modified_at":float}, ... ]
|
||
"""
|
||
metadata = []
|
||
# 获取指定文件夹下所有实体列表
|
||
entries = sftp.listdir_attr(folder)
|
||
for entry in entries:
|
||
# 若实体模式掩码为空则跳过
|
||
if (entry_mode_mask := entry.st_mode) is None:
|
||
continue
|
||
|
||
# 构建云端路径
|
||
remote_path = f"{folder}/{entry.filename}"
|
||
# 若实体为文件夹则递归获取该文件夹下所有元数据
|
||
if entry_mode_mask & 0o040000:
|
||
metadata.extend(get_remote_vault_metadata(remote_path))
|
||
|
||
# 若实体文件为MD文件则构建文档源数据并添加到列表
|
||
elif entry.filename.lower().endswith(".md"):
|
||
metadata.append(
|
||
{"remote_path": remote_path, "modified_at": entry.st_mtime or 0}
|
||
)
|
||
return metadata
|
||
|
||
# 获取云端库指定文件夹下所有元数据
|
||
remote_vault_metadata = get_remote_vault_metadata(REMOTE_DIRECTORY)
|
||
for metadata in remote_vault_metadata:
|
||
# 本地路径
|
||
local_path: Path = LOCAL_DIRECTORY / Path(metadata["remote_path"]).relative_to(
|
||
REMOTE_DIRECTORY
|
||
)
|
||
local_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
is_download_needed = True # 需要下载
|
||
modified_at = metadata["modified_at"]
|
||
# 若本地文档存在且修改时间戳与云端文档修改时间戳差异小于阈值则无需下载
|
||
if local_path.exists():
|
||
if abs(local_path.stat().st_mtime - modified_at) < 60:
|
||
is_download_needed = False
|
||
|
||
if is_download_needed:
|
||
# 下载文档
|
||
sftp.get(metadata["remote_path"], local_path.as_posix())
|
||
# 同步修改时间戳
|
||
utime(local_path, times=(modified_at, modified_at))
|
||
|
||
sftp.close()
|
||
ssh.close()
|
||
|
||
|
||
def process_documents() -> None:
|
||
"""
|
||
处理文档
|
||
"""
|
||
# 实例化标题分割器
|
||
header_splitter = MarkdownHeaderTextSplitter(
|
||
headers_to_split_on=[
|
||
("#", "h1"),
|
||
("##", "h2"),
|
||
("###", "h3"),
|
||
("####", "h4"),
|
||
("#####", "h5"),
|
||
("######", "h6"),
|
||
],
|
||
strip_headers=True, # 不保留标题
|
||
)
|
||
|
||
# 实例化分词器（已保存预训练模型，本地加载）
|
||
tokenizer = AutoTokenizer.from_pretrained(
|
||
Path(__file__).parent / "baai_bge_m3_tokenizer"
|
||
)
|
||
|
||
def count_tokens(text: str) -> int:
|
||
"""
|
||
统计词元数
|
||
:param text: 文本
|
||
:return: token 数量
|
||
"""
|
||
return len(
|
||
tokenizer.encode(text, add_special_tokens=False)
|
||
) # 首尾不添加模型专用特殊标记
|
||
|
||
with connect(
|
||
host="www.liubiren.cloud",
|
||
port=5201,
|
||
user="root",
|
||
password="PG198752",
|
||
dbname="database",
|
||
) as connection:
|
||
document_sync_records: dict[str, float] = {} # 同步记录字典
|
||
with connection.cursor() as cursor:
|
||
cursor.execute("SELECT uid, modified_at FROM document_sync_records;")
|
||
for uid, modified_at in cursor.fetchall():
|
||
document_sync_records[uid] = modified_at
|
||
|
||
for local_path in LOCAL_DIRECTORY.rglob("*.md"): # 递归遍历本地库所有文档
|
||
# 读取文档
|
||
document = load(local_path)
|
||
# 解析文档唯一标识（需Obsidian已安装NoteUIDGenerator插件，用于生成唯一标识）
|
||
document_id = str(document.metadata["uid"])
|
||
|
||
if modified_at := document_sync_records.get(document_id):
|
||
# 若本地文档修改时间戳与数据库文档修改时间戳差异小于阈值则跳过
|
||
if abs(local_path.stat().st_mtime - modified_at) < 60:
|
||
continue
|
||
|
||
# 若本地文档修改时间戳与数据库文档修改时间戳差异大于阈值则删除删除该文档所有文档块
|
||
cursor.execute(
|
||
"DELETE FROM document_chunks WHERE document_id = %s",
|
||
(document_id,),
|
||
)
|
||
connection.commit()
|
||
|
||
# 文档标题
|
||
document_title = local_path.stem
|
||
|
||
index = 0
|
||
chunks: list[Chunk] = []
|
||
# 先根据标题分割文档，若分块词元大小超过阈值则使用递归字符分割器进一步分割
|
||
for chunk in header_splitter.split_text(document.content):
|
||
# --- 构建分块标题 ---
|
||
|
||
headers = []
|
||
# 若h1标题为空则使用文档标题
|
||
if not (h1 := chunk.metadata.get("h1")):
|
||
headers.append(document_title.strip())
|
||
else:
|
||
headers.append(h1.strip())
|
||
if h2 := chunk.metadata.get("h2"):
|
||
headers.append(h2.strip())
|
||
if h3 := chunk.metadata.get("h3"):
|
||
headers.append(h3.strip())
|
||
if h4 := chunk.metadata.get("h4"):
|
||
headers.append(h4.strip())
|
||
if h5 := chunk.metadata.get("h5"):
|
||
headers.append(h5.strip())
|
||
if h6 := chunk.metadata.get("h6"):
|
||
headers.append(h6.strip())
|
||
header = f"标题层级：{" > ".join(headers)}\n\n"
|
||
|
||
# --- 构建分块内容 ---
|
||
# 若分块内容为空则跳过
|
||
if not (content := chunk.page_content.strip()):
|
||
continue
|
||
|
||
# 删除行尾空格
|
||
content = sub(r"[ \t]+$", "", content, flags=MULTILINE)
|
||
# 将全角空格替换为半角空格
|
||
content = content.replace(" ", " ")
|
||
# 压缩空格
|
||
content = sub(r" +", " ", content)
|
||
# 压缩换行符
|
||
content = sub(r"\n{3,}", r"\n\n", content)
|
||
# 删除分割线标记
|
||
content = content.replace("---", "")
|
||
|
||
# 返回文本
|
||
return_text = header + content
|
||
|
||
# 删除斜体标记
|
||
content = sub(r"\*(.*?)\*", r"\1", content)
|
||
# 删除加粗标记
|
||
content = sub(r"\*\*(.*?)\*\*", r"\1", content)
|
||
# 删除代码块标记
|
||
content = sub(r"`(.*?)`", r"\1", content)
|
||
# 删除文档链接标记
|
||
content = sub(r"\[\[(.*?)\]\]", r"\1", content)
|
||
# 删除图片标记（直接删除）
|
||
content = sub(r"!\[.*?\]\(.*?\)", "", content)
|
||
# 删除普通连接标记（删除URL，保留文本）
|
||
content = sub(r"\[(.*?)\]\(.*?\)", r"\1", content)
|
||
|
||
# 初始化分块内容列表
|
||
contents = [content]
|
||
# 若文本词元数超过限制则使用递归字符分割器进一步分割分块内容
|
||
if count_tokens(return_text) > CHUNK_SIZE:
|
||
# 实例化递归字符分割器
|
||
recursive_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=CHUNK_SIZE - count_tokens(header),
|
||
chunk_overlap=60, # 分块重叠
|
||
length_function=count_tokens,
|
||
separators=["\n\n", "\n", "。", "？", "！", ".", "?", "!", ""],
|
||
)
|
||
contents = recursive_splitter.split_text(content)
|
||
|
||
for content in contents:
|
||
# 若内容为空则跳过
|
||
if not content.strip():
|
||
continue
|
||
|
||
chunks.append(
|
||
Chunk(
|
||
index=index,
|
||
document_id=document_id,
|
||
embedding_text=header + content,
|
||
return_text=return_text,
|
||
return_text_fingerprint=blake2s(
|
||
return_text.encode(), digest_size=16
|
||
).hexdigest(),
|
||
)
|
||
)
|
||
index += 1
|
||
|
||
for chunk in chunks:
|
||
print(chunk.return_text)
|
||
print("\n")
|
||
exit()
|
||
|
||
|
||
def main():
|
||
# 同步本地库和云端库
|
||
sync_vault()
|
||
# 处理文档
|
||
process_documents()
|
||
|
||
|
||
main()
|
||
|
||
|
||
exit()
|
||
|
||
|
||
register_vector(db_connection)
|
||
|
||
|
||
def sync_documents1() -> None:
|
||
"""
|
||
同步文档
|
||
:return: None
|
||
"""
|
||
# 获取在线文档字典
|
||
online_documents = get_online_documents()
|
||
# 获取离线文档字典
|
||
offline_documents = get_offline_documents()
|
||
|
||
# 待删除文档唯一标识集合
|
||
pending_delete_documents = set(offline_documents.keys() - online_documents.keys())
|
||
# 待下载文档唯一标识字典
|
||
pending_download_documents: Dict[str, int] = {}
|
||
for document_id, document_updated_at in online_documents.items():
|
||
# 若在线文档唯一标识于离线文档字典不存在则将该文档唯一标识添加至待下载文档唯一标识字典
|
||
if document_id not in offline_documents:
|
||
pending_download_documents[document_id] = document_updated_at
|
||
continue
|
||
|
||
# 若在线文档唯一标识于离线文档字典存在且更新时间戳不一致则将该文档唯一标识添加至待删除/下载文档唯一标识集合
|
||
if (
|
||
document_id in offline_documents
|
||
and document_updated_at != offline_documents[document_id]
|
||
):
|
||
pending_delete_documents.add(document_id)
|
||
pending_download_documents[document_id] = document_updated_at
|
||
continue
|
||
|
||
# 批量删除待删除文档
|
||
if pending_delete_documents:
|
||
with db_connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
DELETE FROM documents
|
||
WHERE id = ANY(%s);
|
||
""",
|
||
(list(pending_delete_documents),),
|
||
)
|
||
db_connection.commit()
|
||
|
||
for document_id, document_updated_at in pending_download_documents.items():
|
||
# 下载文档
|
||
document_content = download_document(document_id)
|
||
print(document_content)
|
||
exit()
|
||
|
||
with db_connection.cursor() as cursor:
|
||
# 写入文档表
|
||
cursor.execute(
|
||
"""
|
||
INSERT INTO documents (id, updated_at)
|
||
VALUES (%s, %s)
|
||
""",
|
||
(document_id, document_updated_at),
|
||
)
|
||
# 写入分块表
|
||
cursor.executemany(
|
||
"""
|
||
INSERT INTO chunks (
|
||
document_id, idx,
|
||
heading_chunk_h1, heading_chunk_h2, heading_chunk_h3, heading_chunk_h4, heading_chunk_h5, heading_chunk_heading, heading_chunk_content,
|
||
content_chunk_content, content_chunk_size,
|
||
embedding_text, embedding_vector
|
||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);
|
||
""",
|
||
[
|
||
(
|
||
chunk.document_id,
|
||
chunk.idx,
|
||
chunk.heading_chunk_h1,
|
||
chunk.heading_chunk_h2,
|
||
chunk.heading_chunk_h3,
|
||
chunk.heading_chunk_h4,
|
||
chunk.heading_chunk_h5,
|
||
chunk.heading_chunk_heading,
|
||
chunk.heading_chunk_content,
|
||
chunk.content_chunk_content,
|
||
chunk.content_chunk_size,
|
||
chunk.embedding_text,
|
||
embedding_vector,
|
||
)
|
||
for chunk, embedding_vector in zip(
|
||
chunks,
|
||
embedder.embed_documents_sync(
|
||
embedding_texts
|
||
).embeddings, # 生成词嵌向量
|
||
)
|
||
], # 构建记录
|
||
)
|
||
db_connection.commit()
|
||
|
||
|
||
def query(
|
||
question: str,
|
||
) -> List[Tuple[float, str, str]]:
|
||
"""
|
||
查询
|
||
:param question: 问题
|
||
:return: 相似度最高的前 k 个分块（每个分块包含相似度、文档 ID、标题、内容）
|
||
"""
|
||
# 生成词嵌向量
|
||
embedding_vector = embedder.embed_query_sync(query=question).embeddings[0]
|
||
|
||
with db_connection.cursor() as cursor:
|
||
cursor.execute(
|
||
"""
|
||
SELECT
|
||
1 - (embedding_vector <=> %s::vector) AS similarity_score,
|
||
heading_chunk_heading,
|
||
heading_chunk_content
|
||
FROM chunks
|
||
WHERE 1 - (embedding_vector <=> %s::vector) >= %s
|
||
ORDER BY embedding_vector <=> %s::vector
|
||
LIMIT %s;
|
||
""",
|
||
(
|
||
embedding_vector,
|
||
embedding_vector,
|
||
0.45,
|
||
embedding_vector,
|
||
5,
|
||
),
|
||
)
|
||
rows = cursor.fetchall()
|
||
return [
|
||
(similarity_score, heading_chunk_heading, heading_chunk_content)
|
||
for similarity_score, heading_chunk_heading, heading_chunk_content in rows
|
||
]
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 同步文档
|
||
# sync_documents()
|
||
print(query("需求评审有哪些流程"))
|