---
language: 
  - "zh"
tags:
  - "information-extraction"
  - "token-classification"
  - "uie"
  - "bert"
  - "onnx"
library_name: "onnxruntime"
pipeline_tag: "feature-extraction"
---

# 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]` 的浮点张量，表示每个位置作为实体结束的概率

## 使用示例

```python
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
# ...
```