Improve train/fit progress bar status

sergioburdisso · Feb 22, 2020 · 2029c37 · 2029c37
1 parent 8d163d9
commit 2029c37
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/pyss3/__init__.py b/pyss3/__init__.py
@@ -2059,12 +2059,16 @@ def fit(self, x_train, y_train, n_grams=1, prep=True, leave_pbar=True):
 
         Print.info("about to start training", offset=1)
         Print.verbosity_region_begin(VERBOSITY.NORMAL)
-        for i in tqdm(range(len(x_train)), desc="Training",
-                      leave=leave_pbar, disable=Print.is_quiet()):
+        progress_bar = tqdm(total=len(x_train), desc="Training",
+                            leave=leave_pbar, disable=Print.is_quiet())
+        for i in range(len(x_train)):
+            progress_bar.set_description_str("Training on '%s'" % str(y_train[i]))
             self.learn(
                 x_train[i], y_train[i],
                 n_grams=n_grams, prep=prep, update=False
             )
+            progress_bar.update(1)
+        progress_bar.close()
         self.__prune_tries__()
         Print.verbosity_region_end()
         Print.info("finished --time: %.1fs" % (time() - stime), offset=1)