dstc9_ner / README.md
wilsontam's picture
Update README.md
498f30f
|
raw
history blame
1.15 kB
metadata
language: en
tags:
  - dstc9
widget:
  - text: i want to book the hilton hotel near china town.
  - text: can you reserve A & B restaurant for me?

Only restaurant, hotel, and attraction names are tagged based on the following data and knowledge base.

Data link: https://github.com/alexa/alexa-with-dstc9-track1-dataset

Label map:

"O": 0 "B-hotel": 1 "I-hotel": 2 "B-restaurant": 3 "I-restaurant": 4 "B-attraction": 5 "I-attraction": 6

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

model_path = "wilsontam/dstc9_ner"

config = AutoConfig.from_pretrained(
  model_path,
  num_labels=7,
)   
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