101 lines
2.7 KiB
Python
101 lines
2.7 KiB
Python
"""Knowledge text chunking helpers.
|
|
|
|
Parser and markdown conversion logic has been moved to ``yuxi.plugins.parser``.
|
|
This module only keeps chunking-related utilities.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
from langchain_community.document_loaders import (
|
|
CSVLoader,
|
|
JSONLoader,
|
|
TextLoader,
|
|
UnstructuredHTMLLoader,
|
|
UnstructuredMarkdownLoader,
|
|
UnstructuredWordDocumentLoader,
|
|
)
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
|
|
def chunk_with_parser(file_path, params=None):
|
|
"""
|
|
使用文件解析器将文件切分成固定大小的块
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
params: 参数
|
|
"""
|
|
params = params or {}
|
|
chunk_size = int(params.get("chunk_size", 500))
|
|
chunk_overlap = int(params.get("chunk_overlap", 100))
|
|
|
|
file_type = Path(file_path).suffix.lower()
|
|
|
|
# 选择合适的加载器
|
|
if file_type in [".txt"]:
|
|
loader = TextLoader(file_path)
|
|
|
|
elif file_type in [".md"]:
|
|
loader = UnstructuredMarkdownLoader(file_path)
|
|
|
|
elif file_type in [".docx", ".doc"]:
|
|
loader = UnstructuredWordDocumentLoader(file_path)
|
|
|
|
elif file_type in [".html", ".htm"]:
|
|
loader = UnstructuredHTMLLoader(file_path)
|
|
|
|
elif file_type in [".json"]:
|
|
loader = JSONLoader(file_path, jq_schema=".")
|
|
|
|
elif file_type in [".csv"]:
|
|
loader = CSVLoader(file_path)
|
|
|
|
else:
|
|
raise ValueError(f"不支持的文件类型: {file_type}")
|
|
|
|
# 加载文档
|
|
docs = loader.load()
|
|
|
|
# 创建文本分割器
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=chunk_size,
|
|
chunk_overlap=chunk_overlap,
|
|
separators=["\n\n", "\n", ".", " ", ""],
|
|
)
|
|
|
|
# 分割文档
|
|
nodes = text_splitter.split_documents(docs)
|
|
|
|
# 添加序号信息到metadata
|
|
for i, node in enumerate(nodes):
|
|
if node.metadata is None:
|
|
node.metadata = {}
|
|
node.metadata["chunk_idx"] = i
|
|
|
|
return nodes
|
|
|
|
|
|
def chunk_text(text, params=None):
|
|
"""
|
|
将文本切分成固定大小的块
|
|
"""
|
|
params = params or {}
|
|
chunk_size = int(params.get("chunk_size", 500))
|
|
chunk_overlap = int(params.get("chunk_overlap", 100))
|
|
|
|
# 创建文本分割器
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n\n", "\n", ".", " ", ""]
|
|
)
|
|
|
|
# 分割文档
|
|
nodes = text_splitter.split_text(text)
|
|
|
|
# 添加序号信息到metadata
|
|
nodes = [{"text": node, "metadata": {"chunk_idx": i}} for i, node in enumerate(nodes)]
|
|
return nodes
|
|
|
|
|
|
def chunk(text_or_path, params=None):
|
|
raise NotImplementedError("chunk is deprecated, use chunk_with_parser or chunk_text instead")
|