Codex
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
"""
|
|
按句子分块策略
|
|
|
|
按句号/问号/感叹号等句子边界分割文本，
|
|
然后将小句子合并到不超过 chunk_size 的分块中。
|
|
参考 Dify 的 sentence 分块模式。
|
|
"""
|
|
import logging
|
|
import re
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from .base import BaseChunker, ChunkResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 句子分隔符正则（中英文句号/问号/感叹号/分号）
|
|
SENTENCE_PATTERN = re.compile(
|
|
r'(?<=[。！？；.!?;])\s*'
|
|
)
|
|
|
|
|
|
class SentenceChunker(BaseChunker):
|
|
"""
|
|
按句子分块器
|
|
|
|
先按句子边界分割文本，再将相邻句子合并为不超过 chunk_size 的分块。
|
|
保证每个分块都是完整句子的组合，不会在句子中间截断。
|
|
"""
|
|
|
|
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
|
"""按句子分块"""
|
|
if not text or not text.strip():
|
|
return []
|
|
|
|
text = self._clean_text(text)
|
|
metadata = metadata or {}
|
|
|
|
# 按句子边界分割
|
|
sentences = SENTENCE_PATTERN.split(text)
|
|
sentences = [s.strip() for s in sentences if s.strip()]
|
|
|
|
if not sentences:
|
|
return [ChunkResult(content=text, position=0, metadata={**metadata})]
|
|
|
|
# 合并句子为分块（不超过 chunk_size）
|
|
chunks = []
|
|
current = ""
|
|
position = 0
|
|
|
|
for sentence in sentences:
|
|
# 如果单个句子就超过 chunk_size，强制作为独立分块
|
|
if len(sentence) > self.chunk_size:
|
|
if current:
|
|
chunks.append(current)
|
|
current = ""
|
|
chunks.append(sentence)
|
|
continue
|
|
|
|
test = current + sentence if not current else current + " " + sentence
|
|
if len(test) <= self.chunk_size:
|
|
current = test
|
|
else:
|
|
if current:
|
|
chunks.append(current)
|
|
current = sentence
|
|
|
|
if current:
|
|
chunks.append(current)
|
|
|
|
# 合并过小的分块
|
|
chunks = self._merge_small_chunks(chunks)
|
|
|
|
# 添加重叠
|
|
if self.chunk_overlap > 0 and len(chunks) > 1:
|
|
chunks = self._add_sentence_overlap(chunks)
|
|
|
|
# 构建结果
|
|
results = []
|
|
for i, content in enumerate(chunks):
|
|
if content.strip():
|
|
results.append(ChunkResult(
|
|
content=content.strip(),
|
|
position=i,
|
|
metadata={**metadata},
|
|
))
|
|
|
|
return results
|
|
|
|
def _add_sentence_overlap(self, chunks: List[str]) -> List[str]:
|
|
"""为相邻分块添加句子级重叠"""
|
|
result = [chunks[0]]
|
|
for i in range(1, len(chunks)):
|
|
prev = chunks[i - 1]
|
|
# 从前一个分块取最后一个句子作为重叠
|
|
prev_sentences = SENTENCE_PATTERN.split(prev)
|
|
prev_sentences = [s.strip() for s in prev_sentences if s.strip()]
|
|
if prev_sentences:
|
|
overlap = prev_sentences[-1]
|
|
if len(overlap) <= self.chunk_overlap:
|
|
combined = overlap + " " + chunks[i]
|
|
if len(combined) <= self.chunk_size * 1.2:
|
|
result.append(combined)
|
|
continue
|
|
result.append(chunks[i])
|
|
return result
|