diff --git a/src/heretic/main.py b/src/heretic/main.py index 1262c221..693c3d0a 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -131,26 +131,31 @@ def obtain_merge_strategy(settings: Settings, model: Model) -> str | None: ) print() - strategy = prompt_select( - "How do you want to proceed?", - choices=[ - Choice( - title="Merge LoRA into full model" - + ( - "" - if settings.quantization == QuantizationMethod.NONE - else " (requires sufficient RAM)" + strategy = prompt_select( + "How do you want to proceed?", + choices=[ + Choice( + title="Merge LoRA into full model" + + ( + "" + if settings.quantization == QuantizationMethod.NONE + else " (requires sufficient RAM)" + ), + value="merge", ), - value="merge", - ), - Choice( - title="Save LoRA adapter only (can be merged later)", - value="adapter", - ), - ], - ) + Choice( + title="Cancel", + value="cancel", + ), + ], + ) + + if strategy == "cancel": + return None - return strategy + return strategy + else: + return "merge" def run():