Repository: tensorflow/addons

Issue:
Have TQDMProgressBar work in evaluate mode as well
**Describe the feature and the current behavior/state.**

Right now, the `TQDMProgressBar` callback can only be used when training a model.
When evaluating nothing happens. And even if I try to "cheat" by manually setting the `on_test_batch_(begin|end)` methods, it fails with an error about `mode`.

**Relevant information**
- Are you willing to contribute it (yes/no): don't know
- Are you willing to maintain it going forward? (yes/no): no
- Is there a relevant academic paper? (if so, where): no
- Is there already an implementation in another framework? (if so, where): no
- Was it part of tf.contrib? (if so, where): no

**Which API type would this fall under (layer, metric, optimizer, etc.)**
callbacks

**Who will benefit with this feature?**
everyone using `tqdm` to evaluate in a notebook: I want something fancier than the default evaluate progressbar

**Any other info.**
It failed in both tf2.2 and tf2.1

You can find the colab gist [here](https://colab.research.google.com/drive/1eHI-6QhttMhdMYeH8VP0SwfyB0b46ugj).


Discussion:
I don't think it would be hard to do because keras exposes methods for that: e.g. https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/Callback#on_test_batch_begin . Contributions welcome!
@gabrieldemarmiesse , please review PR [#38447](https://github.com/tensorflow/tensorflow/pull/38447)
As a small note, I believe the TQDM callback should have a cleanup as we removed the "samples" mode and a refactoring would be welcome before a PR to add this feature. There is definitly some dead code/renaming to do.
Yep. Got it. Let me do that too.

Resolving patch:
diff --git a/tensorflow_addons/callbacks/tqdm_progress_bar.py b/tensorflow_addons/callbacks/tqdm_progress_bar.py
--- a/tensorflow_addons/callbacks/tqdm_progress_bar.py
+++ b/tensorflow_addons/callbacks/tqdm_progress_bar.py
@@ -93,71 +93,74 @@ def __init__(
         self.last_update_time = time.time()
         self.overall_progress_tqdm = None
         self.epoch_progress_tqdm = None
+        self.is_training = False
         self.num_epochs = None
         self.logs = None
         super().__init__()
 
-    def on_train_begin(self, logs=None):
-        self.num_epochs = self.params["epochs"]
-
-        if self.show_overall_progress:
-            self.overall_progress_tqdm = self.tqdm(
-                desc="Training",
-                total=self.num_epochs,
-                bar_format=self.overall_bar_format,
-                leave=self.leave_overall_progress,
-                dynamic_ncols=True,
-                unit="epochs",
-            )
-
-        # set counting mode
-        self.mode = "steps"
-        self.total_steps = self.params["steps"]
-
-    def on_train_end(self, logs={}):
-        if self.show_overall_progress:
-            self.overall_progress_tqdm.close()
-
-    def on_epoch_begin(self, epoch, logs={}):
-        current_epoch_description = "Epoch {epoch}/{num_epochs}".format(
-            epoch=epoch + 1, num_epochs=self.num_epochs
-        )
-
-        if self.show_epoch_progress:
-            print(current_epoch_description)
-            self.epoch_progress_tqdm = self.tqdm(
-                total=self.total_steps,
-                bar_format=self.epoch_bar_format,
-                leave=self.leave_epoch_progress,
-                dynamic_ncols=True,
-                unit=self.mode,
-            )
-
+    def _initialize_progbar(self, hook, epoch, logs=None):
         self.num_samples_seen = 0
         self.steps_to_update = 0
         self.steps_so_far = 0
         self.logs = defaultdict(float)
-
-    def on_epoch_end(self, epoch, logs={}):
-
-        if self.show_epoch_progress:
-            metrics = self.format_metrics(logs)
-            self.epoch_progress_tqdm.desc = metrics
-
-            # set miniters and mininterval to 0 so last update displays
-            self.epoch_progress_tqdm.miniters = 0
-            self.epoch_progress_tqdm.mininterval = 0
-
-            # update the rest of the steps in epoch progress bar
-            self.epoch_progress_tqdm.update(
-                self.total_steps - self.epoch_progress_tqdm.n
+        self.num_epochs = self.params["epochs"]
+        self.mode = "steps"
+        self.total_steps = self.params["steps"]
+        if hook == "train_overall":
+            if self.show_overall_progress:
+                self.overall_progress_tqdm = self.tqdm(
+                    desc="Training",
+                    total=self.num_epochs,
+                    bar_format=self.overall_bar_format,
+                    leave=self.leave_overall_progress,
+                    dynamic_ncols=True,
+                    unit="epochs",
+                )
+        elif hook == "test":
+            if self.show_epoch_progress:
+                self.epoch_progress_tqdm = self.tqdm(
+                    total=self.total_steps,
+                    desc="Evaluating",
+                    bar_format=self.epoch_bar_format,
+                    leave=self.leave_epoch_progress,
+                    dynamic_ncols=True,
+                    unit=self.mode,
+                )
+        elif hook == "train_epoch":
+            current_epoch_description = "Epoch {epoch}/{num_epochs}".format(
+                epoch=epoch + 1, num_epochs=self.num_epochs
             )
-            self.epoch_progress_tqdm.close()
-
-        if self.show_overall_progress:
-            self.overall_progress_tqdm.update(1)
-
-    def on_batch_end(self, batch, logs={}):
+            if self.show_epoch_progress:
+                print(current_epoch_description)
+                self.epoch_progress_tqdm = self.tqdm(
+                    total=self.total_steps,
+                    bar_format=self.epoch_bar_format,
+                    leave=self.leave_epoch_progress,
+                    dynamic_ncols=True,
+                    unit=self.mode,
+                )
+
+    def _clean_up_progbar(self, hook, logs):
+        if hook == "train_overall":
+            if self.show_overall_progress:
+                self.overall_progress_tqdm.close()
+        else:
+            if hook == "test":
+                metrics = self.format_metrics(logs, self.num_samples_seen)
+            else:
+                metrics = self.format_metrics(logs)
+            if self.show_epoch_progress:
+                self.epoch_progress_tqdm.desc = metrics
+                # set miniters and mininterval to 0 so last update displays
+                self.epoch_progress_tqdm.miniters = 0
+                self.epoch_progress_tqdm.mininterval = 0
+                # update the rest of the steps in epoch progress bar
+                self.epoch_progress_tqdm.update(
+                    self.total_steps - self.epoch_progress_tqdm.n
+                )
+                self.epoch_progress_tqdm.close()
+
+    def _update_progbar(self, logs):
         if self.mode == "samples":
             batch_size = logs["size"]
         else:
@@ -167,8 +170,7 @@ def on_batch_end(self, batch, logs={}):
         self.steps_to_update += 1
         self.steps_so_far += 1
 
-        if self.steps_so_far < self.total_steps:
-
+        if self.steps_so_far <= self.total_steps:
             for metric, value in logs.items():
                 self.logs[metric] += value * batch_size
 
@@ -187,6 +189,37 @@ def on_batch_end(self, batch, logs={}):
                 # update timestamp for last update
                 self.last_update_time = now
 
+    def on_train_begin(self, logs=None):
+        self.is_training = True
+        self._initialize_progbar("train_overall", None, logs)
+
+    def on_train_end(self, logs={}):
+        self.is_training = False
+        self._clean_up_progbar("train_overall", logs)
+
+    def on_test_begin(self, logs={}):
+        if not self.is_training:
+            self._initialize_progbar("test", None, logs)
+
+    def on_test_end(self, logs={}):
+        if not self.is_training:
+            self._clean_up_progbar("test", self.logs)
+
+    def on_epoch_begin(self, epoch, logs={}):
+        self._initialize_progbar("train_epoch", epoch, logs)
+
+    def on_epoch_end(self, epoch, logs={}):
+        self._clean_up_progbar("train_epoch", logs)
+        if self.show_overall_progress:
+            self.overall_progress_tqdm.update(1)
+
+    def on_test_batch_end(self, batch, logs={}):
+        if not self.is_training:
+            self._update_progbar(logs)
+
+    def on_batch_end(self, batch, logs={}):
+        self._update_progbar(logs)
+
     def format_metrics(self, logs={}, factor=1):
         """Format metrics in logs into a string.
 


Test patch:
diff --git a/tensorflow_addons/callbacks/tests/tqdm_progress_bar_test.py b/tensorflow_addons/callbacks/tests/tqdm_progress_bar_test.py
--- a/tensorflow_addons/callbacks/tests/tqdm_progress_bar_test.py
+++ b/tensorflow_addons/callbacks/tests/tqdm_progress_bar_test.py
@@ -120,3 +120,14 @@ def test_tqdm_progress_bar_validation(capsys):
     fit_stderr = capsys.readouterr().err
     assert re.search(r"val_loss: [0-9]\.[0-9][0-9][0-9][0-9]", fit_stderr)
     assert re.search(r"val_acc: [0-9]\.[0-9][0-9][0-9][0-9]", fit_stderr)
+
+
+def test_tqdm_progress_bar_evaluate(capsys):
+    x, y, model = get_data_and_model()
+
+    pb = tfa.callbacks.TQDMProgressBar()
+    capsys.readouterr()  # flush the buffer
+    model.evaluate(x, y, callbacks=[pb], verbose=0)
+    evaluate_stderr = capsys.readouterr().err
+    assert re.search(r"loss: [0-9]\.[0-9][0-9][0-9][0-9]", evaluate_stderr)
+    assert re.search(r"acc: [0-9]\.[0-9][0-9][0-9][0-9]", evaluate_stderr)