From 3d5e645c13a364792ecd8219425ba0fcd1ed8270 Mon Sep 17 00:00:00 2001 From: Philipp Emanuel Weidmann Date: Sun, 12 Oct 2025 12:53:40 +0530 Subject: [PATCH] Handle Ctrl+C gracefully --- src/heretic/main.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index 933e90e..687fec0 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -22,6 +22,7 @@ from accelerate.utils import ( ) from huggingface_hub import ModelCard from pydantic import ValidationError +from rich.traceback import install from .config import Settings from .evaluator import Evaluator @@ -29,7 +30,7 @@ from .model import Model from .utils import get_readme_intro, load_prompts, print -def main(): +def run(): # Modified "Pagga" font from https://budavariam.github.io/asciiart-text/ print(f"[cyan]█░█░█▀▀░█▀▄░█▀▀░▀█▀░█░█▀▀[/] v{version('heretic')}") print("[cyan]█▀█░█▀▀░█▀▄░█▀▀░░█░░█░█░░[/]") @@ -377,3 +378,22 @@ def main(): except Exception as error: print(f"[red]Error: {error}[/]") + + +def main(): + # Install Rich traceback handler. + install() + + try: + run() + except BaseException as error: + # Transformers appears to handle KeyboardInterrupt (or BaseException) + # internally in some places, which can re-raise a different error in the handler, + # masking the root cause. We therefore check both the error itself and its context. + if isinstance(error, KeyboardInterrupt) or isinstance( + error.__context__, KeyboardInterrupt + ): + print() + print("[red]Shutting down...[/]") + else: + raise