ow

2026-05-26 14:03:34 +08:00
1 parent 293a11d863
commit e60118eb53
@@ -1,24 +1,29 @@
 using CloudBuilder.Core.DatabaseAccessor.Entity;
 using CloudBuilder.Core.Service;
-using DocumentFormat.OpenXml.Wordprocessing;
-using EnvDTE80;
-using Irony.Parsing;
 using Microsoft.ML;
 using Microsoft.ML.Data;
 using Microsoft.ML.Transforms.Onnx;
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
-using System.Text;
-using System.Threading.Tasks;

 namespace CloudBuilder.AI.Service.Utility
 {
+    /// <summary>
+    /// 基于 ONNX 模型的命名实体识别器，支持提取人名及人物角色（人名+职位）。
+    /// </summary>
    public class ChineseNameExtractor
    {
        private BertChineseTokenizer tokenizer;
        private OnnxScoringEstimator pipeline;
        private PredictionEngine<NerInput, NerOutput> engine;
+
+        /// <summary>
+        /// 开启后会将每个 token 的预测标签及置信度输出到控制台，方便调试。
+        /// </summary>
+        public bool DebugMode { get; set; } = false;
+
        public ChineseNameExtractor(IApplicationService service)
        {
            InitMLContext(service);
@@ -26,54 +31,62 @@ namespace CloudBuilder.AI.Service.Utility

        private void InitMLContext(IApplicationService service)
        {
-            // 1. 初始化 MLContext
            var mlContext = new MLContext();
-
            string modelDir = service.Configuration["FileServiceSettings:OnnxDirectory"];
-            // 2. 检查必要文件
            string modelPath = Path.Combine(modelDir, "model.onnx");
            string vocabPath = Path.Combine(modelDir, "vocab.txt");

            if (!File.Exists(modelPath) || !File.Exists(vocabPath))
-            {
-                throw new ValidatedException("请确保 model.onnx 和 vocab.txt 在程序目录下。");
-            }
+                throw new Exception("请确保 model.onnx 和 vocab.txt 在指定目录下。");

-            // 3. 加载词汇表与标签映射
            tokenizer = new BertChineseTokenizer(vocabPath);

-            // 4. 定义模型输入输出的精确形状
-            //    通过 Netron (https://netron.app/) 查看模型结构后，你可以确认这些形状。
            var shapeDictionary = new Dictionary<string, int[]>
-                {
-                    { "input_ids", new[] { 1, 128 } },
-                    { "attention_mask", new[] { 1, 128 } },
-                    { "token_type_ids", new[] { 1, 128 } },
-                    { "logits", new[] { 1, 128, 29 } }   // 29 个标签
+            {
+                { "input_ids", new[] { 1, 128 } },
+                { "attention_mask", new[] { 1, 128 } },
+                { "token_type_ids", new[] { 1, 128 } },
+                { "logits", new[] { 1, 128, tokenizer.LabelCount } }
            };

-            // 5. 构建管道（使用空数据视图和 shapeDictionary）
            var emptyData = mlContext.Data.LoadFromEnumerable(new List<NerInput>());
            pipeline = mlContext.Transforms.ApplyOnnxModel(
-               modelFile: modelPath,
-               inputColumnNames: new[] { "input_ids", "attention_mask", "token_type_ids" },
-               outputColumnNames: new[] { "logits" },
-               shapeDictionary: shapeDictionary,
-               gpuDeviceId: null,        // 若要用 GPU，可改为 0 等
-               fallbackToCpu: true
-           );
+                modelFile: modelPath,
+                inputColumnNames: new[] { "input_ids", "attention_mask", "token_type_ids" },
+                outputColumnNames: new[] { "logits" },
+                shapeDictionary: shapeDictionary,
+                gpuDeviceId: null,
+                fallbackToCpu: true
+            );

-            // 6. 拟合并创建预测引擎
-            //Console.WriteLine("正在加载 ONNX 模型...");
            var transformer = pipeline.Fit(emptyData);
            engine = mlContext.Model.CreatePredictionEngine<NerInput, NerOutput>(transformer);
-            //Console.WriteLine("模型加载完成。");
        }

+        /// <summary>
+        /// 提取纯人名（标签为 B-PER / I-PER 的连续实体）。
+        /// </summary>
        public List<string> ExtractNames(string text)
        {
-            var tokenized = tokenizer.Tokenize(text);
+            return ExtractPersonEntities(text, onlyPureName: true);
+        }

+        /// <summary>
+        /// 提取人物角色（人名 + 紧接其后的职位，例如“岳堂主”、“张经理”等）。
+        /// </summary>
+        public List<string> ExtractPersonRoles(string text)
+        {
+            return ExtractPersonEntities(text, onlyPureName: false);
+        }
+
+        /// <summary>
+        /// 核心实体提取方法，支持合并人名与其后的职位标签。
+        /// </summary>
+        /// <param name="text">输入文本</param>
+        /// <param name="onlyPureName">true 只提取纯人名；false 提取人名+职位复合实体</param>
+        private List<string> ExtractPersonEntities(string text, bool onlyPureName)
+        {
+            var tokenized = tokenizer.Tokenize(text);
            var input = new NerInput
            {
                InputIds = tokenized.InputIds,
@@ -82,35 +95,26 @@ namespace CloudBuilder.AI.Service.Utility
            };

            var prediction = engine.Predict(input);
-
-            return ExtractNames(prediction, tokenized, tokenizer.IdToLabel);
-        }
-        /// <summary>
-        /// 根据模型输出的 logits 和分词信息，提取出人名（BIO 格式）。
-        /// </summary>
-        /// <param name="output">模型预测输出。</param>
-        /// <param name="tokenized">分词后的结果。</param>
-        /// <param name="idToLabel">标签ID到名称的映射数组。</param>
-        /// <returns>识别出的人名列表。</returns>
-        private List<string> ExtractNames(NerOutput output, TokenizedResult tokenized, string[] idToLabel)
-        {
-            var names = new List<string>();
-            string currentName = "";
+            var idToLabel = tokenizer.IdToLabel;
            int seqLen = 128;
            int numLabels = idToLabel.Length;

+            var entities = new List<string>();
+            string currentEntity = "";
+            bool insideEntity = false;
+            bool isPersonEntity = false;   // 当前实体是否以 B-PER 开始
+
            for (int i = 0; i < seqLen; i++)
            {
-                // 跳过填充 token（attention_mask 为 0）
-                if (tokenized.AttentionMask[i] == 0)
-                    continue;
+                if (tokenized.AttentionMask[i] == 0) continue;   // 跳过 padding

+                // 找到当前 token 的预测标签
                int startIdx = i * numLabels;
                float maxVal = float.MinValue;
                int maxIdx = 0;
                for (int j = 0; j < numLabels; j++)
                {
-                    float val = output.Logits[startIdx + j];
+                    float val = prediction.Logits[startIdx + j];
                    if (val > maxVal)
                    {
                        maxVal = val;
@@ -121,40 +125,99 @@ namespace CloudBuilder.AI.Service.Utility
                string label = idToLabel[maxIdx];
                string token = tokenized.Tokens[i];

-                if (label == "B-PER")
+                if (DebugMode)
+                    Console.WriteLine($"{token} -> {label} (置信度: {maxVal:F3})");
+
+                // 判断当前标签类型
+                bool isPerStart = label == "B-PER";
+                bool isPerInside = label == "I-PER";
+                bool isPositionStart = label == "B-POSITION";
+                bool isPositionInside = label == "I-POSITION";
+
+                // ---- 仅提取纯人名的逻辑 ----
+                if (onlyPureName)
                {
-                    if (currentName != "")
-                        names.Add(currentName);
-                    currentName = token;
+                    if (isPerStart)
+                    {
+                        if (insideEntity && currentEntity.Length > 0) entities.Add(currentEntity);
+                        currentEntity = token;
+                        insideEntity = true;
+                        isPersonEntity = true;
+                    }
+                    else if (isPerInside && insideEntity && isPersonEntity)
+                    {
+                        currentEntity += token;
+                    }
+                    else
+                    {
+                        if (insideEntity && currentEntity.Length > 0)
+                        {
+                            entities.Add(currentEntity);
+                            currentEntity = "";
+                            insideEntity = false;
+                            isPersonEntity = false;
+                        }
+                    }
+                    continue;
                }
-                else if (label == "I-PER")
+
+                // ---- 提取人物角色（人名+职位）的逻辑 ----
+                if (isPerStart)
                {
-                    currentName += token;
+                    // 无论之前是否有未结束的实体，都结束之前的，并开始一个新的人名
+                    if (insideEntity && currentEntity.Length > 0)
+                        entities.Add(currentEntity);
+
+                    currentEntity = token;
+                    insideEntity = true;
+                    isPersonEntity = true;
+                }
+                else if (isPositionStart)
+                {
+                    if (insideEntity && isPersonEntity)
+                    {
+                        // 当前正在收集人名，且遇到 B-POSITION，则合并职位
+                        currentEntity += token;
+                        // 保持 isPersonEntity = true 不变，这样后续的 I-POSITION 也能继续合并
+                    }
+                    else
+                    {
+                        // 不在人名内部，则视作新实体（提取单独的职位，如“堂主”）
+                        if (insideEntity && currentEntity.Length > 0)
+                            entities.Add(currentEntity);
+                        currentEntity = token;
+                        insideEntity = true;
+                        isPersonEntity = false;   // 标记为纯职位实体
+                    }
+                }
+                else if ((isPerInside || isPositionInside) && insideEntity)
+                {
+                    // 无论是 I-PER 还是 I-POSITION，只要在实体内部都追加
+                    currentEntity += token;
                }
                else
                {
-                    if (currentName != "")
+                    // 遇到其他标签（O 或其他 B- 开头），结束当前实体
+                    if (insideEntity && currentEntity.Length > 0)
                    {
-                        names.Add(currentName);
-                        currentName = "";
+                        entities.Add(currentEntity);
+                        currentEntity = "";
+                        insideEntity = false;
+                        isPersonEntity = false;
                    }
+                    // 注意：不处理非目标标签本身，因为它们不属于角色
                }
            }

-            if (currentName != "")
-                names.Add(currentName);
-            return names;
-        }
+            // 收尾最后一个实体
+            if (insideEntity && currentEntity.Length > 0)
+                entities.Add(currentEntity);

+            return entities;
+        }
    }

-
-    // ========== 数据模型定义 ==========
-    // 以下类的 VectorType 和 ColumnName 必须与 shapeDictionary 及 ONNX 模型定义完全一致。
-
-    /// <summary>
-    /// ONNX 模型的输入。假设 batch=1，序列长度=128，三个 int64 张量。
-    /// </summary>
+    // ==================== 数据模型 ====================
    public class NerInput
    {
        [VectorType(1, 128)]
@@ -170,78 +233,66 @@ namespace CloudBuilder.AI.Service.Utility
        public long[] TokenTypeIds { get; set; }
    }

-    /// <summary>
-    /// ONNX 模型的输出。假设 logits 形状为 [1, 128, num_labels]。
-    /// </summary>
    public class NerOutput
    {
-        [VectorType(1, 128, 29)]   // 与 shapeDictionary 一致
+        [VectorType(1, 128, 29)]   // 维度需与模型输出一致
        [ColumnName("logits")]
        public float[] Logits { get; set; }
    }

-    // ========== 分词器与词汇表加载 ==========
-
-    /// <summary>
-    /// 简易的 BERT 中文分词器，按字进行分割，适用于中文 NER 任务。
-    /// 特殊 token：[CLS] = 101, [SEP] = 102, [PAD] = 0。
-    /// </summary>
+    // ==================== 分词器 ====================
    public class BertChineseTokenizer
    {
        private readonly Dictionary<string, int> _tokenToId;
+
        /// <summary>
-        /// 标签ID到名称的映射，其顺序和内容必须与模型输出完全一致。
-        /// 此处的映射基于 bert-base-chinese-finetuned-ner 模型常见的标签体系。
+        /// 标签映射，务必与模型训练时的标签顺序一致。
        /// </summary>
        public readonly string[] IdToLabel;
+        public int LabelCount => IdToLabel.Length;

        public BertChineseTokenizer(string vocabPath, int maxSeqLength = 128)
        {
-            // 加载词汇表
            _tokenToId = new Dictionary<string, int>();
            var lines = File.ReadAllLines(vocabPath);
            for (int i = 0; i < lines.Length; i++)
                _tokenToId[lines[i]] = i;

-            // 默认标签映射，顺序务必与模型输出一致
+            // 请根据实际模型输出维度调整此数组！
+            // 下面的顺序仅为示例，必须与模型 logits 第2维一一对应。
            IdToLabel = new string[] {
-    "O",           // 0
-    "B-ORG",       // 1
-    "I-ORG",       // 2
-    "B-PER",       // 3
-    "I-PER",       // 4
-    "B-TIME",      // 5
-    "I-TIME",      // 6
-    "B-LOC",       // 7
-    "I-LOC",       // 8
-    "B-POSITION",  // 9
-    "I-POSITION",  // 10
-    "B-COMPANY",   // 11
-    "I-COMPANY",   // 12
-    "B-GAME",      // 13
-    "I-GAME",      // 14
-    "B-GOVERNMENT",// 15
-    "I-GOVERNMENT",// 16
-    "B-SCENE",     // 17
-    "I-SCENE",     // 18
-    "B-SUBJECT",   // 19
-    "I-SUBJECT",   // 20
-    "B-CREATION",  // 21
-    "I-CREATION",  // 22
-    "B-FOOD",      // 23
-    "I-FOOD",      // 24
-    "B-MOVIE",     // 25
-    "I-MOVIE",     // 26
-    "[CLS]",       // 27
-    "[SEP]"        // 28
-};
+                "O",            // 0
+                "B-ORG",        // 1
+                "I-ORG",        // 2
+                "B-PER",        // 3
+                "I-PER",        // 4
+                "B-TIME",       // 5
+                "I-TIME",       // 6
+                "B-LOC",        // 7
+                "I-LOC",        // 8
+                "B-POSITION",   // 9
+                "I-POSITION",   // 10
+                "B-COMPANY",    // 11
+                "I-COMPANY",    // 12
+                "B-GAME",       // 13
+                "I-GAME",       // 14
+                "B-GOVERNMENT", // 15
+                "I-GOVERNMENT", // 16
+                "B-SCENE",      // 17
+                "I-SCENE",      // 18
+                "B-SUBJECT",    // 19
+                "I-SUBJECT",    // 20
+                "B-CREATION",   // 21
+                "I-CREATION",   // 22
+                "B-FOOD",       // 23
+                "I-FOOD",       // 24
+                "B-MOVIE",      // 25
+                "I-MOVIE",      // 26
+                "[CLS]",        // 27   （如果模型没有这些特殊标签，请删除）
+                "[SEP]"         // 28
+            };
        }

-        /// <summary>
-        /// 对中文句子进行分词，并生成模型所需的 input_ids 等。
-        /// </summary>
-        /// <param name="text">输入的中文句子。</param>
-        /// <returns>分词后的结果，包括 token IDs 和 tokens 本身。</returns>
        public TokenizedResult Tokenize(string text)
        {
            const int maxLen = 128;
@@ -250,28 +301,27 @@ namespace CloudBuilder.AI.Service.Utility
            var tokenTypeIds = new long[maxLen];
            var tokens = new string[maxLen];

-            // 将所有字符当作 token（BERT 中文模型常用）
            var chars = text.ToCharArray();
            int charIndex = 0;

            for (int i = 0; i < maxLen; i++)
            {
-                if (i == 0)  // [CLS]
+                if (i == 0)
                {
                    inputIds[i] = 101;
                    tokens[i] = "[CLS]";
                }
-                else if (i == chars.Length + 1)  // [SEP]
+                else if (i == chars.Length + 1)
                {
                    inputIds[i] = 102;
                    tokens[i] = "[SEP]";
                }
-                else if (i > chars.Length + 1)  // padding
+                else if (i > chars.Length + 1)
                {
                    inputIds[i] = 0;
                    tokens[i] = "[PAD]";
                }
-                else  // 实际字符
+                else
                {
                    string c = chars[charIndex].ToString();
                    charIndex++;
@@ -282,15 +332,13 @@ namespace CloudBuilder.AI.Service.Utility
                    }
                    else
                    {
-                        // 如果词汇表没有这个字，用 [UNK] token
                        inputIds[i] = _tokenToId["[UNK]"];
                        tokens[i] = "[UNK]";
                    }
                }

-                // attention_mask：有效 token 为 1，padding 为 0
                attentionMask[i] = (i <= chars.Length + 1) ? 1 : 0;
-                tokenTypeIds[i] = 0; // 单句都为 0
+                tokenTypeIds[i] = 0;
            }

            return new TokenizedResult
@@ -303,9 +351,6 @@ namespace CloudBuilder.AI.Service.Utility
        }
    }

-    /// <summary>
-    /// 存储分词后的结果。
-    /// </summary>
    public class TokenizedResult
    {
        public long[] InputIds { get; set; }
@@ -313,4 +358,4 @@ namespace CloudBuilder.AI.Service.Utility
        public long[] TokenTypeIds { get; set; }
        public string[] Tokens { get; set; }
    }
-}
+}