317 lines
11 KiB
C#
317 lines
11 KiB
C#
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; }
|
||
}
|
||
}
|