feat: TQDM added to trainers (#1593)

Co-authored-by: Fedor Ignatov <[email protected]>

Author: dimakarp1996

deeppavlov/core/trainers/fit_trainer.py

@@ -19,6 +19,8 @@
 from logging import getLogger
 from typing import Tuple, Dict, Union, Optional, Iterable, Any, Collection
 
+from tqdm import tqdm
+
 from deeppavlov.core.commands.infer import build_model
 from deeppavlov.core.common.chainer import Chainer
 from deeppavlov.core.common.params import from_params
@@ -90,7 +92,7 @@ def fit_chainer(self, iterator: Union[DataFittingIterator, DataLearningIterator]
                     targets = [targets]
 
                 if self.batch_size > 0 and callable(getattr(component, 'partial_fit', None)):
-                    for i, (x, y) in enumerate(iterator.gen_batches(self.batch_size, shuffle=False)):
+                    for i, (x, y) in tqdm(enumerate(iterator.gen_batches(self.batch_size, shuffle=False))):
                         preprocessed = self._chainer.compute(x, y, targets=targets)
                         # noinspection PyUnresolvedReferences
                         component.partial_fit(*preprocessed)
@@ -160,7 +162,7 @@ def test(self, data: Iterable[Tuple[Collection[Any], Collection[Any]]],
 
         data = islice(data, self.max_test_batches)
 
-        for x, y_true in data:
+        for x, y_true in tqdm(data):
             examples += len(x)
             y_predicted = list(self._chainer.compute(list(x), list(y_true), targets=expected_outputs))
             if len(expected_outputs) == 1:

deeppavlov/core/trainers/nn_trainer.py

@@ -20,12 +20,15 @@
 from pathlib import Path
 from typing import List, Tuple, Union, Optional, Iterable
 
+from tqdm import tqdm
+
 from deeppavlov.core.common.errors import ConfigError
+from deeppavlov.core.common.log_events import get_tb_writer
 from deeppavlov.core.common.registry import register
 from deeppavlov.core.data.data_learning_iterator import DataLearningIterator
 from deeppavlov.core.trainers.fit_trainer import FitTrainer
 from deeppavlov.core.trainers.utils import parse_metrics, NumpyArrayEncoder
-from deeppavlov.core.common.log_events import get_tb_writer
+
 log = getLogger(__name__)
 report_log = getLogger('train_report')
 
@@ -273,7 +276,7 @@ def train_on_batches(self, iterator: DataLearningIterator) -> None:
         while True:
             impatient = False
             self._send_event(event_name='before_train')
-            for x, y_true in iterator.gen_batches(self.batch_size, data_type='train'):
+            for x, y_true in tqdm(iterator.gen_batches(self.batch_size, data_type='train')):
                 self.last_result = self._chainer.train_on_batch(x, y_true)
                 if self.last_result is None:
                     self.last_result = {}