fix use undefined ner_model_dir arg bug

macanv · Feb 20, 2019 · 4b7a243 · 4b7a243
1 parent ba954b1
commit 4b7a243
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bert_base/server/__init__.py b/bert_base/server/__init__.py
@@ -92,7 +92,7 @@ def __init__(self, args):
             with Pool(processes=1) as pool:
                 # optimize the graph, must be done in another process
                 from .graph import optimize_ner_model
-                num_labels, label2id, id2label = init_predict_var(self.args.ner_model_dir)
+                num_labels, label2id, id2label = init_predict_var(self.args.model_dir)
                 self.num_labels = num_labels + 1
                 self.id2label = id2label
                 self.graph_path = pool.apply(optimize_ner_model, (self.args, num_labels))