Browse Source

feat: 添加 UniversalNER 模型支持

新增功能:
- 支持 zeffmuks/universal-ner Ollama 模型
- 自动检测模型类型并切换提取策略
- UniversalNER 使用专用 Prompt 格式: "文本. 实体类型"
- 为每种实体类型分别调用模型提取

使用方式:
NER_MODEL=ollama OLLAMA_MODEL=zeffmuks/universal-ner \
  uvicorn app.main:app --host 0.0.0.0 --port 8001

模型对比:
- universal-ner: NER 专用,速度快,输出简洁
- qwen2.5:7b: 通用 LLM,中文强,输出详细
何文松 1 tháng trước cách đây
mục cha
commit
68d3d79433

+ 8 - 19
python-services/ner-service/.env.example

@@ -3,10 +3,9 @@
 # ============================================
 # NER 模型配置
 # ============================================
-# 可选值: rule / ollama / spacy / transformers / api
+# 可选值: rule / ollama
 # - rule: 基于规则的简单 NER(开发测试用,速度快但准确率低)
 # - ollama: 使用本地 Ollama LLM(推荐生产环境,准确率高)
-# - api: 使用远程 API(如百炼、DeepSeek)
 NER_MODEL=ollama
 
 # ============================================
@@ -15,11 +14,13 @@ NER_MODEL=ollama
 # Ollama 服务地址
 OLLAMA_URL=http://localhost:11434
 
-# 使用的模型(推荐中文 NER)
-# - qwen2.5:7b(推荐,中文能力最强)
-# - qwen2.5:14b(更强,需要更多显存)
-# - llama3.1:8b(英文较好)
-OLLAMA_MODEL=qwen2.5:7b
+# 使用的模型(推荐)
+# 通用 LLM 模式:
+#   - qwen2.5:7b(推荐中文,能力最强)
+#   - qwen2.5:14b(更强,需要更多显存)
+# UniversalNER 专用模式:
+#   - zeffmuks/universal-ner(NER 专用模型,速度快)
+OLLAMA_MODEL=zeffmuks/universal-ner
 
 # 请求超时时间(秒)
 OLLAMA_TIMEOUT=120
@@ -33,18 +34,6 @@ CHUNK_SIZE=2000
 # 分块重叠字符数(避免实体被截断)
 CHUNK_OVERLAP=200
 
-# ============================================
-# API 配置(当 NER_MODEL=api 时使用)
-# ============================================
-# API 基础 URL
-# API_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1
-
-# API 密钥
-# API_KEY=your-api-key
-
-# API 模型
-# API_MODEL=qwen-plus
-
 # ============================================
 # 日志配置
 # ============================================

+ 9 - 5
python-services/ner-service/README.md

@@ -136,11 +136,15 @@ Content-Type: application/json
 
 ### 推荐模型
 
-| 模型 | 大小 | 中文能力 | 推荐场景 |
-|------|------|----------|----------|
-| qwen2.5:7b | ~4.7GB | ⭐⭐⭐⭐⭐ | 推荐,中文最强 |
-| qwen2.5:14b | ~9GB | ⭐⭐⭐⭐⭐ | 显存充足时 |
-| llama3.1:8b | ~4.7GB | ⭐⭐⭐ | 英文为主 |
+| 模型 | 大小 | 中文能力 | 特点 | 推荐场景 |
+|------|------|----------|------|----------|
+| **zeffmuks/universal-ner** | ~4.2GB | ⭐⭐⭐ | NER 专用,速度快 | **首选** |
+| qwen2.5:7b | ~4.7GB | ⭐⭐⭐⭐⭐ | 通用 LLM,中文强 | 需要复杂理解 |
+| qwen2.5:14b | ~9GB | ⭐⭐⭐⭐⭐ | 更强的通用能力 | 显存充足时 |
+
+**UniversalNER vs 通用 LLM**:
+- UniversalNER: 专为 NER 优化,输出简洁,速度快
+- Qwen: 通用能力强,可输出详细信息(如位置),但较慢
 
 ## 测试
 

+ 4 - 0
python-services/ner-service/app/config.py

@@ -29,6 +29,10 @@ class Settings(BaseSettings):
     ollama_model: str = "qwen2.5:7b"  # 推荐中文 NER 使用 qwen2.5
     ollama_timeout: int = 120  # 秒
     
+    # UniversalNER 专用配置(当 ollama_model 包含 'universal-ner' 时自动启用)
+    # 模型名: zeffmuks/universal-ner
+    universal_ner_model: str = "zeffmuks/universal-ner"
+    
     # 文本分块配置(用于长文本处理)
     chunk_size: int = 2000  # 每个分块的最大字符数
     chunk_overlap: int = 200  # 分块重叠字符数

+ 154 - 2
python-services/ner-service/app/services/ollama_service.py

@@ -22,7 +22,11 @@ class OllamaService:
         self.timeout = settings.ollama_timeout
         self.chunk_size = settings.chunk_size
         self.chunk_overlap = settings.chunk_overlap
-        logger.info(f"初始化 Ollama 服务: url={self.base_url}, model={self.model}")
+        
+        # 检测是否使用 UniversalNER
+        self.is_universal_ner = "universal-ner" in self.model.lower()
+        
+        logger.info(f"初始化 Ollama 服务: url={self.base_url}, model={self.model}, universal_ner={self.is_universal_ner}")
     
     def _split_text(self, text: str) -> List[Dict[str, Any]]:
         """
@@ -200,10 +204,25 @@ class OllamaService:
         使用 Ollama LLM 提取实体
         
         支持长文本自动分块处理
+        自动检测是否使用 UniversalNER 并切换提取策略
         """
         if not text or not text.strip():
             return []
         
