Views
No views yet
input_ids: 形状为 [batch_size, sequence_length] 的整型张量attention_mask: 形状为 [batch_size, sequence_length] 的整型张量token_type_ids: 形状为 [batch_size, sequence_length] 的整型张量batch_size: 批处理大小，可变sequence_length: 序列长度，可变int64start_prob: 形状为 [batch_size, sequence_length] 的浮点张量，表示每个位置作为实体开始的概率end_prob: 形状为 [batch_size, sequence_length] 的浮点张量，表示每个位置作为实体结束的概率1import onnxruntime as ort
2
3tokenizer = AutoTokenizer.from_pretrained("xusenlin/uie-base")
4session = ort.InferenceSession("path/to/model.onnx")
5
6
7intput = "张三与B公司签订了一份合同，约定了合同金额为100万元，合同期限为一年。"
8schema = ["人名", "公司", "金额", "时间"]
9input_ids_tensor, attention_mask, token_type_ids, offsets_mapping = tokenizer(
10 intput,
11 schema[0],
12 return_tensors="pt",
13 return_offsets_mapping=True,
14 add_special_tokens=True
15)
16inputs = {
17 "input_ids": input_ids_tensor, # shape: [batch_size, sequence_length]
18 "attention_mask": attention_mask, # shape: [batch_size, sequence_length]
19 "token_type_ids": token_type_ids # shape: [batch_size, sequence_length]
20}
21outputs = session.run(None, inputs)
22start_probs, end_probs = outputs
23
24# use offsets_mapping and start_probs, end_probs to get the entities
25# ...