Added Progressbar, REPLACED stream to logger

main/api.py

@@ -58,11 +58,12 @@ class LLMApi:
             model_path = self.models_path / model_name.split('/')[-1]
 
             # Download and save model
-            self.logger.info(f"Downloading model: {model_name}")
+            self.logger.enable_stream_to_logger()
             model = AutoModelForCausalLM.from_pretrained(model_name, progress_bar=True)
 
-            self.logger.info(f"Downloading tokenizer for model: {model_name}")
+            # Download and save tokenizer
             tokenizer = AutoTokenizer.from_pretrained(model_name, progress_bar=True)
+            self.logger.disable_stream_to_logger()
 
             self.logger.info(f"Saving model to {model_path}")
             model.save_pretrained(model_path)

main/utils/logging.py

@@ -1,6 +1,31 @@
 import logging
+import sys
 from pathlib import Path
 
+class StreamToLogger:
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+    def __init__(self, logger, log_level=logging.INFO):
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+        self.enabled = True
+
+    def write(self, buf):
+        if self.enabled:
+            for line in buf.rstrip().splitlines():
+                self.logger.log(self.log_level, line.rstrip())
+
+    def flush(self):
+        pass
+
+    def enable(self):
+        self.enabled = True
+
+    def disable(self):
+        self.enabled = False
+
 def setup_logger(config: dict, name: str = None) -> logging.Logger:
     """Set up logger with configuration from config file."""
     logger = logging.getLogger(name or __name__)
@@ -26,4 +51,12 @@ def setup_logger(config: dict, name: str = None) -> logging.Logger:
     logger.addHandler(file_handler)
     logger.addHandler(console_handler)
 
+    # Redirect stdout to logger
+    stream_to_logger = StreamToLogger(logger, logging.INFO)
+    sys.stdout = stream_to_logger
+
+    # Add methods to enable/disable StreamToLogger
+    logger.enable_stream_to_logger = stream_to_logger.enable
+    logger.disable_stream_to_logger = stream_to_logger.disable
+
     return logger