让大家久等了，BERT推理加速终于开源了

pip install lightseq transformers

output_lightseq_model_name = "lightseq-bert-base-chinese"
input_huggingface_bert_model = "bert-base-chinese"

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
import lightseq.inference as lsi

tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
hf_model = AutoModelForMaskedLM.from_pretrained("bert-base-chinese")
hf_model.to("cuda:0")
ls_model = lsi.Bert("lightseq-bert-base-chinese.hdf5", 128)

while True:
    raw_text = input("请输入中文句子，要预测的字符用#代替：\n> ")
    input_text = raw_text.replace("#", "[MASK]")
    inputs = tokenizer(input_text, return_tensors="pt")
    input_ids = inputs["input_ids"]
    mask = inputs["attention_mask"]

    outputs = ls_model.infer(input_ids, mask)
    logits = hf_model.cls(torch.Tensor(outputs).to(dtype=torch.float, device="cuda:0"))
    output_ids = logits.argmax(axis=2)
    res_text = tokenizer.batch_decode(output_ids)

    res_text = res_text[0][1:-1].replace(" ", "")
    output_text = list(raw_text)
    for i in range(len(raw_text)):
        if raw_text[i] == "#":
            output_text[i] = res_text[i]
    print("> " + "".join(output_text))

请输入中文句子，要预测的字符用#代替：
> 巴黎是#国的首都。
> 巴黎是法国的首都。