├── data
└── dgre
│ └── raw_data
│ └── process.py
├── evaluate.py
├── generic.py
├── README.md
├── doccano.py
├── export_model.py
├── doccano.md
├── model.py
├── finetune.py
├── convert.py
├── uie_predictor.py
└── utils.py
/data/dgre/raw_data/process.py:
--------------------------------------------------------------------------------
1 | import json
2 | import codecs
3 |
4 | data_path = "train.json"
5 |
6 | with codecs.open(data_path, 'r', encoding="utf-8") as fp:
7 | data = fp.readlines()
8 | output_file = open("../mid_data/doccnao_train.json', 'w', encoding='utf-8')
9 |
10 | for did, d in enumerate(data):
11 | d = eval(d)
12 | tmp = {}
13 | tmp["id"] = d['ID']
14 | tmp['text'] = d['text']
15 | tmp['relations'] = []
16 | tmp['entities'] = []
17 | ent_id = 0
18 | for rel_id,spo in enumerate(d['spo_list']):
19 | rel_tmp = {}
20 | ent_tmp = {}
21 | rel_tmp['id'] = rel_id
22 | ent_tmp['id'] = ent_id
23 | h = spo['h']
24 | ent_tmp['start_offset'] = h['pos'][0]
25 | ent_tmp['end_offset'] = h['pos'][1]
26 | ent_tmp['label'] = "主体"
27 | if ent_tmp not in tmp['entities']:
28 | tmp['entities'].append(ent_tmp)
29 | from_id = ent_id
30 | ent_id += 1
31 | else:
32 | ind = tmp['entities'].index(ent_tmp)
33 | from_id = tmp['entities'][ind]['id']
34 | ent_id = len(tmp['entities']) + 1
35 |
36 | t = spo['t']
37 | ent_tmp = {}
38 | ent_tmp['id'] = ent_id
39 | ent_tmp['start_offset'] = t['pos'][0]
40 | ent_tmp['end_offset'] = t['pos'][1]
41 | ent_tmp['label'] = "客体"
42 | if ent_tmp not in tmp['entities']:
43 | tmp['entities'].append(ent_tmp)
44 | to_id = ent_id
45 | ent_id += 1
46 | else:
47 | ind = tmp['entities'].index(ent_tmp)
48 | to_id = tmp['entities'][ind]['id']
49 | ent_id = len(tmp['entities']) + 1
50 |
51 | rel_tmp['from_id'] = from_id
52 | rel_tmp['to_id'] = to_id
53 | rel_tmp['type'] = spo['relation']
54 |
55 | tmp['relations'].append(rel_tmp)
56 | output_file.write(json.dumps(tmp, ensure_ascii=False) + "\n")
--------------------------------------------------------------------------------
/evaluate.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) 2022 Heiheiyoyo. All Rights Reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | from model import UIE
16 | import argparse
17 | import torch
18 | from utils import SpanEvaluator, IEDataset, logger, tqdm
19 | from transformers import BertTokenizerFast
20 | from torch.utils.data import DataLoader
21 |
22 |
23 | @torch.no_grad()
24 | def evaluate(model, metric, data_loader, device='gpu', loss_fn=None, show_bar=True):
25 | """
26 | Given a dataset, it evals model and computes the metric.
27 | Args:
28 | model(obj:`torch.nn.Module`): A model to classify texts.
29 | metric(obj:`Metric`): The evaluation metric.
30 | data_loader(obj:`torch.utils.data.DataLoader`): The dataset loader which generates batches.
31 | """
32 | return_loss = False
33 | if loss_fn is not None:
34 | return_loss = True
35 | model.eval()
36 | metric.reset()
37 | loss_list = []
38 | loss_sum = 0
39 | loss_num = 0
40 | if show_bar:
41 | data_loader = tqdm(
42 | data_loader, desc="Evaluating", unit='batch')
43 | for batch in data_loader:
44 | input_ids, token_type_ids, att_mask, start_ids, end_ids = batch
45 | if device == 'gpu':
46 | input_ids = input_ids.cuda()
47 | token_type_ids = token_type_ids.cuda()
48 | att_mask = att_mask.cuda()
49 | outputs = model(input_ids=input_ids,
50 | token_type_ids=token_type_ids,
51 | attention_mask=att_mask)
52 | start_prob, end_prob = outputs[0], outputs[1]
53 | if device == 'gpu':
54 | start_prob, end_prob = start_prob.cpu(), end_prob.cpu()
55 | start_ids = start_ids.type(torch.float32)
56 | end_ids = end_ids.type(torch.float32)
57 |
58 | if return_loss:
59 | # Calculate loss
60 | loss_start = loss_fn(start_prob, start_ids)
61 | loss_end = loss_fn(end_prob, end_ids)
62 | loss = (loss_start + loss_end) / 2.0
63 | loss = float(loss)
64 | loss_list.append(loss)
65 | loss_sum += loss
66 | loss_num += 1
67 | if show_bar:
68 | data_loader.set_postfix(
69 | {
70 | 'dev loss': f'{loss_sum / loss_num:.5f}'
71 | }
72 | )
73 |
74 | # Calcalate metric
75 | num_correct, num_infer, num_label = metric.compute(start_prob, end_prob,
76 | start_ids, end_ids)
77 | metric.update(num_correct, num_infer, num_label)
78 | precision, recall, f1 = metric.accumulate()
79 | model.train()
80 | if return_loss:
81 | loss_avg = sum(loss_list) / len(loss_list)
82 | return loss_avg, precision, recall, f1
83 | else:
84 | return precision, recall, f1
85 |
86 |
87 | def do_eval():
88 | tokenizer = BertTokenizerFast.from_pretrained(args.model_path)
89 | model = UIE.from_pretrained(args.model_path)
90 | if args.device == 'gpu':
91 | model = model.cuda()
92 |
93 | test_ds = IEDataset(args.test_path, tokenizer=tokenizer,
94 | max_seq_len=args.max_seq_len)
95 |
96 | test_data_loader = DataLoader(
97 | test_ds, batch_size=args.batch_size, shuffle=False)
98 | metric = SpanEvaluator()
99 | precision, recall, f1 = evaluate(
100 | model, metric, test_data_loader, args.device)
101 | logger.info("Evaluation precision: %.5f, recall: %.5f, F1: %.5f" %
102 | (precision, recall, f1))
103 |
104 |
105 | if __name__ == "__main__":
106 | # yapf: disable
107 | parser = argparse.ArgumentParser()
108 |
109 | parser.add_argument("-m", "--model_path", type=str, required=True,
110 | help="The path of saved model that you want to load.")
111 | parser.add_argument("-t", "--test_path", type=str, required=True,
112 | help="The path of test set.")
113 | parser.add_argument("-b", "--batch_size", type=int, default=16,
114 | help="Batch size per GPU/CPU for training.")
115 | parser.add_argument("--max_seq_len", type=int, default=512,
116 | help="The maximum total input sequence length after tokenization.")
117 | parser.add_argument("-D", '--device', choices=['cpu', 'gpu'], default="gpu",
118 | help="Select which device to run model, defaults to gpu.")
119 |
120 | args = parser.parse_args()
121 |
122 | do_eval()
123 |
--------------------------------------------------------------------------------
/generic.py:
--------------------------------------------------------------------------------
1 | import inspect
2 | from collections import OrderedDict, UserDict
3 | from collections.abc import MutableMapping
4 | from contextlib import ExitStack
5 | from dataclasses import fields
6 | from enum import Enum
7 | from typing import Any, ContextManager, List, Tuple
8 |
9 | import numpy as np
10 |
11 |
12 | def is_tensor(x):
13 | """
14 | Tests if `x` is a `torch.Tensor`, `tf.Tensor`, `jaxlib.xla_extension.DeviceArray` or `np.ndarray`.
15 | """
16 | if is_torch_fx_proxy(x):
17 | return True
18 | if is_torch_available():
19 | import torch
20 |
21 | if isinstance(x, torch.Tensor):
22 | return True
23 | if is_tf_available():
24 | import tensorflow as tf
25 |
26 | if isinstance(x, tf.Tensor):
27 | return True
28 |
29 | if is_flax_available():
30 | import jax.numpy as jnp
31 | from jax.core import Tracer
32 |
33 | if isinstance(x, (jnp.ndarray, Tracer)):
34 | return True
35 |
36 | return isinstance(x, np.ndarray)
37 |
38 | class ModelOutput(OrderedDict):
39 | """
40 | Base class for all model outputs as dataclass. Has a `__getitem__` that allows indexing by integer or slice (like a
41 | tuple) or strings (like a dictionary) that will ignore the `None` attributes. Otherwise behaves like a regular
42 | python dictionary.
43 |
44 | You can't unpack a `ModelOutput` directly. Use the [`~utils.ModelOutput.to_tuple`] method to convert it to a tuple
45 | before.
46 |
47 | """
48 |
49 | def __post_init__(self):
50 | class_fields = fields(self)
51 |
52 | # Safety and consistency checks
53 | if not len(class_fields):
54 | raise ValueError(f"{self.__class__.__name__} has no fields.")
55 | if not all(field.default is None for field in class_fields[1:]):
56 | raise ValueError(f"{self.__class__.__name__} should not have more than one required field.")
57 |
58 | first_field = getattr(self, class_fields[0].name)
59 | other_fields_are_none = all(getattr(self, field.name) is None for field in class_fields[1:])
60 |
61 | if other_fields_are_none and not is_tensor(first_field):
62 | if isinstance(first_field, dict):
63 | iterator = first_field.items()
64 | first_field_iterator = True
65 | else:
66 | try:
67 | iterator = iter(first_field)
68 | first_field_iterator = True
69 | except TypeError:
70 | first_field_iterator = False
71 |
72 | # if we provided an iterator as first field and the iterator is a (key, value) iterator
73 | # set the associated fields
74 | if first_field_iterator:
75 | for element in iterator:
76 | if (
77 | not isinstance(element, (list, tuple))
78 | or not len(element) == 2
79 | or not isinstance(element[0], str)
80 | ):
81 | break
82 | setattr(self, element[0], element[1])
83 | if element[1] is not None:
84 | self[element[0]] = element[1]
85 | elif first_field is not None:
86 | self[class_fields[0].name] = first_field
87 | else:
88 | for field in class_fields:
89 | v = getattr(self, field.name)
90 | if v is not None:
91 | self[field.name] = v
92 |
93 | def __delitem__(self, *args, **kwargs):
94 | raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
95 |
96 | def setdefault(self, *args, **kwargs):
97 | raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
98 |
99 | def pop(self, *args, **kwargs):
100 | raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
101 |
102 | def update(self, *args, **kwargs):
103 | raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
104 |
105 | def __getitem__(self, k):
106 | if isinstance(k, str):
107 | inner_dict = {k: v for (k, v) in self.items()}
108 | return inner_dict[k]
109 | else:
110 | return self.to_tuple()[k]
111 |
112 | def __setattr__(self, name, value):
113 | if name in self.keys() and value is not None:
114 | # Don't call self.__setitem__ to avoid recursion errors
115 | super().__setitem__(name, value)
116 | super().__setattr__(name, value)
117 |
118 | def __setitem__(self, key, value):
119 | # Will raise a KeyException if needed
120 | super().__setitem__(key, value)
121 | # Don't call self.__setattr__ to avoid recursion errors
122 | super().__setattr__(key, value)
123 |
124 | def to_tuple(self) -> Tuple[Any]:
125 | """
126 | Convert self to a tuple containing all the attributes/keys that are not `None`.
127 | """
128 | return tuple(self[k] for k in self.keys())
129 |
130 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # pytorch_uie_re
2 | 基于pytorch的百度UIE关系抽取，源代码来源：[here](https://github.com/heiheiyoyo/uie_pytorch)。
3 |
4 | 百度UIE通用信息抽取的样例一般都是使用doccano标注数据，这里介绍如何使用通用数据集，利用UIE进行微调。
5 |
6 | # 依赖
7 |
8 | ```
9 | torch>=1.7.0
10 | transformers==4.20.0
11 | colorlog
12 | colorama
13 | ```
14 |
15 | # 步骤
16 |
17 | - 1、数据放在data下面，比如已经放置的dgre（[工业知识图谱关系抽取-高端装备制造知识图谱自动化构建 竞赛 - DataFountain](https://www.datafountain.cn/competitions/584)），raw_data下面是原始的数据，新建一个process.py，将数据处理成类似mid_data下面的数据，```python process.py```即:
18 |
19 | ```python
20 | {"id": "AT0004", "text": "617号汽车故障报告故障现象一辆吉利车装用MR479发动机，行驶里程为23709公里，驾驶员反映该车在行驶中无异响，但在起步和换挡过程中车身有抖动现象，并且听到离合器内部有异响。", "relations": [{"id": 0, "from_id": 0, "to_id": 1, "type": "部件故障"}, {"id": 1, "from_id": 2, "to_id": 3, "type": "部件故障"}], "entities": [{"id": 0, "start_offset": 80, "end_offset": 83, "label": "主体"}, {"id": 1, "start_offset": 86, "end_offset": 88, "label": "客体"}, {"id": 2, "start_offset": 68, "end_offset": 70, "label": "主体"}, {"id": 3, "start_offset": 71, "end_offset": 73, "label": "客体"}]}
21 | ```
22 |
23 | - 2、将mid_data下面的数据使用doccano.py转换成final_data下的数据，具体指令是：
24 |
25 | ```python
26 | python doccano.py \
27 | --doccano_file ./data/dgre/mid_data/doccano_train.json \
28 | --task_type "ext" \ # ext表示抽取任务
29 | --splits 0.9 0.1 0.0 \ # 训练、验证、测试数据的比例。训练，不对数据进行切分，因此将第一位设置为1.0
30 | --save_dir ./data/dgre/final_data/ \
31 | --negative_ratio 1 # 生成负样本的比率
32 | ```
33 |
34 | 最终会在final_data下生成train.txt、dev.txt。
35 |
36 | - 3、将paddle版本的模型转换为pytorch版的模型：
37 |
38 | ```python
39 | python convert.py --input_model=uie-base --output_model=uie_base_pytorch --no_validate_output
40 | ```
41 |
42 | 其中input_model可选的模型可参考convert.py里面。output_model是我们要保存的模型路径，下面会用到。之后我们可以测试下转换的效果：
43 |
44 | ```python
45 | from uie_predictor import UIEPredictor
46 | from pprint import pprint
47 |
48 | schema = ['时间', '选手', '赛事名称'] # Define the schema for entity extraction
49 | ie = UIEPredictor('./uie_base_pytorch', schema=schema)
50 | pprint(ie("2月8日上午北京冬奥会自由式滑雪女子大跳台决赛中中国选手谷爱凌以188.25分获得金牌！")) # Better print results using pprint
51 | ```
52 |
53 | - 4、开始微调：
54 |
55 | ```python
56 | python finetune.py \
57 | --train_path "./data/dgre/final_data/train.txt" \
58 | --dev_path "./data/dgre/final_data/dev.txt" \
59 | --save_dir "./checkpoint/dgre" \
60 | --learning_rate 1e-5 \
61 | --batch_size 8 \
62 | --max_seq_len 512 \
63 | --num_epochs 3 \
64 | --model "uie_base_pytorch" \
65 | --seed 1000 \
66 | --logging_steps 464 \
67 | --valid_steps 464 \
68 | --device "gpu" \
69 | --max_model_num 1
70 | ```
71 |
72 | 训练完成后，会在同目录下生成checkpoint/dgre/model_best/。
73 |
74 | - 5、进行验证：
75 |
76 | ```python
77 | python evaluate.py \
78 | --model_path "./checkpoint/dgre/model_best" \
79 | --test_path "./data/dgre/final_data/dev.txt" \
80 | --batch_size 16 \
81 | --max_seq_len 512
82 | ```
83 |
84 | - 6、使用训练好的模型进行预测：
85 |
86 | ```python
87 | from uie_predictor import UIEPredictor
88 | from pprint import pprint
89 |
90 | schema = [{"主体":["部件故障", "性能故障", "检测工具", "组成"]}, "客体"] # Define the schema for entity extraction
91 | ie = UIEPredictor('./checkpoint/dgre/model_best', schema=schema)
92 | text = "分析诊断首先用故障诊断仪读取故障码为12，其含义是电控系统正常。考虑到电喷发动机控制是由怠速马达来实现的，所以先拆下怠速马达，发现其阀头上粘有大量胶质油污。用化油器清洗剂清洗后，装车试验，故障依旧。接着清洗喷油咀，故障仍未排除。最后把节气门体拆下来清洗。在操作过程中发现：一根插在节气门体下部真空管上的胶管已断裂，造成节气门后腔与大气相通，影响怠速运转稳定。这条胶管应该是连接在节气门进气管和气门室盖排气孔之间特制的丁字胶管的一部分，但该车没有使用特制的丁字胶管，它用一条直通胶管将节气门进气管和气门室盖排气孔连起来。维修方案把节气门体清洗干净后装车，再用一条专用特制的丁字形的三通胶管把节气门进气管、气门室盖排气孔和节气门体下部真空管接好，然后启动发动机，加速收油，发动机转速平稳下降"
93 | res = ie(text)
94 | pprint(res) # Better print results using pprint
95 |
96 | """
97 | [{'主体': [{'end': 164,
98 | 'probability': 0.44656947,
99 | 'relations': {'性能故障': [{'end': 169,
100 | 'probability': 0.9079101,
101 | 'start': 167,
102 | 'text': '相通'}],
103 | '部件故障': [{'end': 169,
104 | 'probability': 0.91306436,
105 | 'start': 167,
106 | 'text': '相通'}]},
107 | 'start': 159,
108 | 'text': '节气门后腔'},
109 | {'end': 153,
110 | 'probability': 0.92779523,
111 | 'relations': {'性能故障': [{'end': 156,
112 | 'probability': 0.9924409,
113 | 'start': 154,
114 | 'text': '断裂'}],
115 | '部件故障': [{'end': 156,
116 | 'probability': 0.9910217,
117 | 'start': 154,
118 | 'text': '断裂'}]},
119 | 'start': 151,
120 | 'text': '胶管'},
121 | {'end': 68,
122 | 'probability': 0.32421115,
123 | 'relations': {'性能故障': [{'end': 77,
124 | 'probability': 0.6647082,
125 | 'start': 69,
126 | 'text': '粘有大量胶质油污'}],
127 | '部件故障': [{'end': 77,
128 | 'probability': 0.8483382,
129 | 'start': 69,
130 | 'text': '粘有大量胶质油污'}]},
131 | 'start': 66,
132 | 'text': '阀头'}],
133 | '客体': [{'end': 77, 'probability': 0.5508156, 'start': 69, 'text': '粘有大量胶质油污'},
134 | {'end': 156, 'probability': 0.9614242, 'start': 154, 'text': '断裂'},
135 | {'end': 169,
136 | 'probability': 0.56304514,
137 | 'start': 164,
138 | 'text': '与大气相通'}]}]
139 | """
140 | ```
141 |
142 | 会发现，关系抽取会有关系的重复，可以多训练几个epoch看看。
143 |
144 | # 补充
145 |
146 | - 标签名最好是使用中文。
147 | - 可使用不同大小的模型进行训练和推理，以达到精度和速度的平衡。
148 |
--------------------------------------------------------------------------------
/doccano.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) 2022 Heiheiyoyo. All Rights Reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | import os
16 | import time
17 | import argparse
18 | import json
19 | from decimal import Decimal
20 | import numpy as np
21 |
22 | from utils import set_seed, convert_ext_examples, convert_cls_examples, logger
23 |
24 |
25 | def do_convert():
26 | set_seed(args.seed)
27 |
28 | tic_time = time.time()
29 | if not os.path.exists(args.doccano_file):
30 | raise ValueError("Please input the correct path of doccano file.")
31 |
32 | if not os.path.exists(args.save_dir):
33 | os.makedirs(args.save_dir)
34 |
35 | if len(args.splits) != 0 and len(args.splits) != 3:
36 | raise ValueError("Only []/ len(splits)==3 accepted for splits.")
37 |
38 | def _check_sum(splits):
39 | return Decimal(str(splits[0])) + Decimal(str(splits[1])) + Decimal(
40 | str(splits[2])) == Decimal("1")
41 |
42 | if len(args.splits) == 3 and not _check_sum(args.splits):
43 | raise ValueError(
44 | "Please set correct splits, sum of elements in splits should be equal to 1."
45 | )
46 |
47 | with open(args.doccano_file, "r", encoding="utf-8") as f:
48 | raw_examples = f.readlines()
49 |
50 | def _create_ext_examples(examples,
51 | negative_ratio=0,
52 | shuffle=False,
53 | is_train=True):
54 | entities, relations = convert_ext_examples(
55 | examples, negative_ratio, is_train=is_train)
56 | examples = entities + relations
57 | if shuffle:
58 | indexes = np.random.permutation(len(examples))
59 | examples = [examples[i] for i in indexes]
60 | return examples
61 |
62 | def _create_cls_examples(examples, prompt_prefix, options, shuffle=False):
63 | examples = convert_cls_examples(examples, prompt_prefix, options)
64 | if shuffle:
65 | indexes = np.random.permutation(len(examples))
66 | examples = [examples[i] for i in indexes]
67 | return examples
68 |
69 | def _save_examples(save_dir, file_name, examples):
70 | count = 0
71 | save_path = os.path.join(save_dir, file_name)
72 | if not examples:
73 | logger.info("Skip saving %d examples to %s." % (0, save_path))
74 | return
75 | with open(save_path, "w", encoding="utf-8") as f:
76 | for example in examples:
77 | f.write(json.dumps(example, ensure_ascii=False) + "\n")
78 | count += 1
79 | logger.info("Save %d examples to %s." % (count, save_path))
80 |
81 | if len(args.splits) == 0:
82 | if args.task_type == "ext":
83 | examples = _create_ext_examples(raw_examples, args.negative_ratio,
84 | args.is_shuffle)
85 | else:
86 | examples = _create_cls_examples(raw_examples, args.prompt_prefix,
87 | args.options, args.is_shuffle)
88 | _save_examples(args.save_dir, "train.txt", examples)
89 | else:
90 | if args.is_shuffle:
91 | indexes = np.random.permutation(len(raw_examples))
92 | raw_examples = [raw_examples[i] for i in indexes]
93 |
94 | i1, i2, _ = args.splits
95 | p1 = int(len(raw_examples) * i1)
96 | p2 = int(len(raw_examples) * (i1 + i2))
97 |
98 | if args.task_type == "ext":
99 | train_examples = _create_ext_examples(
100 | raw_examples[:p1], args.negative_ratio, args.is_shuffle)
101 | dev_examples = _create_ext_examples(
102 | raw_examples[p1:p2], -1, is_train=False)
103 | test_examples = _create_ext_examples(
104 | raw_examples[p2:], -1, is_train=False)
105 | else:
106 | train_examples = _create_cls_examples(
107 | raw_examples[:p1], args.prompt_prefix, args.options)
108 | dev_examples = _create_cls_examples(
109 | raw_examples[p1:p2], args.prompt_prefix, args.options)
110 | test_examples = _create_cls_examples(
111 | raw_examples[p2:], args.prompt_prefix, args.options)
112 |
113 | _save_examples(args.save_dir, "train.txt", train_examples)
114 | _save_examples(args.save_dir, "dev.txt", dev_examples)
115 | _save_examples(args.save_dir, "test.txt", test_examples)
116 |
117 | logger.info('Finished! It takes %.2f seconds' % (time.time() - tic_time))
118 |
119 |
120 | if __name__ == "__main__":
121 | # yapf: disable
122 | parser = argparse.ArgumentParser()
123 |
124 | parser.add_argument("-d", "--doccano_file", default="./data/doccano.json",
125 | type=str, help="The doccano file exported from doccano platform.")
126 | parser.add_argument("-s", "--save_dir", default="./data",
127 | type=str, help="The path of data that you wanna save.")
128 | parser.add_argument("--negative_ratio", default=5, type=int,
129 | help="Used only for the extraction task, the ratio of positive and negative samples, number of negtive samples = negative_ratio * number of positive samples")
130 | parser.add_argument("--splits", default=[0.8, 0.1, 0.1], type=float, nargs="*",
131 | help="The ratio of samples in datasets. [0.6, 0.2, 0.2] means 60%% samples used for training, 20%% for evaluation and 20%% for test.")
132 | parser.add_argument("--task_type", choices=['ext', 'cls'], default="ext", type=str,
133 | help="Select task type, ext for the extraction task and cls for the classification task, defaults to ext.")
134 | parser.add_argument("--options", default=["正向", "负向"], type=str, nargs="+",
135 | help="Used only for the classification task, the options for classification")
136 | parser.add_argument("--prompt_prefix", default="情感倾向", type=str,
137 | help="Used only for the classification task, the prompt prefix for classification")
138 | parser.add_argument("--is_shuffle", default=True, type=bool,
139 | help="Whether to shuffle the labeled dataset, defaults to True.")
140 | parser.add_argument("--seed", type=int, default=1000,
141 | help="random seed for initialization")
142 |
143 | args = parser.parse_args()
144 | # yapf: enable
145 |
146 | do_convert()
147 |
--------------------------------------------------------------------------------
/export_model.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) 2022 Heiheiyoyo. All Rights Reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | import argparse
16 | import os
17 | from itertools import chain
18 | from typing import List, Union
19 | import shutil
20 | from pathlib import Path
21 |
22 | import numpy as np
23 | import torch
24 | from transformers import (BertTokenizer, PreTrainedModel,
25 | PreTrainedTokenizerBase)
26 |
27 | from model import UIE
28 | from utils import logger
29 |
30 |
31 | def validate_onnx(tokenizer: PreTrainedTokenizerBase, pt_model: PreTrainedModel, onnx_path: Union[Path, str], strict: bool = True, atol: float = 1e-05):
32 |
33 | # 验证模型
34 | from onnxruntime import InferenceSession, SessionOptions
35 | from transformers import AutoTokenizer
36 |
37 | logger.info("Validating ONNX model...")
38 | if strict:
39 | ref_inputs = tokenizer('装备', "印媒所称的“印度第一艘国产航母”—“维克兰特”号",
40 | add_special_tokens=True,
41 | truncation=True,
42 | max_length=512,
43 | return_tensors="pt")
44 | else:
45 | batch_size = 2
46 | seq_length = 6
47 | dummy_input = [" ".join([tokenizer.unk_token])
48 | * seq_length] * batch_size
49 | ref_inputs = dict(tokenizer(dummy_input, return_tensors="pt"))
50 | # ref_inputs =
51 | ref_outputs = pt_model(**ref_inputs)
52 | ref_outputs_dict = {}
53 |
54 | # We flatten potential collection of outputs (i.e. past_keys) to a flat structure
55 | for name, value in ref_outputs.items():
56 | # Overwriting the output name as "present" since it is the name used for the ONNX outputs
57 | # ("past_key_values" being taken for the ONNX inputs)
58 | if name == "past_key_values":
59 | name = "present"
60 | ref_outputs_dict[name] = value
61 |
62 | # Create ONNX Runtime session
63 | options = SessionOptions()
64 | session = InferenceSession(str(onnx_path), options, providers=[
65 | "CPUExecutionProvider"])
66 |
67 | # We flatten potential collection of inputs (i.e. past_keys)
68 | onnx_inputs = {}
69 | for name, value in ref_inputs.items():
70 | onnx_inputs[name] = value.numpy()
71 | onnx_named_outputs = ['start_prob', 'end_prob']
72 | # Compute outputs from the ONNX model
73 | onnx_outputs = session.run(onnx_named_outputs, onnx_inputs)
74 |
75 | # Check we have a subset of the keys into onnx_outputs against ref_outputs
76 | ref_outputs_set, onnx_outputs_set = set(
77 | ref_outputs_dict.keys()), set(onnx_named_outputs)
78 | if not onnx_outputs_set.issubset(ref_outputs_set):
79 | logger.info(
80 | f"\t-[x] ONNX model output names {onnx_outputs_set} do not match reference model {ref_outputs_set}"
81 | )
82 |
83 | raise ValueError(
84 | "Outputs doesn't match between reference model and ONNX exported model: "
85 | f"{onnx_outputs_set.difference(ref_outputs_set)}"
86 | )
87 | else:
88 | logger.info(
89 | f"\t-[✓] ONNX model output names match reference model ({onnx_outputs_set})")
90 |
91 | # Check the shape and values match
92 | for name, ort_value in zip(onnx_named_outputs, onnx_outputs):
93 | ref_value = ref_outputs_dict[name].detach().numpy()
94 |
95 | logger.info(f'\t- Validating ONNX Model output "{name}":')
96 |
97 | # Shape
98 | if not ort_value.shape == ref_value.shape:
99 | logger.info(
100 | f"\t\t-[x] shape {ort_value.shape} doesn't match {ref_value.shape}")
101 | raise ValueError(
102 | "Outputs shape doesn't match between reference model and ONNX exported model: "
103 | f"Got {ref_value.shape} (reference) and {ort_value.shape} (ONNX)"
104 | )
105 | else:
106 | logger.info(
107 | f"\t\t-[✓] {ort_value.shape} matches {ref_value.shape}")
108 |
109 | # Values
110 | if not np.allclose(ref_value, ort_value, atol=atol):
111 | logger.info(f"\t\t-[x] values not close enough (atol: {atol})")
112 | raise ValueError(
113 | "Outputs values doesn't match between reference model and ONNX exported model: "
114 | f"Got max absolute difference of: {np.amax(np.abs(ref_value - ort_value))}"
115 | )
116 | else:
117 | logger.info(f"\t\t-[✓] all values close (atol: {atol})")
118 |
119 |
120 | def export_onnx(args: argparse.Namespace, tokenizer: PreTrainedTokenizerBase, model: PreTrainedModel, device: torch.device, input_names: List[str], output_names: List[str]):
121 | with torch.no_grad():
122 | model = model.to(device)
123 | model.eval()
124 | model.config.return_dict = True
125 | model.config.use_cache = False
126 |
127 | # Create folder
128 | if not args.output_path.exists():
129 | args.output_path.mkdir(parents=True)
130 | save_path = args.output_path / "inference.onnx"
131 |
132 | dynamic_axes = {name: {0: 'batch', 1: 'sequence'}
133 | for name in chain(input_names, output_names)}
134 |
135 | # Generate dummy input
136 | batch_size = 2
137 | seq_length = 6
138 | dummy_input = [" ".join([tokenizer.unk_token])
139 | * seq_length] * batch_size
140 | inputs = dict(tokenizer(dummy_input, return_tensors="pt"))
141 |
142 | if save_path.exists():
143 | logger.warning(f'Overwrite model {save_path.as_posix()}')
144 | save_path.unlink()
145 |
146 | torch.onnx.export(model,
147 | (inputs,),
148 | save_path,
149 | input_names=input_names,
150 | output_names=output_names,
151 | dynamic_axes=dynamic_axes,
152 | do_constant_folding=True,
153 | opset_version=11
154 | )
155 |
156 | if not os.path.exists(save_path):
157 | logger.error(f'Export Failed!')
158 |
159 | return save_path
160 |
161 |
162 | def main():
163 | parser = argparse.ArgumentParser()
164 | parser.add_argument("-m", "--model_path", type=Path, required=True,
165 | default='./checkpoint/model_best', help="The path to model parameters to be loaded.")
166 | parser.add_argument("-o", "--output_path", type=Path, default=None,
167 | help="The path of model parameter in static graph to be saved.")
168 | args = parser.parse_args()
169 |
170 | if args.output_path is None:
171 | args.output_path = args.model_path
172 |
173 | tokenizer = BertTokenizer.from_pretrained(args.model_path)
174 | model = UIE.from_pretrained(args.model_path)
175 | device = torch.device('cpu')
176 | input_names = [
177 | 'input_ids',
178 | 'token_type_ids',
179 | 'attention_mask',
180 | ]
181 | output_names = [
182 | 'start_prob',
183 | 'end_prob'
184 | ]
185 |
186 | logger.info("Export Tokenizer Config...")
187 |
188 | export_tokenizer(args)
189 |
190 | logger.info("Export ONNX Model...")
191 |
192 | save_path = export_onnx(
193 | args, tokenizer, model, device, input_names, output_names)
194 | validate_onnx(tokenizer, model, save_path)
195 |
196 | logger.info(f"All good, model saved at: {save_path.as_posix()}")
197 |
198 |
199 | def export_tokenizer(args):
200 | for tokenizer_fine in ['tokenizer_config.json', 'special_tokens_map.json', 'vocab.txt']:
201 | file_from = args.model_path / tokenizer_fine
202 | file_to = args.output_path/tokenizer_fine
203 | if file_from.resolve() == file_to.resolve():
204 | continue
205 | shutil.copyfile(file_from, file_to)
206 |
207 |
208 | if __name__ == "__main__":
209 |
210 | main()
211 |
--------------------------------------------------------------------------------
/doccano.md:
--------------------------------------------------------------------------------
1 | # doccano
2 |
3 | **目录**
4 |
5 | * [1. 安装](#安装)
6 | * [2. 项目创建](#项目创建)
7 | * [3. 数据上传](#数据上传)
8 | * [4. 标签构建](#标签构建)
9 | * [5. 任务标注](#任务标注)
10 | * [6. 数据导出](#数据导出)
11 | * [7. 数据转换](#数据转换)
12 |
13 |
14 |
15 | ## 1. 安装
16 |
17 | 参考[doccano官方文档](https://github.com/doccano/doccano) 完成doccano的安装与初始配置。
18 |
19 | **以下标注示例用到的环境配置：**
20 |
21 | - doccano 1.6.2
22 |
23 |
24 |
25 | ## 2. 项目创建
26 |
27 | UIE支持抽取与分类两种类型的任务，根据实际需要创建一个新的项目：
28 |
29 | #### 2.1 抽取式任务项目创建
30 |
31 | 创建项目时选择**序列标注**任务，并勾选**Allow overlapping entity**及**Use relation Labeling**。适配**命名实体识别、关系抽取、事件抽取、评价观点抽取**等任务。
32 |
33 |
34 |
35 |
42 |
43 |
59 |
60 |
75 |
76 |
81 |
82 |
89 |
90 |
103 |
104 |
115 |
116 |
127 |
128 |
139 |
140 |
149 |
150 |