Clean logging, Refactor

lerobot/common/utils.py

@@ -1,4 +1,6 @@
+import logging
 import random
+from datetime import datetime
 
 import numpy as np
 import torch
@@ -10,3 +12,34 @@ def set_seed(seed):
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
+
+
+def init_logging():
+    def custom_format(record):
+        dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        fnameline = f"{record.pathname}:{record.lineno}"
+        message = f"{record.levelname} {dt} {fnameline[-15:]:>15} {record.msg}"
+        return message
+
+    logging.basicConfig(level=logging.INFO)
+
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+
+    formatter = logging.Formatter()
+    formatter.format = custom_format
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logging.getLogger().addHandler(console_handler)
+
+
+def format_number_KMB(num):
+    suffixes = ["", "K", "M", "B", "T", "Q"]
+    divisor = 1000.0
+
+    for suffix in suffixes:
+        if abs(num) < divisor:
+            return f"{num:.0f}{suffix}"
+        num /= divisor
+
+    return num