wilsontam commited on
Commit
8b7011b
1 Parent(s): 498f30f

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +11 -1
README.md CHANGED
@@ -28,9 +28,19 @@ import json
28
 
29
  model_path = "wilsontam/dstc9_ner"
30
 
 
 
 
 
 
 
 
 
 
 
31
  config = AutoConfig.from_pretrained(
32
  model_path,
33
- num_labels=7,
34
  )
35
  model = AutoModelForTokenClassification.from_pretrained(
36
  model_path,
 
28
 
29
  model_path = "wilsontam/dstc9_ner"
30
 
31
+ label_map = {
32
+ "LABEL_0": "O",
33
+ "LABEL_1": "B-hotel",
34
+ "LABEL_2": "I-hotel",
35
+ "LABEL_3": "B-restaurant",
36
+ "LABEL_4": "I-restaurant",
37
+ "LABEL_5": "B-attraction",
38
+ "LABEL_6": "I-attraction",
39
+ }
40
+
41
  config = AutoConfig.from_pretrained(
42
  model_path,
43
+ num_labels=len(label_map),
44
  )
45
  model = AutoModelForTokenClassification.from_pretrained(
46
  model_path,