Merge pull request #182 from wenlisong/master

issue#101: fix bug
macanv · Jul 24, 2019 · 348c2ed · 348c2ed
2 parents 9c9e7b1 + 32fd3c9
commit 348c2ed
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/bert_base/server/graph.py b/bert_base/server/graph.py
@@ -286,7 +286,7 @@ def optimize_ner_model(args, num_labels,  logger=None):
                 from bert_base.train.models import create_model
                 (total_loss, logits, trans, pred_ids) = create_model(
                     bert_config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, segment_ids=None,
-                    labels=None, num_labels=num_labels, use_one_hot_embeddings=False, dropout_rate=1.0)
+                    labels=None, num_labels=num_labels, use_one_hot_embeddings=False, dropout_rate=1.0, lstm_size=args.lstm_size)
                 pred_ids = tf.identity(pred_ids, 'pred_ids')
                 saver = tf.train.Saver()
 

diff --git a/bert_base/server/helper.py b/bert_base/server/helper.py
@@ -98,6 +98,8 @@ def get_args_parser():
                         help='masking the embedding on [CLS] and [SEP] with zero. \
                         When pooling_strategy is in {CLS_TOKEN, FIRST_TOKEN, SEP_TOKEN, LAST_TOKEN} \
                         then the embedding is preserved, otherwise the embedding is masked to zero before pooling')
+    group2.add_argument('-lstm_size', type=int, default=128,
+                        help='size of lstm units.')
 
     group3 = parser.add_argument_group('Serving Configs',
                                        'config how server utilizes GPU/CPU resources')