diff --git a/src/heretic/main.py b/src/heretic/main.py index 241415b..fabfde3 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -1,8 +1,20 @@ # SPDX-License-Identifier: AGPL-3.0-or-later # Copyright (C) 2025-2026 Philipp Emanuel Weidmann + contributors +import os import sys +# Enable expandable segments to reduce memory fragmentation on multi-GPU setups. +if ( + "PYTORCH_ALLOC_CONF" not in os.environ + and "PYTORCH_CUDA_ALLOC_CONF" not in os.environ +): + os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" + +# Silence warning spam from PyTorch's C++ logging. +if "TORCH_CPP_LOG_LEVEL" not in os.environ: + os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR" + # Ensure standard output/error use UTF-8 instead of system default charmap (e.g. cp1252 on Windows). for stream in (sys.stdout, sys.stderr): if ( @@ -11,6 +23,7 @@ for stream in (sys.stdout, sys.stderr): ): stream.reconfigure(encoding="utf-8") # type: ignore + from .config import Settings @@ -23,6 +36,7 @@ def _is_help_invocation() -> bool: if _is_help_invocation(): Settings() + # FIXME: Rich progress bars are currently disabled because of rendering issues # when used from multiple threads in parallel (e.g. by huggingface_hub). """ @@ -35,7 +49,6 @@ patch_tqdm() import logging import math -import os import random import re import time @@ -180,13 +193,6 @@ def obtain_export_strategy( def run(): - # Enable expandable segments to reduce memory fragmentation on multi-GPU setups. - if ( - "PYTORCH_ALLOC_CONF" not in os.environ - and "PYTORCH_CUDA_ALLOC_CONF" not in os.environ - ): - os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" - # Modified "Pagga" font from https://budavariam.github.io/asciiart-text/ print(f"[cyan]█░█░█▀▀░█▀▄░█▀▀░▀█▀░█░█▀▀[/] v{version('heretic-llm')}") print(