+        # 根据模型类型选择提取策略
+        if self.is_universal_ner:
+            return await self._extract_with_universal_ner(text, entity_types)
+        else:
+            return await self._extract_with_general_llm(text, entity_types)
+    
+    async def _extract_with_general_llm(
+        self, 
+        text: str, 
+        entity_types: Optional[List[str]] = None
+    ) -> List[EntityInfo]:
+        """
+        使用通用 LLM(如 Qwen)提取实体
+        """
         # 分割长文本
         chunks = self._split_text(text)
         
@@ -235,9 +254,142 @@ class OllamaService:
             
             logger.info(f"分块 {i+1} 提取实体: {len(entities)} 个")
         
-        logger.info(f"Ollama NER 提取完成: 总实体数={len(all_entities)}")
+        logger.info(f"通用 LLM NER 提取完成: 总实体数={len(all_entities)}")
+        return all_entities
+    
+    async def _extract_with_universal_ner(
+        self, 
+        text: str, 
+        entity_types: Optional[List[str]] = None
+    ) -> List[EntityInfo]:
+        """
+        使用 UniversalNER 模型提取实体
+        
+        UniversalNER 的 Prompt 格式: "文本内容. 实体类型英文名"
+        返回格式: ["实体1", "实体2", ...]
+        """
+        # 实体类型映射(中文类型 -> UniversalNER 英文类型)
+        type_mapping = {
+            "PERSON": ["person", "people", "human"],
+            "ORG": ["organization", "company", "institution"],
+            "LOC": ["location", "place", "address"],
+            "DATE": ["date", "time"],
+            "NUMBER": ["number", "quantity", "measurement"],
+            "DEVICE": ["device", "equipment", "instrument"],
+            "PROJECT": ["project", "program"],
+            "METHOD": ["method", "standard", "specification"],
+        }
+        
+        types_to_extract = entity_types or list(type_mapping.keys())
+        
+        # 分割长文本
+        chunks = self._split_text(text)
+        
+        all_entities = []
+        seen_entities = set()  # 用于去重
+        
+        for i, chunk in enumerate(chunks):
+            chunk_text = chunk["text"]
+            chunk_start = chunk["start_pos"]
+            
+            logger.info(f"UniversalNER 处理分块 {i+1}/{len(chunks)}: 长度={len(chunk_text)}")
+            
+            # 对每种实体类型分别提取
+            for entity_type in types_to_extract:
+                if entity_type not in type_mapping:
+                    continue
+                
+                # 使用第一个英文类型名
+                english_type = type_mapping[entity_type][0]
+                
+                # UniversalNER 的 Prompt 格式
+                prompt = f"{chunk_text} {english_type}"
+                
+                # 调用 Ollama
+                response = await self._call_ollama(prompt)
+                
+                if not response:
+                    continue
+                
+                # 解析 UniversalNER 响应(返回格式如: ["实体1", "实体2"])
+                entities = self._parse_universal_ner_response(
+                    response, entity_type, chunk_text, chunk_start
+                )
+                
+                # 去重
+                for entity in entities:
+                    entity_key = f"{entity.type}:{entity.name}"
+                    if entity_key not in seen_entities:
+                        seen_entities.add(entity_key)
+                        all_entities.append(entity)
+            
+            logger.info(f"分块 {i+1} UniversalNER 提取实体: {len([e for e in all_entities if e not in seen_entities])} 个")
+        
+        logger.info(f"UniversalNER 提取完成: 总实体数={len(all_entities)}")
         return all_entities
     
+    def _parse_universal_ner_response(
+        self, 
+        response: str, 
+        entity_type: str,
+        original_text: str,
+        chunk_start_pos: int = 0
+    ) -> List[EntityInfo]:
+        """
+        解析 UniversalNER 的响应
+        
+        UniversalNER 返回格式: ["实体1", "实体2", ...]
+        """
+        entities = []
+        
+        try:
+            # 清理响应,提取 JSON 数组
+            response = response.strip()
+            
+            # 尝试找到 JSON 数组
+            json_match = re.search(r'\[[\s\S]*?\]', response)
+            if not json_match:
+                logger.debug(f"UniversalNER 响应中未找到数组: {response[:100]}")
+                return entities
+            
+            json_str = json_match.group()
+            entity_names = json.loads(json_str)
+            
+            if not isinstance(entity_names, list):
+                return entities
+            
+            for name in entity_names:
+                if not isinstance(name, str) or len(name) < 2:
+                    continue
+                
+                name = name.strip()
+                
+                # 在原文中查找位置
+                pos = original_text.find(name)
+                char_start = pos + chunk_start_pos if pos >= 0 else 0
+                char_end = char_start + len(name) if pos >= 0 else 0
+                
+                entity = EntityInfo(
+                    name=name,
+                    type=entity_type,
+                    value=name,
+                    position=PositionInfo(
+                        char_start=char_start,
+                        char_end=char_end,
+                        line=1
+                    ),
+                    confidence=0.85,  # UniversalNER 置信度
+                    temp_id=str(uuid.uuid4())[:8]
+                )
+                entities.append(entity)
+                
+        except json.JSONDecodeError as e:
+            logger.debug(f"UniversalNER JSON 解析失败: {e}, response={response[:100]}")
+        except Exception as e:
+            logger.error(f"解析 UniversalNER 响应失败: {e}")
+        
+        return entities
+    
     async def check_health(self) -> bool:
         """
         检查 Ollama 服务是否可用