Added Progressbar, removed stream to logger

main/api.py

@@ -58,12 +58,11 @@ class LLMApi:
             model_path = self.models_path / model_name.split('/')[-1]
 
             # Download and save model
-            self.logger.enable_stream_to_logger()
+            self.logger.info(f"Downloading model: {model_name}")
             model = AutoModelForCausalLM.from_pretrained(model_name, progress_bar=True)
 
-            # Download and save tokenizer
+            self.logger.info(f"Downloading tokenizer for model: {model_name}")
             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,31 +1,6 @@
 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__)
@@ -51,12 +26,4 @@ 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