from transformers import AutoConfig, AutoModelForTokenClassification, BertTokenizer
from transformers import TokenClassificationPipeline
import json

model_path = "wilsontam/dstc9_ner"

label_map = {
"LABEL_0": "O",
"LABEL_1": "B-hotel",
"LABEL_2": "I-hotel",
"LABEL_3": "B-restaurant",
"LABEL_4": "I-restaurant",
"LABEL_5": "B-attraction",
"LABEL_6": "I-attraction",
}

config = AutoConfig.from_pretrained(
  model_path,
  num_labels=len(label_map),
)   
model = AutoModelForTokenClassification.from_pretrained(
  model_path,
  from_tf=False,
  config=config,
)   
tokenizer = BertTokenizer.from_pretrained(
  model_path,
)

# device=-1: cpu, device=0: gpu
pipeline = TokenClassificationPipeline(model, tokenizer, device=-1)

tokens = pipeline(["i want to book the hilton hotel near china town.", "can you reserve A & B restaurant for me?"])

Credit: Jia-Chen Jason Gu, Wilson Tam