lingyue-zhibao/python-services/ner-service/app/services/ner_service.py

"""
NER 服务实现

支持多种模式：
1. rule - 基于规则的简单 NER（默认，用于开发测试）
2. spacy - 使用 spaCy 模型
3. transformers - 使用 Transformers 模型
4. api - 调用外部 API（如 DeepSeek/Qwen）
"""
import re
import uuid
from typing import List, Optional
from loguru import logger

from ..config import settings
from ..models import EntityInfo, PositionInfo


class NerService:
    """NER 服务"""
    
    def __init__(self):
        self.model_type = settings.ner_model
        logger.info(f"初始化 NER 服务: model_type={self.model_type}")
    
    async def extract_entities(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        从文本中提取实体
        
        Args:
            text: 待提取的文本
            entity_types: 指定要提取的实体类型，为空则提取所有类型
            
        Returns:
            实体列表
        """
        if not text or not text.strip():
            return []
        
        if self.model_type == "rule":
            return await self._extract_by_rules(text, entity_types)
        elif self.model_type == "spacy":
            return await self._extract_by_spacy(text, entity_types)
        elif self.model_type == "transformers":
            return await self._extract_by_transformers(text, entity_types)
        elif self.model_type == "api":
            return await self._extract_by_api(text, entity_types)
        else:
            logger.warning(f"未知的模型类型: {self.model_type}，使用规则模式")
            return await self._extract_by_rules(text, entity_types)
    
    async def _extract_by_rules(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        基于规则的 NER 提取
        用于开发测试阶段，后续可替换为更高级的模型
        """
        entities = []
        
        # 规则定义
        rules = {
            "DATE": [
                # 中文日期格式
                r'(\d{4}年\d{1,2}月\d{1,2}日)',
                r'(\d{4}年\d{1,2}月)',
                r'(\d{4}-\d{1,2}-\d{1,2})',
                r'(\d{4}/\d{1,2}/\d{1,2})',
            ],
            "NUMBER": [
                # 带单位的数值
                r'(\d+\.?\d*\s*(?:万元|元|米|公里|千米|平方米|㎡|吨|kg|g|个|台|套|件|次|人|天|小时|分钟|秒|%|百分比))',
                # 百分比
                r'(\d+\.?\d*%)',
                # 纯数值（较大的数）
                r'(?<![a-zA-Z])(\d{4,}(?:\.\d+)?)(?![a-zA-Z])',
            ],
            "ORG": [
                # 机构/公司名称
                r'([\u4e00-\u9fa5]{2,10}(?:公司|集团|院|所|局|部|厅|委|会|中心|协会|学会|银行|医院|学校|大学|学院))',
                # xx省/市/县/区
                r'([\u4e00-\u9fa5]{2,6}(?:省|市|县|区|镇|乡|村)(?:人民)?(?:政府|委员会)?)',
            ],
            "LOC": [
                # 地点
                r'([\u4e00-\u9fa5]{2,6}(?:省|市|县|区|镇|乡|村|路|街|巷|号|楼|栋|单元|室))',
                # 常见地名后缀
                r'([\u4e00-\u9fa5]{2,8}(?:工业园|开发区|高新区|科技园|产业园))',
            ],
            "PERSON": [
                # 人名（简单规则：姓+名）
                r'(?:(?:张|王|李|赵|刘|陈|杨|黄|周|吴|徐|孙|马|朱|胡|郭|何|林|罗|高|郑|梁|谢|唐|许|邓|冯|韩|曹|曾|彭|萧|蔡|潘|田|董|袁|于|余|叶|蒋|杜|苏|魏|程|吕|丁|沈|任|姚|卢|傅|钟|姜|崔|谭|廖|范|汪|陆|金|石|戴|贾|韦|夏|邱|方|侯|邹|熊|孟|秦|白|江|阎|薛|尹|段|雷|黎|史|龙|陶|贺|顾|毛|郝|龚|邵|万|钱|严|赖|覃|洪|武|莫|孔)[\u4e00-\u9fa5]{1,2})(?:总|经理|主任|工程师|教授|博士|先生|女士|同志)?',
            ],
            "DEVICE": [
                # 设备名称
                r'([\u4e00-\u9fa5]{2,10}(?:设备|仪器|仪表|机器|装置|系统|探测器|传感器|检测仪|分析仪|监测仪))',
            ],
            "PROJECT": [
                # 项目名称
                r'([\u4e00-\u9fa5]{2,20}(?:项目|工程|计划|方案|课题))',
            ],
        }
        
        # 过滤实体类型
        if entity_types:
            rules = {k: v for k, v in rules.items() if k in entity_types}
        
        # 执行规则匹配
        seen_entities = set()  # 用于去重
        
        for entity_type, patterns in rules.items():
            for pattern in patterns:
                for match in re.finditer(pattern, text):
                    entity_text = match.group(1) if match.groups() else match.group(0)
                    entity_text = entity_text.strip()
                    
                    # 去重
                    entity_key = f"{entity_type}:{entity_text}"
                    if entity_key in seen_entities:
                        continue
                    seen_entities.add(entity_key)
                    
                    # 计算行号
                    line_num = text[:match.start()].count('\n') + 1
                    
                    # 获取上下文
                    context_start = max(0, match.start() - 20)
                    context_end = min(len(text), match.end() + 20)
                    context = text[context_start:context_end]
                    if context_start > 0:
                        context = "..." + context
                    if context_end < len(text):
                        context = context + "..."
                    
                    entity = EntityInfo(
                        name=entity_text,
                        type=entity_type,
                        value=entity_text,
                        position=PositionInfo(
                            char_start=match.start(),
                            char_end=match.end(),
                            line=line_num
                        ),
                        context=context,
                        confidence=0.8,  # 规则匹配默认置信度
                        temp_id=str(uuid.uuid4())[:8]
                    )
                    entities.append(entity)
        
        logger.info(f"规则 NER 提取完成: entity_count={len(entities)}")
        return entities
    
    async def _extract_by_spacy(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        使用 spaCy 进行 NER 提取
        """
        # TODO: 实现 spaCy NER
        logger.warning("spaCy NER 尚未实现，回退到规则模式")
        return await self._extract_by_rules(text, entity_types)
    
    async def _extract_by_transformers(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        使用 Transformers 模型进行 NER 提取
        """
        # TODO: 实现 Transformers NER
        logger.warning("Transformers NER 尚未实现，回退到规则模式")
        return await self._extract_by_rules(text, entity_types)
    
    async def _extract_by_api(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        调用外部 API 进行 NER 提取
        """
        # TODO: 实现 API NER（调用 DeepSeek/Qwen）
        logger.warning("API NER 尚未实现，回退到规则模式")
        return await self._extract_by_rules(text, entity_types)


# 创建单例
ner_service = NerService()