From 1792003420bd46126a1e7dc23c8c0eb3ef2ee286 Mon Sep 17 00:00:00 2001
From: macan <ma_cancan@163.com>
Date: Fri, 1 Mar 2019 16:51:36 +0800
Subject: [PATCH] fix ner model as service bug

---
 bert_base/server/__init__.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/bert_base/server/__init__.py b/bert_base/server/__init__.py
index a67cc7e..26cb51a 100644
--- a/bert_base/server/__init__.py
+++ b/bert_base/server/__init__.py
@@ -95,7 +95,7 @@ def __init__(self, args):
                 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))
+                self.graph_path = pool.apply(optimize_ner_model, (self.args, self.num_labels))
             if self.graph_path:
                 self.logger.info('optimized graph is stored at: %s' % self.graph_path)
             else:
@@ -509,7 +509,7 @@ def _run(self, sink, *receivers):
             elif self.mode == 'NER':
                 pred_label_result, pred_ids_result = ner_result_to_json(r['encodes'], self.id2label)
                 rst = send_ndarray(sink, r['client_id'], pred_label_result)
-                # print(rst)
+                # print('rst:', rst)
                 logger.info('job done\tsize: %s\tclient: %s' % (r['encodes'].shape, r['client_id']))
             elif self.mode == 'CLASS':
                 pred_label_result = [self.id2label.get(x, -1) for x in r['encodes']]
@@ -661,6 +661,7 @@ def ner_result_to_json(predict_ids, id2label):
     """
     NER识别结果转化为真实标签结果进行返回
     :param predict_ids:
+    :param id2label
     :return:
     """
     if False: