diff --git a/hate_crack/attacks.py b/hate_crack/attacks.py index 4acb2dd..f24de78 100644 --- a/hate_crack/attacks.py +++ b/hate_crack/attacks.py @@ -574,13 +574,20 @@ def passgpt_attack(ctx: Any) -> None: if not base: base = default_model + from hate_crack.passgpt_train import _detect_device + + detected = _detect_device() + device_labels = {"cuda": "cuda", "mps": "mps (Apple Silicon)", "cpu": "cpu"} + device_options = ["cuda", "mps", "cpu"] print("\n\tSelect training device:") - print("\t (1) cuda (Recommended)") - print("\t (2) mps (Apple Silicon)") - print("\t (3) cpu") - device_choice = input("\n\tDevice [1]: ").strip() - device_map = {"1": "cuda", "2": "mps", "3": "cpu", "": "cuda"} - device = device_map.get(device_choice, "cuda") + for i, dev in enumerate(device_options, 1): + label = device_labels[dev] + suffix = " (detected)" if dev == detected else "" + print(f"\t ({i}) {label}{suffix}") + default_idx = device_options.index(detected) + 1 + device_choice = input(f"\n\tDevice [{default_idx}]: ").strip() + device_map = {"1": "cuda", "2": "mps", "3": "cpu", "": detected} + device = device_map.get(device_choice, detected) result = ctx.hcatPassGPTTrain(training_file, base, device=device) if result is None: diff --git a/tests/test_passgpt_attack.py b/tests/test_passgpt_attack.py index 58d532d..fd81e02 100644 --- a/tests/test_passgpt_attack.py +++ b/tests/test_passgpt_attack.py @@ -261,11 +261,14 @@ class TestPassGPTAttackHandler: ctx.select_file_with_autocomplete.return_value = "/tmp/wordlist.txt" ctx.hcatPassGPTTrain.return_value = "/home/user/.hate_crack/passgpt/wordlist" - # "T" for train, "" for default base model, "" for default device (cuda), "" for default max candidates + # "T" for train, "" for default base model, "" for default device (auto-detected), "" for default max candidates inputs = iter(["T", "", "", ""]) with ( patch("builtins.input", side_effect=inputs), patch("hate_crack.attacks.os.path.isdir", return_value=False), + patch( + "hate_crack.passgpt_train._detect_device", return_value="cuda" + ), ): from hate_crack.attacks import passgpt_attack @@ -283,11 +286,14 @@ class TestPassGPTAttackHandler: ctx.select_file_with_autocomplete.return_value = "/tmp/wordlist.txt" ctx.hcatPassGPTTrain.return_value = None - # "T" for train, "" for default base model, "" for default device (cuda) + # "T" for train, "" for default base model, "" for default device (auto-detected) inputs = iter(["T", "", ""]) with ( patch("builtins.input", side_effect=inputs), patch("hate_crack.attacks.os.path.isdir", return_value=False), + patch( + "hate_crack.passgpt_train._detect_device", return_value="cuda" + ), ): from hate_crack.attacks import passgpt_attack