This commit is contained in:
owenchen committed 2026-05-22 11:47:05 +08:00
1 parent 0c79fb3519
commit 16bd851b33
4 files changed
+443

No files matched your search

@@ -0,0 +1,316 @@
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.Linq;
using System.Text;
using System.Threading.Tasks;
namespace CloudBuilder.AI.Service.Utility
{
public class ChineseNameExtractor
{
private BertChineseTokenizer tokenizer;
private OnnxScoringEstimator pipeline;
private PredictionEngine<NerInput, NerOutput> engine;
public ChineseNameExtractor(IApplicationService service)
{
InitMLContext(service);
}
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 在程序目录下。");
}
// 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 个标签
};
// 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
);
// 6. 拟合并创建预测引擎
//Console.WriteLine("正在加载 ONNX 模型...");
var transformer = pipeline.Fit(emptyData);
engine = mlContext.Model.CreatePredictionEngine<NerInput, NerOutput>(transformer);
//Console.WriteLine("模型加载完成。");
}
public List<string> ExtractNames(string text)
{
var tokenized = tokenizer.Tokenize(text);
var input = new NerInput
{
InputIds = tokenized.InputIds,
AttentionMask = tokenized.AttentionMask,
TokenTypeIds = tokenized.TokenTypeIds
};
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 = "";
int seqLen = 128;
int numLabels = idToLabel.Length;
for (int i = 0; i < seqLen; i++)
{
// 跳过填充 token(attention_mask 为 0)
if (tokenized.AttentionMask[i] == 0)
continue;
int startIdx = i * numLabels;
float maxVal = float.MinValue;
int maxIdx = 0;
for (int j = 0; j < numLabels; j++)
{
float val = output.Logits[startIdx + j];
if (val > maxVal)
{
maxVal = val;
maxIdx = j;
}
}
string label = idToLabel[maxIdx];
string token = tokenized.Tokens[i];
if (label == "B-PER")
{
if (currentName != "")
names.Add(currentName);
currentName = token;
}
else if (label == "I-PER")
{
currentName += token;
}
else
{
if (currentName != "")
{
names.Add(currentName);
currentName = "";
}
}
}
if (currentName != "")
names.Add(currentName);
return names;
}
}
// ========== 数据模型定义 ==========
// 以下类的 VectorType 和 ColumnName 必须与 shapeDictionary 及 ONNX 模型定义完全一致。
/// <summary>
/// ONNX 模型的输入。假设 batch=1,序列长度=128,三个 int64 张量。
/// </summary>
public class NerInput
{
[VectorType(1, 128)]
[ColumnName("input_ids")]
public long[] InputIds { get; set; }
[VectorType(1, 128)]
[ColumnName("attention_mask")]
public long[] AttentionMask { get; set; }
[VectorType(1, 128)]
[ColumnName("token_type_ids")]
public long[] TokenTypeIds { get; set; }
}
/// <summary>
/// ONNX 模型的输出。假设 logits 形状为 [1, 128, num_labels]。
/// </summary>
public class NerOutput
{
[VectorType(1, 128, 29)] // 与 shapeDictionary 一致
[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 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;
// 默认标签映射,顺序务必与模型输出一致
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
};
}
/// <summary>
/// 对中文句子进行分词,并生成模型所需的 input_ids 等。
/// </summary>
/// <param name="text">输入的中文句子。</param>
/// <returns>分词后的结果,包括 token IDs 和 tokens 本身。</returns>
public TokenizedResult Tokenize(string text)
{
const int maxLen = 128;
var inputIds = new long[maxLen];
var attentionMask = new long[maxLen];
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]
{
inputIds[i] = 101;
tokens[i] = "[CLS]";
}
else if (i == chars.Length + 1) // [SEP]
{
inputIds[i] = 102;
tokens[i] = "[SEP]";
}
else if (i > chars.Length + 1) // padding
{
inputIds[i] = 0;
tokens[i] = "[PAD]";
}
else // 实际字符
{
string c = chars[charIndex].ToString();
charIndex++;
if (_tokenToId.TryGetValue(c, out int id))
{
inputIds[i] = id;
tokens[i] = c;
}
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
}
return new TokenizedResult
{
InputIds = inputIds,
AttentionMask = attentionMask,
TokenTypeIds = tokenTypeIds,
Tokens = tokens
};
}
}
/// <summary>
/// 存储分词后的结果。
/// </summary>
public class TokenizedResult
{
public long[] InputIds { get; set; }
public long[] AttentionMask { get; set; }
public long[] TokenTypeIds { get; set; }
public string[] Tokens { get; set; }
}
}