use colortimelog

setup.py

@@ -15,7 +15,7 @@
 
 import setuptools
 
-VERSION = '0.1.0'
+VERSION = '0.1.1'
 
 with open('README.md', 'r') as file_object:
   LONG_DESCRIPTION = file_object.read()

tests/integration_test.py

@@ -109,8 +109,8 @@ def test_four_clusters(self):
     predicted_label = model.predict(test_sequence, inference_args)
 
     # run evaluation
-    model.logger.print(
-        3, 'Asserting the equivalence between'
+    model.logger.info(
+        'Asserting the equivalence between'
         '\nGround truth: {}\nPredicted: {}'.format(
             test_cluster_id, predicted_label))
     accuracy = uisrnn.compute_sequence_match_accuracy(
@@ -125,8 +125,8 @@ def test_four_clusters(self):
     predicted_label = loaded_model.predict(test_sequence, inference_args)
 
     # run evaluation with loaded model
-    model.logger.print(
-        3, 'Asserting the equivalence between'
+    model.logger.info(
+        'Asserting the equivalence between'
         '\nGround truth: {}\nPredicted: {}'.format(
             test_cluster_id, predicted_label))
     accuracy = uisrnn.compute_sequence_match_accuracy(
@@ -140,13 +140,13 @@ def test_four_clusters(self):
     model.fit(train_sequence[:100, :], train_cluster_id[:100], training_args)
     transition_bias_2 = model.transition_bias
     self.assertNotAlmostEqual(transition_bias_1, transition_bias_2)
-    model.logger.print(
-        3, 'Asserting transition_bias changed from {} to {}'.format(
+    model.logger.info(
+        'Asserting transition_bias changed from {} to {}'.format(
             transition_bias_1, transition_bias_2))
 
     # run evaluation
-    model.logger.print(
-        3, 'Asserting the equivalence between'
+    model.logger.info(
+        'Asserting the equivalence between'
         '\nGround truth: {}\nPredicted: {}'.format(
             test_cluster_id, predicted_label))
     accuracy = uisrnn.compute_sequence_match_accuracy(

uisrnn/arguments.py

@@ -87,12 +87,12 @@ def parse_arguments():
            'value is None, we will estimate it from training data.')
   model_parser.add_argument(
       '--verbosity',
-      default=2,
+      default=3,
       type=int,
       help='How verbose will the logging information be. Higher value '
       'represents more verbose information. A general guideline: '
-      '0 for errors; 1 for finishing important steps; '
-      '2 for finishing less important steps; 3 or above for debugging '
+      '0 for fatals; 1 for errors; 2 for finishing important steps; '
+      '3 for finishing less important steps; 4 or above for debugging '
       'information.')
   model_parser.add_argument(
       '--enable_cuda',

uisrnn/uisrnn.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """The UIS-RNN model."""
 
+import colortimelog
 import functools
 import numpy as np
 import torch
@@ -103,7 +104,7 @@ def __init__(self, args):
     self.transition_bias = args.transition_bias
     self.transition_bias_denominator = 0.0
     self.crp_alpha = args.crp_alpha
-    self.logger = utils.Logger(args.verbosity)
+    self.logger = colortimelog.Logger(args.verbosity)
 
   def _get_optimizer(self, optimizer, learning_rate):
     """Get optimizer for UISRNN.

uisrnn/utils.py

@@ -21,23 +21,6 @@
 from torch import autograd
 
 
-class Logger:
-  """A class for printing logging information to screen."""
-
-  def __init__(self, verbosity):
-    self._verbosity = verbosity
-
-  def print(self, level, message):
-    """Print a message if level is not higher than verbosity.
-
-    Args:
-      level: the level of this message, smaller value means more important
-      message: the message to be printed
-    """
-    if level <= self._verbosity:
-      print(message)
-
-
 def generate_random_string(length=6):
   """Generate a random string of upper case letters and digits.