bluryar / uie-base-ONNX

UIE ONNX模型使用指南

本指南介绍如何使用和导出UIE (Universal Information Extraction) 模型的ONNX版本。

模型说明

UIE模型是一个通用信息抽取模型，基于BERT架构。该模型可以导出为ONNX格式以实现更快的推理速度。

模型输入

模型接受以下输入张量：

• input_ids : 形状为 [batch_size, sequence_length] 的整型张量
• attention_mask : 形状为 [batch_size, sequence_length] 的整型张量
• token_type_ids : 形状为 [batch_size, sequence_length] 的整型张量

其中：

• batch_size : 批处理大小，可变
• sequence_length : 序列长度，可变
• 所有输入张量的数据类型均为 int64

模型输出

模型输出以下张量：

• start_prob : 形状为 [batch_size, sequence_length] 的浮点张量，表示每个位置作为实体开始的概率
• end_prob : 形状为 [batch_size, sequence_length] 的浮点张量，表示每个位置作为实体结束的概率

使用示例

import onnxruntime as ort

tokenizer = AutoTokenizer.from_pretrained("xusenlin/uie-base")
session = ort.InferenceSession("path/to/model.onnx")


intput = "张三与B公司签订了一份合同，约定了合同金额为100万元，合同期限为一年。"
schema = ["人名", "公司", "金额", "时间"]
input_ids_tensor, attention_mask, token_type_ids, offsets_mapping = tokenizer(
    intput,
    schema[0],
    return_tensors="pt",
    return_offsets_mapping=True,
    add_special_tokens=True
)
inputs = {
    "input_ids": input_ids_tensor, # shape: [batch_size, sequence_length]
    "attention_mask": attention_mask, # shape: [batch_size, sequence_length]
    "token_type_ids": token_type_ids # shape: [batch_size, sequence_length]
}
outputs = session.run(None, inputs)
start_probs, end_probs = outputs

# use offsets_mapping and start_probs, end_probs to get the entities
# ...