Codex
93 lines
2.4 KiB
Python
93 lines
2.4 KiB
Python
"""
|
|
分块策略基类
|
|
"""
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ChunkResult:
|
|
"""分块结果"""
|
|
content: str
|
|
position: int = 0
|
|
char_count: int = 0
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
if not self.char_count:
|
|
self.char_count = len(self.content)
|
|
|
|
|
|
class BaseChunker(ABC):
|
|
"""
|
|
分块策略基类
|
|
|
|
所有分块策略必须继承此类并实现 chunk 方法
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
chunk_size: int = 500,
|
|
chunk_overlap: int = 50,
|
|
separator: Optional[str] = None,
|
|
):
|
|
self.chunk_size = chunk_size
|
|
self.chunk_overlap = chunk_overlap
|
|
self.separator = separator
|
|
|
|
@abstractmethod
|
|
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
|
"""
|
|
将文本分块
|
|
|
|
Args:
|
|
text: 原始文本
|
|
metadata: 文档元数据
|
|
|
|
Returns:
|
|
分块结果列表
|
|
"""
|
|
pass
|
|
|
|
def _clean_text(self, text: str) -> str:
|
|
"""清理文本：去除多余空白"""
|
|
import re
|
|
# 合并连续空行为单个空行
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
# 去除行尾空白
|
|
text = '\n'.join(line.rstrip() for line in text.split('\n'))
|
|
return text.strip()
|
|
|
|
def _merge_small_chunks(self, chunks: List[str], min_size: int = 50) -> List[str]:
|
|
"""合并过小的分块"""
|
|
if not chunks:
|
|
return []
|
|
|
|
merged = []
|
|
buffer = ""
|
|
|
|
for chunk in chunks:
|
|
if not chunk.strip():
|
|
continue
|
|
if buffer and len(buffer) + len(chunk) <= self.chunk_size:
|
|
buffer = buffer + "\n" + chunk
|
|
elif buffer and len(buffer) < min_size:
|
|
buffer = buffer + "\n" + chunk
|
|
else:
|
|
if buffer:
|
|
merged.append(buffer)
|
|
buffer = chunk
|
|
|
|
if buffer:
|
|
# 最后一个 buffer 如果太小，合并到前一个
|
|
if merged and len(buffer) < min_size:
|
|
merged[-1] = merged[-1] + "\n" + buffer
|
|
else:
|
|
merged.append(buffer)
|
|
|
|
return merged
|