""" ONNX 推理模块 提供基于 ONNX 的 TRPG 日志命名实体识别推理功能。 """ import os from typing import List, Dict, Any, Optional from pathlib import Path try: import numpy as np import onnxruntime as ort from transformers import AutoTokenizer except ImportError as e: raise ImportError( "依赖未安装。请运行: pip install onnxruntime transformers numpy" ) from e # 默认模型路径（相对于包安装位置） DEFAULT_MODEL_DIR = Path(__file__).parent.parent.parent.parent / "models" / "trpg-final" # 远程模型 URL（用于自动下载） MODEL_URL = "https://github.com/HydroRoll-Team/base-model/releases/download/v0.1.0/model.onnx" class TRPGParser: """ TRPG 日志解析器（基于 ONNX） Args: model_path: ONNX 模型路径，默认使用内置模型 tokenizer_path: tokenizer 配置路径，默认与 model_path 相同 device: 推理设备，"cpu" 或 "cuda" Examples: >>> parser = TRPGParser() >>> result = parser.parse("风雨 2024-06-08 21:44:59 剧烈的疼痛...") >>> print(result['metadata']['speaker']) '风雨' """ def __init__( self, model_path: Optional[str] = None, tokenizer_path: Optional[str] = None, device: str = "cpu", ): # 确定模型路径 if model_path is None: model_path = self._get_default_model_path() if tokenizer_path is None: tokenizer_path = Path(model_path).parent self.model_path = Path(model_path) self.tokenizer_path = Path(tokenizer_path) self.device = device # 加载模型 self._load_model() def _get_default_model_path(self) -> str: """获取默认模型路径""" # 1. 尝试相对于项目根目录 project_root = Path(__file__).parent.parent.parent.parent local_model = project_root / "models" / "trpg-final" / "model.onnx" if local_model.exists(): return str(local_model) # 2. 尝试用户数据目录 from pathlib import Path user_model_dir = Path.home() / ".cache" / "basemodel" / "models" / "trpg-final" user_model = user_model_dir / "model.onnx" if user_model.exists(): return str(user_model) # 3. 抛出错误，提示下载 raise FileNotFoundError( f"模型文件未找到。请从 {MODEL_URL} 下载模型到 {user_model_dir}\n" f"或运行: python -m basemodel.download_model" ) def _load_model(self): """加载 ONNX 模型和 Tokenizer""" # 加载 tokenizer self.tokenizer = AutoTokenizer.from_pretrained( str(self.tokenizer_path), local_files_only=True, ) # 加载 ONNX 模型 providers = ["CPUExecutionProvider"] if self.device == "cuda" and "CUDAExecutionProvider" in ort.get_available_providers(): providers.insert(0, "CUDAExecutionProvider") self.session = ort.InferenceSession( str(self.model_path), providers=providers, ) # 加载标签映射 import json config_path = self.tokenizer_path / "config.json" if config_path.exists(): with open(config_path, "r", encoding="utf-8") as f: config = json.load(f) self.id2label = {int(k): v for k, v in config.get("id2label", {}).items()} else: # 默认标签 self.id2label = { 0: "O", 1: "B-action", 2: "I-action", 3: "B-comment", 4: "I-comment", 5: "B-dialogue", 6: "I-dialogue", 7: "B-speaker", 8: "I-speaker", 9: "B-timestamp", 10: "I-timestamp", } def parse(self, text: str) -> Dict[str, Any]: """ 解析单条 TRPG 日志 Args: text: 待解析的日志文本 Returns: 包含 metadata 和 content 的字典 - metadata: speaker, timestamp - content: dialogue, action, comment 列表 Examples: >>> parser = TRPGParser() >>> result = parser.parse("风雨 2024-06-08 21:44:59 剧烈的疼痛...") >>> result['metadata']['speaker'] '风雨' """ # Tokenize inputs = self.tokenizer( text, return_tensors="np", return_offsets_mapping=True, padding="max_length", truncation=True, max_length=128, ) # 推理 outputs = self.session.run( ["logits"], { "input_ids": inputs["input_ids"].astype(np.int64), "attention_mask": inputs["attention_mask"].astype(np.int64), }, ) # 后处理 logits = outputs[0][0] predictions = np.argmax(logits, axis=-1) offsets = inputs["offset_mapping"][0] # 聚合实体 entities = self._group_entities(predictions, offsets, logits) # 构建结果 result = {"metadata": {}, "content": []} for ent in entities: if ent["start"] >= len(text) or ent["end"] > len(text): continue raw_text = text[ent["start"]: ent["end"]] clean_text = self._clean_text(raw_text, ent["type"]) if not clean_text.strip(): continue if ent["type"] in ["timestamp", "speaker"]: result["metadata"][ent["type"]] = clean_text elif ent["type"] in ["dialogue", "action", "comment"]: result["content"].append({ "type": ent["type"], "content": clean_text, "confidence": round(ent["score"], 3), }) return result def _group_entities(self, predictions, offsets, logits): """将 token 级别的预测聚合为实体""" entities = [] current = None for i in range(len(predictions)): start, end = offsets[i] if start == end: # special tokens continue pred_id = int(predictions[i]) label = self.id2label.get(pred_id, "O") if label == "O": if current: entities.append(current) current = None continue tag_type = label[2:] if len(label) > 2 else "O" if label.startswith("B-"): if current: entities.append(current) current = { "type": tag_type, "start": int(start), "end": int(end), "score": float(np.max(logits[i])), } elif label.startswith("I-") and current and current["type"] == tag_type: current["end"] = int(end) else: if current: entities.append(current) current = None if current: entities.append(current) return entities def _clean_text(self, text: str, group: str) -> str: """清理提取的文本""" import re text = text.strip() # 移除周围符号 if group == "comment": text = re.sub(r"^[（(]+|[）)]+$", "", text) elif group == "dialogue": text = re.sub(r'^[""''「」『』]+|[""""」』『』]+$', "", text) elif group == "action": text = re.sub(r"^[*＃]+|[*＃]+$", "", text) # 修复时间戳 if group == "timestamp" and text and text[0].isdigit(): if len(text) > 2 and text[2] == "-": text = "20" + text return text def parse_batch(self, texts: List[str]) -> List[Dict[str, Any]]: """ 批量解析多条日志 Args: texts: 日志文本列表 Returns: 解析结果列表 """ return [self.parse(text) for text in texts] # 便捷函数 def parse_line(text: str, model_path: Optional[str] = None) -> Dict[str, Any]: """ 解析单条日志的便捷函数 Args: text: 日志文本 model_path: 可选的模型路径 Returns: 解析结果字典 """ parser = TRPGParser(model_path=model_path) return parser.parse(text) def parse_lines(texts: List[str], model_path: Optional[str] = None) -> List[Dict[str, Any]]: """ 批量解析日志的便捷函数 Args: texts: 日志文本列表 model_path: 可选的模型路径 Returns: 解析结果列表 """ parser = TRPGParser(model_path=model_path) return parser.parse_batch(texts) __all__ = ["TRPGParser", "parse_line", "parse_lines"]