AIyaodianzhushou/data-pipeline/processors/cleaner.py

"""
文本清洗与纠错
处理 PDF 提取文本中的格式问题
"""
import re
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class TextCleaner:
    def __init__(self):
        self.replacements = [
            (r"(?<!\n)\n(?!\n)", ""),   # 移除不当换行
            (r"\x00", ""),               # 移除空字符
            (r"([a-z])\n([a-z])", r"\1\2"),  # 英文单词断行连接
            (r"(\d)\n(\d)", r"\1\2"),    # 数字断行连接
            (r"\n{3,}", "\n\n"),         # 多个连续换行缩为两个
            (r"\t", " "),                # 制表符替换
            (r" {2,}", " "),             # 多个空格缩为一个
        ]

    def clean(self, text: str) -> str:
        result = text
        for pattern, replacement in self.replacements:
            result = re.sub(pattern, replacement, result)
        return result.strip()

    def clean_drug_entry(self, text: str) -> str:
        cleaned = self.clean(text)
        cleaned = self._fix_common_ocr_errors(cleaned)
        return cleaned

    def _fix_common_ocr_errors(self, text: str) -> str:
        ocr_fixes = {
            "咤": "啶",
            "唾": "唑",
            "茶碱": "茶碱",
        }
        for wrong, correct in ocr_fixes.items():
            text = text.replace(wrong, correct)
        return text