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"([a-z])\n([a-z])", r"\1\2"),  # 英文单词断行连接
            (r"(\d)\n(\d)", r"\1\2"),    # 数字断行连接
            (r"(?<!\n)\n(?!\n)", ""),   # 移除不当换行（段落内的单换行）
            (r"\x00", ""),               # 移除空字符
            (r"\n{3,}", "\n\n"),         # 多个连续换行缩为两个
            (r"\t", " "),                # 制表符替换
            (r" {2,}", " "),             # 多个空格缩为一个
        ]

    def clean(self, text: str) -> str:
        result = self._normalize_width(text)
        result = self._normalize_units(result)
        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 纠错字典，覆盖常见字形混淆和编码问题"""
        ocr_fixes = {
            # ---- 已有 ----
            "咤": "啶",
            "唾": "唑",
            "茶碱": "茶碱",
            # ---- 化学药常见 OCR 错误 ----
            "甘": "苷",          # 苦杏仁甘→苦杏仁苷
            "甙": "苷",          # 异体字规范
            "溜": "馏",          # 蒸溜→蒸馏
            "腊": "蜡",          # 腊→蜡
            "酿": "酶",          # 酿→酶 (OCR)
            "容": "密",          # 容封→密封
            "陚": "酯",          # 化学式
            "肔": "肽",
            "唼": "哌",
            "噻昉": "噻吩",
            "涣": "溴",
            "氮卓": "氮䓬",
            "畄": "留",
            "貭": "质",
            "千燥": "干燥",
            "堿": "碱",
        }
        for wrong, correct in ocr_fixes.items():
            text = text.replace(wrong, correct)
        return text

    def _normalize_width(self, text: str) -> str:
        """全角/半角字符规范化，确保标点和数字一致性"""
        # 全角 ASCII 标点 → 半角（保留中文逗号句号）
        width_map = {
            "！": "!", "？": "?", "＂": '"', "＇": "'",
            "：": ":", "；": ";", "（": "(", "）": ")",
            "【": "[", "】": "]", "～": "~",
            # 全角数字 → ASCII
            "０": "0", "１": "1", "２": "2", "３": "3", "４": "4",
            "５": "5", "６": "6", "７": "7", "８": "8", "９": "9",
            # 全角大写字母 → ASCII
            "Ａ": "A", "Ｂ": "B", "Ｃ": "C", "Ｄ": "D", "Ｅ": "E",
            "Ｆ": "F", "Ｇ": "G", "Ｈ": "H", "Ｉ": "I", "Ｊ": "J",
            "Ｋ": "K", "Ｌ": "L", "Ｍ": "M", "Ｎ": "N", "Ｏ": "O",
            "Ｐ": "P", "Ｑ": "Q", "Ｒ": "R", "Ｓ": "S", "Ｔ": "T",
            "Ｕ": "U", "Ｖ": "V", "Ｗ": "W", "Ｘ": "X", "Ｙ": "Y", "Ｚ": "Z",
            "ａ": "a", "ｂ": "b", "ｃ": "c", "ｄ": "d", "ｅ": "e",
            "ｆ": "f", "ｇ": "g", "ｈ": "h", "ｉ": "i", "ｊ": "j",
            "ｋ": "k", "ｌ": "l", "ｍ": "m", "ｎ": "n", "ｏ": "o",
            "ｐ": "p", "ｑ": "q", "ｒ": "r", "ｓ": "s", "ｔ": "t",
            "ｕ": "u", "ｖ": "v", "ｗ": "w", "ｘ": "x", "ｙ": "y", "ｚ": "z",
        }
        for full, half in width_map.items():
            text = text.replace(full, half)
        return text

    def _normalize_units(self, text: str) -> str:
        """药品计量单位规范化"""
        unit_fixes = [
            (r'\bug\b', 'μg'),        # ug → μg
            (r'\buL\b', 'μL'),        # uL → μL
            (r'\bum\b', 'μm'),        # um → μm
            (r'(?i)\bml\b', 'mL'),    # ML/ml/Ml → mL
            (r'(?i)\bul\b', 'μL'),    # ul → μL
            (r'(\d+)h\b', r'\1小时'),  # 24h → 24小时
            (r'(\d+)min\b', r'\1分钟'),  # 30min → 30分钟
        ]
        for pattern, replacement in unit_fixes:
            text = re.sub(pattern, replacement, text)
        return text