uxoxo/eb2ab

1from queue import Queue, Empty2import time3import logging4 5 6class RedirectConsole:7    def __init__(self, log_buffer: Queue, real_output):8        self.log_buffer = log_buffer  # Queue buffer for the log9        self.real_output = real_output  # Real terminal (sys.__stdout__ or sys.__stderr__)10 11        # Setup for transformers logging12        self.setup_transformers_logger()13 14    def write(self, message: str):15        # Write to the real terminal16        self.real_output.write(message)17        self.real_output.flush()18 19        # Write to the log buffer20        self.log_buffer.put(message)21 22    def flush(self):23        self.real_output.flush()24 25    def isatty(self) -> bool:26        return self.real_output.isatty()27 28    def poll_logs(self, stop_event):29        logs = ""30        errors = ""31        while not stop_event.is_set() or not self.log_buffer.empty():32            try:33                # Read logs from the buffer without blocking34                log = self.log_buffer.get_nowait()35                if "An error occurred" in log:36                    errors += log  # Capture error messages separately37                logs += log38            except Empty:39                pass  # No logs in the buffer40            yield logs, errors  # Yield updated logs and errors41            time.sleep(0.1)  # Prevent tight looping42 43    def setup_transformers_logger(self):44        # Configure the `transformers` logger45        transformers_logger = logging.getLogger("transformers")46        transformers_logger.setLevel(logging.WARNING)  # Capture warnings and above47 48        # Create a handler that writes to this instance49        handler = logging.StreamHandler(self)50        handler.setFormatter(logging.Formatter("%(message)s"))  # Simplified format51        transformers_logger.addHandler(handler)