88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
文本分块器 - 将长文本分割成小块
|
|||
|
|
"""
|
|||
|
|
import re
|
|||
|
|
from config import CHUNK_SIZE, CHUNK_OVERLAP
|
|||
|
|
|
|||
|
|
def split_text(text, chunk_size=None, chunk_overlap=None):
|
|||
|
|
"""
|
|||
|
|
将文本分割成小块
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
text: 要分割的文本
|
|||
|
|
chunk_size: 每块的最大字符数
|
|||
|
|
chunk_overlap: 块之间的重叠字符数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
文本块列表
|
|||
|
|
"""
|
|||
|
|
if chunk_size is None:
|
|||
|
|
chunk_size = CHUNK_SIZE
|
|||
|
|
if chunk_overlap is None:
|
|||
|
|
chunk_overlap = CHUNK_OVERLAP
|
|||
|
|
|
|||
|
|
if not text or not text.strip():
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 清理文本
|
|||
|
|
text = text.strip()
|
|||
|
|
text = re.sub(r'\n{3,}', '\n\n', text) # 多个换行合并
|
|||
|
|
text = re.sub(r' {2,}', ' ', text) # 多个空格合并
|
|||
|
|
|
|||
|
|
# 按段落分割
|
|||
|
|
paragraphs = re.split(r'\n\n+', text)
|
|||
|
|
|
|||
|
|
chunks = []
|
|||
|
|
current_chunk = ""
|
|||
|
|
|
|||
|
|
for para in paragraphs:
|
|||
|
|
para = para.strip()
|
|||
|
|
if not para:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 如果当前段落本身就超过chunk_size，需要进一步分割
|
|||
|
|
if len(para) > chunk_size:
|
|||
|
|
# 先保存当前chunk
|
|||
|
|
if current_chunk:
|
|||
|
|
chunks.append(current_chunk.strip())
|
|||
|
|
current_chunk = ""
|
|||
|
|
|
|||
|
|
# 按句子分割长段落
|
|||
|
|
sentences = re.split(r'([。！？.!?])', para)
|
|||
|
|
temp_chunk = ""
|
|||
|
|
|
|||
|
|
for i in range(0, len(sentences), 2):
|
|||
|
|
sentence = sentences[i]
|
|||
|
|
if i + 1 < len(sentences):
|
|||
|
|
sentence += sentences[i + 1]
|
|||
|
|
|
|||
|
|
if len(temp_chunk) + len(sentence) <= chunk_size:
|
|||
|
|
temp_chunk += sentence
|
|||
|
|
else:
|
|||
|
|
if temp_chunk:
|
|||
|
|
chunks.append(temp_chunk.strip())
|
|||
|
|
temp_chunk = sentence[-chunk_overlap:] + sentence if len(sentence) > chunk_overlap else sentence
|
|||
|
|
|
|||
|
|
if temp_chunk:
|
|||
|
|
current_chunk = temp_chunk
|
|||
|
|
else:
|
|||
|
|
# 检查是否可以添加到当前chunk
|
|||
|
|
if len(current_chunk) + len(para) + 1 <= chunk_size:
|
|||
|
|
current_chunk += ("\n" if current_chunk else "") + para
|
|||
|
|
else:
|
|||
|
|
# 保存当前chunk，开始新的
|
|||
|
|
if current_chunk:
|
|||
|
|
chunks.append(current_chunk.strip())
|
|||
|
|
# 保留重叠部分
|
|||
|
|
if len(current_chunk) > chunk_overlap:
|
|||
|
|
current_chunk = current_chunk[-chunk_overlap:] + "\n" + para
|
|||
|
|
else:
|
|||
|
|
current_chunk = para
|
|||
|
|
|
|||
|
|
# 保存最后一个chunk
|
|||
|
|
if current_chunk.strip():
|
|||
|
|
chunks.append(current_chunk.strip())
|
|||
|
|
|
|||
|
|
return chunks
|