Convert PyTorch* BERT-NER Model

Download and Convert the Model to ONNX*

To download a pre-trained model or train the model yourself, refer to the instruction in the BERT-NER model repository. The model with config files is stored in the out_base directory.

To convert the model to ONNX* format, create and run the script with the following content in the root directory of the model repository. If you download the pre-trained model, you need to download ` <>`__ to run the script. The instruction was tested with the repository hash commit e5be564156f194f1becb0d82aeaf6e762d9eb9ed.

import torch

from bert import Ner

ner = Ner("out_base")

input_ids, input_mask, segment_ids, valid_positions = ner.preprocess('Steve went to Paris')
input_ids = torch.tensor([input_ids], dtype=torch.long, device=ner.device)
input_mask = torch.tensor([input_mask], dtype=torch.long, device=ner.device)
segment_ids = torch.tensor([segment_ids], dtype=torch.long, device=ner.device)
valid_ids = torch.tensor([valid_positions], dtype=torch.long, device=ner.device)

ner_model, tknizr, model_config = ner.load_model("out_base")

with torch.no_grad():
    logits = ner_model(input_ids, segment_ids, input_mask, valid_ids)
                  (input_ids, segment_ids, input_mask, valid_ids),
                  input_names=['input_ids', 'segment_ids', 'input_mask', 'valid_ids'],
                      "input_ids": {0: "batch_size"},
                      "segment_ids": {0: "batch_size"},
                      "input_mask": {0: "batch_size"},
                      "valid_ids": {0: "batch_size"},
                      "output": {0: "output"}

The script generates ONNX* model file bert-ner.onnx.

Convert ONNX* BERT-NER model to IR

mo --input_model bert-ner.onnx --input "input_mask[1 128],segment_ids[1 128],input_ids[1 128]"

where 1 is batch_size and 128 is sequence_length.