From c695bfb93c6057c0a8fa1ed5e680eea944e124be Mon Sep 17 00:00:00 2001
From: LiyuanLucasLiu <llychinalz@gmail.com>
Date: Tue, 19 Sep 2017 15:25:02 -0500
Subject: [PATCH] add tqdm to predict

---
 model/predictor.py | 5 ++++-
 seq_wc.py          | 7 +++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/model/predictor.py b/model/predictor.py
index be4f5f7..6deaae2 100644
--- a/model/predictor.py
+++ b/model/predictor.py
@@ -9,6 +9,8 @@
 import torch.autograd as autograd
 import numpy as np
 import itertools
+import sys
+from tqdm import tqdm
 
 from model.crf import CRFDecode_vb
 from model.utils import *
@@ -123,7 +125,8 @@ def output_batch(self, ner_model, features, fout):
         """
         f_len = len(features)
 
-        for ind in range(0, f_len, self.batch_size):
+        for ind in tqdm( range(0, f_len, self.batch_size), mininterval=1,
+                desc=' - Process', leave=False, file=sys.stdout):
             eind = min(f_len, ind + self.batch_size)
             labels = self.apply_model(ner_model, features[ind: eind])
             labels = torch.unbind(labels, 1)
diff --git a/seq_wc.py b/seq_wc.py
index 1b01173..06fa5fd 100644
--- a/seq_wc.py
+++ b/seq_wc.py
@@ -30,6 +30,7 @@
     parser.add_argument('--output_file', default='output.txt', help='path to output file')
     args = parser.parse_args()
 
+    print('loading dictionary')
     with open(args.load_arg, 'r') as f:
         jd = json.load(f)
     jd = jd['args']
@@ -42,15 +43,16 @@
     if args.gpu >= 0:
         torch.cuda.set_device(args.gpu)
 
-    # load corpus
+    # loading corpus
+    print('loading corpus')
     with codecs.open(args.input_file, 'r', 'utf-8') as f:
         lines = f.readlines()
 
     # converting format
-
     features = utils.read_features(lines)
 
     # build model
+    print('loading model')
     ner_model = LM_LSTM_CRF(len(l_map), len(c_map), jd['char_dim'], jd['char_hidden'], jd['char_layers'], jd['word_dim'], jd['word_hidden'], jd['word_layers'], len(f_map), jd['drop_out'], large_CRF=jd['small_crf'], if_highway=jd['high_way'], in_doc_words=in_doc_words, highway_layers = jd['highway_layers'])
 
     ner_model.load_state_dict(checkpoint_file['state_dict'])
@@ -67,5 +69,6 @@
     decode_label = (args.decode_type == 'label')
     predictor = predict_wc(if_cuda, f_map, c_map, l_map, f_map['<eof>'], c_map['\n'], l_map['<pad>'], l_map['<start>'], decode_label, args.batch_size, jd['caseless'])
 
+    print('annotating')
     with open(args.output_file, 'w') as fout:
         predictor.output_batch(ner_model, features, fout)
\ No newline at end of file