Fix progress bar

Author: fealho

ctgan/synthesizers/_utils.py

@@ -55,3 +55,14 @@ def _set_device(enable_gpu, device=None):
 def validate_and_set_device(enable_gpu, cuda):
     enable_gpu = get_enable_gpu_value(enable_gpu, cuda)
     return _set_device(enable_gpu)
+
+
+def _format_score(score):
+    """Format a score as a fixed-length string ``±XX.XX``.
+
+    Values are clipped to the range ``[-99.99, +99.99]`` so the result
+    is always exactly 6 characters.
+    """
+    score = max(-99.99, min(99.99, score))
+    sign = '+' if score >= 0 else '-'
+    return f'{sign}{abs(score):05.2f}'

ctgan/synthesizers/ctgan.py

@@ -12,7 +12,7 @@
 from ctgan.data_sampler import DataSampler
 from ctgan.data_transformer import DataTransformer
 from ctgan.errors import InvalidDataError
-from ctgan.synthesizers._utils import _set_device, validate_and_set_device
+from ctgan.synthesizers._utils import _format_score, _set_device, validate_and_set_device
 from ctgan.synthesizers.base import BaseSynthesizer, random_state
 
 
@@ -379,8 +379,10 @@ def fit(self, train_data, discrete_columns=(), epochs=None):
 
         epoch_iterator = tqdm(range(epochs), disable=(not self._verbose))
         if self._verbose:
-            description = 'Gen. ({gen:.2f}) | Discrim. ({dis:.2f})'
-            epoch_iterator.set_description(description.format(gen=0, dis=0))
+            description = 'Gen. ({gen}) | Discrim. ({dis})'
+            epoch_iterator.set_description(
+                description.format(gen=_format_score(0), dis=_format_score(0))
+            )
 
         steps_per_epoch = max(len(train_data) // self._batch_size, 1)
         for i in epoch_iterator:
@@ -479,7 +481,10 @@ def fit(self, train_data, discrete_columns=(), epochs=None):
 
             if self._verbose:
                 epoch_iterator.set_description(
-                    description.format(gen=generator_loss, dis=discriminator_loss)
+                    description.format(
+                        gen=_format_score(generator_loss),
+                        dis=_format_score(discriminator_loss),
+                    )
                 )
 
     @random_state

ctgan/synthesizers/tvae.py

@@ -10,7 +10,7 @@
 from tqdm import tqdm
 
 from ctgan.data_transformer import DataTransformer
-from ctgan.synthesizers._utils import _set_device, validate_and_set_device
+from ctgan.synthesizers._utils import _format_score, _set_device, validate_and_set_device
 from ctgan.synthesizers.base import BaseSynthesizer, random_state
 
 
@@ -161,8 +161,8 @@ def fit(self, train_data, discrete_columns=()):
         self.loss_values = pd.DataFrame(columns=['Epoch', 'Batch', 'Loss'])
         iterator = tqdm(range(self.epochs), disable=(not self.verbose))
         if self.verbose:
-            iterator_description = 'Loss: {loss:.3f}'
-            iterator.set_description(iterator_description.format(loss=0))
+            iterator_description = 'Loss: {loss}'
+            iterator.set_description(iterator_description.format(loss=_format_score(0)))
 
         for i in iterator:
             loss_values = []
@@ -205,7 +205,7 @@ def fit(self, train_data, discrete_columns=()):
 
             if self.verbose:
                 iterator.set_description(
-                    iterator_description.format(loss=loss.detach().cpu().item())
+                    iterator_description.format(loss=_format_score(loss.detach().cpu().item()))
                 )
 
     @random_state

tests/integration/synthesizer/test_tvae.py

@@ -120,5 +120,5 @@ def test_tvae_save(tmpdir, capsys):
     assert len(loss_values) == 10
     assert set(loss_values.columns) == {'Epoch', 'Batch', 'Loss'}
     assert all(loss_values['Batch'] == 0)
-    last_loss_val = loss_values['Loss'].iloc[-1]
-    assert f'Loss: {round(last_loss_val, 3):.3f}: 100%' in captured_out
+    last_loss_val = max(-99.99, min(99.99, loss_values['Loss'].iloc[-1]))
+    assert f'Loss: {last_loss_val:+06.2f}: 100%' in captured_out

tests/unit/synthesizer/test_tvae.py

@@ -60,6 +60,6 @@ def mock_add(a, b):
 
         # Assert
         tqdm_mock.assert_called_once_with(range(epochs), disable=False)
-        assert iterator_mock.set_description.call_args_list[0] == call('Loss: 0.000')
-        assert iterator_mock.set_description.call_args_list[1] == call('Loss: 1.235')
+        assert iterator_mock.set_description.call_args_list[0] == call('Loss: +00.00')
+        assert iterator_mock.set_description.call_args_list[1] == call('Loss: +01.23')
         assert iterator_mock.set_description.call_count == 2