Spaces:
Build error
Build error
File size: 1,803 Bytes
aa7ea23 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | from queue import Queue, Empty
import time
import logging
class RedirectConsole:
def __init__(self, log_buffer: Queue, real_output):
self.log_buffer = log_buffer # Queue buffer for the log
self.real_output = real_output # Real terminal (sys.__stdout__ or sys.__stderr__)
# Setup for transformers logging
self.setup_transformers_logger()
def write(self, message: str):
# Write to the real terminal
self.real_output.write(message)
self.real_output.flush()
# Write to the log buffer
self.log_buffer.put(message)
def flush(self):
self.real_output.flush()
def isatty(self) -> bool:
return self.real_output.isatty()
def poll_logs(self, stop_event):
logs = ""
errors = ""
while not stop_event.is_set() or not self.log_buffer.empty():
try:
# Read logs from the buffer without blocking
log = self.log_buffer.get_nowait()
if "An error occurred" in log:
errors += log # Capture error messages separately
logs += log
except Empty:
pass # No logs in the buffer
yield logs, errors # Yield updated logs and errors
time.sleep(0.1) # Prevent tight looping
def setup_transformers_logger(self):
# Configure the `transformers` logger
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING) # Capture warnings and above
# Create a handler that writes to this instance
handler = logging.StreamHandler(self)
handler.setFormatter(logging.Formatter("%(message)s")) # Simplified format
transformers_logger.addHandler(handler) |