中文命名实体识别（ERNIE-艾尼）

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import os

import paddle
import paddle.fluid as fluid
import paddlehub as hub
from paddlehub.finetune.evaluate import chunk_eval, calculate_f1


# loading Paddlehub ERNIE pretrained model
# 主要的变动就是这里加上了版本号码
module = hub.Module(name="ernie",version="1.0.2")
inputs, outputs, program = module.context(max_seq_len=128)

# Sentence labeling dataset reader
# 可以通过自定义一个数据集的类
dataset = hub.dataset.MSRA_NER()
reader = hub.reader.SequenceLabelReader(
    dataset=dataset,
    vocab_path=module.get_vocab_path(),
    max_seq_len=128)
inv_label_map = {val: key for key, val in reader.label_map.items()}

# Construct transfer learning network
# Use "sequence_output" for token-level output.
sequence_output = outputs["sequence_output"]

# Setup feed list for data feeder
# Must feed all the tensor of ERNIE's module need
feed_list = [
    inputs["input_ids"].name,
    inputs["position_ids"].name,
    inputs["segment_ids"].name,
    inputs["input_mask"].name,
]

strategy = hub.AdamWeightDecayStrategy(
    weight_decay=0.01,
    warmup_proportion=0.1,
    learning_rate=5e-5,
    lr_scheduler="linear_decay",
    optimizer_name="adam")

# Setup runing config for PaddleHub Finetune API
config = hub.RunConfig(
    use_data_parallel=False,
    use_pyreader=False,
    use_cuda=True,
    batch_size=1,
    enable_memory_optim=False,
    checkpoint_dir='ernie_seq_label_turtorial_demo',
    strategy=strategy)

# Define a sequence labeling finetune task by PaddleHub's API
seq_label_task = hub.SequenceLabelTask(
    data_reader=reader,
    feature=sequence_output,
    feed_list=feed_list,
    max_seq_len=10000,
    num_classes=dataset.num_labels,
    config=config)

# test data
data = [
    ["我们变而以书会友，以书结缘，把欧美、港台流行的食品类图谱、画册、工具书汇集一堂。"],
    ["为了跟踪国际最新食品工艺、流行趋势，大量搜集海外专业书刊资料是提高技艺的捷径。"],
    ["其中线装古籍逾千册；民国出版物几百种；珍本四册、稀见本四百余册，出版时间跨越三百余年。"],
    ["有的古木交柯，春机荣欣，从诗人句中得之，而入画中，观之令人心驰。"],
    ["不过重在晋趣，略增明人气息，妙在集古有道、不露痕迹罢了。"],
]

run_states = seq_label_task.predict(data=data)
results = [run_state.run_results for run_state in run_states]

for num_batch, batch_results in enumerate(results):
    infers = batch_results[0].reshape([-1]).astype(np.int32).tolist()

    np_lens = batch_results[1]

    for index, np_len in enumerate(np_lens):
        labels = infers[index * 128:(index + 1) * 128]

        label_str = ""
        count = 0
        for label_val in labels:
            label_str += inv_label_map[label_val]
            count += 1
            if count == np_len:
                break

        # Drop the label results of CLS and SEP Token
        print(
            "%s\tpredict=%s" %
            (data[num_batch + index][0], label_str[1:-1]))

代码仅在AI STUDIO和Linux环境中可以运行，windows环境中暂时报错