Codex
160 lines
5.2 KiB
Python
160 lines
5.2 KiB
Python
"""
|
|
文档预处理/清洗服务
|
|
|
|
参考 Dify 的 DatasetProcessRule，支持可配置的文本清洗规则。
|
|
在文本提取之后、分块之前执行。
|
|
"""
|
|
import logging
|
|
import re
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 默认预处理规则
|
|
DEFAULT_PROCESS_RULES: Dict[str, Any] = {
|
|
"pre_processing_rules": [
|
|
{"id": "remove_extra_spaces", "enabled": True},
|
|
{"id": "remove_urls_emails", "enabled": False},
|
|
{"id": "remove_html_tags", "enabled": False},
|
|
{"id": "remove_consecutive_newlines", "enabled": True},
|
|
{"id": "remove_trailing_whitespace", "enabled": True},
|
|
],
|
|
}
|
|
|
|
|
|
class CleaningService:
|
|
"""
|
|
文本清洗服务
|
|
|
|
支持的清洗规则：
|
|
- remove_extra_spaces: 合并连续空格为单个空格
|
|
- remove_urls_emails: 移除 URL 和邮箱地址
|
|
- remove_html_tags: 移除 HTML 标签
|
|
- remove_consecutive_newlines: 合并连续空行（3+）为双空行
|
|
- remove_trailing_whitespace: 去除行尾空白
|
|
"""
|
|
|
|
# 规则处理器映射
|
|
RULE_PROCESSORS = {
|
|
"remove_extra_spaces": "_remove_extra_spaces",
|
|
"remove_urls_emails": "_remove_urls_emails",
|
|
"remove_html_tags": "_remove_html_tags",
|
|
"remove_consecutive_newlines": "_remove_consecutive_newlines",
|
|
"remove_trailing_whitespace": "_remove_trailing_whitespace",
|
|
}
|
|
|
|
@classmethod
|
|
def clean(cls, text: str, process_rules: Optional[Dict[str, Any]] = None) -> str:
|
|
"""
|
|
根据预处理规则清洗文本
|
|
|
|
Args:
|
|
text: 原始文本
|
|
process_rules: 预处理规则配置，为 None 则使用默认规则
|
|
|
|
Returns:
|
|
清洗后的文本
|
|
"""
|
|
if not text:
|
|
return text
|
|
|
|
rules = process_rules or DEFAULT_PROCESS_RULES
|
|
pre_rules = rules.get("pre_processing_rules", [])
|
|
|
|
original_length = len(text)
|
|
|
|
for rule in pre_rules:
|
|
rule_id = rule.get("id", "")
|
|
enabled = rule.get("enabled", False)
|
|
|
|
if not enabled:
|
|
continue
|
|
|
|
processor_name = cls.RULE_PROCESSORS.get(rule_id)
|
|
if not processor_name:
|
|
logger.warning(f"未知的预处理规则: {rule_id}")
|
|
continue
|
|
|
|
processor = getattr(cls, processor_name, None)
|
|
if processor:
|
|
text = processor(text)
|
|
|
|
cleaned_length = len(text)
|
|
if original_length != cleaned_length:
|
|
logger.info(
|
|
f"文本清洗完成: {original_length} -> {cleaned_length} 字符 "
|
|
f"(减少 {original_length - cleaned_length})"
|
|
)
|
|
|
|
return text.strip()
|
|
|
|
@staticmethod
|
|
def _remove_extra_spaces(text: str) -> str:
|
|
"""合并连续空格为单个空格（保留换行符）"""
|
|
# 只处理同一行内的连续空格，不影响换行
|
|
lines = text.split('\n')
|
|
cleaned_lines = []
|
|
for line in lines:
|
|
cleaned_lines.append(re.sub(r'[ \t]+', ' ', line))
|
|
return '\n'.join(cleaned_lines)
|
|
|
|
@staticmethod
|
|
def _remove_urls_emails(text: str) -> str:
|
|
"""移除 URL 和邮箱地址"""
|
|
# 移除 URL
|
|
text = re.sub(
|
|
r'https?://[^\s<>"{}|\\^`\[\]]+',
|
|
'',
|
|
text,
|
|
)
|
|
# 移除邮箱
|
|
text = re.sub(
|
|
r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',
|
|
'',
|
|
text,
|
|
)
|
|
return text
|
|
|
|
@staticmethod
|
|
def _remove_html_tags(text: str) -> str:
|
|
"""移除 HTML 标签，保留文本内容"""
|
|
# 移除 script 和 style 标签及其内容
|
|
text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.DOTALL | re.IGNORECASE)
|
|
text = re.sub(r'<style[^>]*>.*?</style>', '', text, flags=re.DOTALL | re.IGNORECASE)
|
|
# 移除所有 HTML 标签
|
|
text = re.sub(r'<[^>]+>', '', text)
|
|
# 解码常见 HTML 实体
|
|
text = text.replace(' ', ' ')
|
|
text = text.replace('<', '<')
|
|
text = text.replace('>', '>')
|
|
text = text.replace('&', '&')
|
|
text = text.replace('"', '"')
|
|
text = text.replace(''', "'")
|
|
return text
|
|
|
|
@staticmethod
|
|
def _remove_consecutive_newlines(text: str) -> str:
|
|
"""合并连续空行（3个以上换行）为双换行"""
|
|
return re.sub(r'\n{3,}', '\n\n', text)
|
|
|
|
@staticmethod
|
|
def _remove_trailing_whitespace(text: str) -> str:
|
|
"""去除每行行尾空白"""
|
|
return '\n'.join(line.rstrip() for line in text.split('\n'))
|
|
|
|
@classmethod
|
|
def get_default_rules(cls) -> Dict[str, Any]:
|
|
"""获取默认预处理规则"""
|
|
return DEFAULT_PROCESS_RULES.copy()
|
|
|
|
@classmethod
|
|
def get_available_rules(cls) -> List[Dict[str, str]]:
|
|
"""获取所有可用的预处理规则"""
|
|
return [
|
|
{"id": "remove_extra_spaces", "label": "合并连续空格"},
|
|
{"id": "remove_urls_emails", "label": "移除 URL 和邮箱"},
|
|
{"id": "remove_html_tags", "label": "移除 HTML 标签"},
|
|
{"id": "remove_consecutive_newlines", "label": "合并连续空行"},
|
|
{"id": "remove_trailing_whitespace", "label": "去除行尾空白"},
|
|
]
|