feat: add training time estimates and device selection to PassGPT menu

Show estimated training times for CUDA/MPS/CPU before starting a
training run. Add device selection prompt with cuda as the default.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Justin Bollinger
2026-02-18 11:27:09 -05:00
co-authored by Claude Opus 4.6
parent e1d6922edc
commit 16860a51d0
3 changed files with 23 additions and 6 deletions
+15 -1
View File
@@ -557,6 +557,11 @@ def passgpt_attack(ctx: Any) -> None:
if choice.upper() == "T":
print("\n\tTrain a new PassGPT model")
print("\n\t--- Estimated Training Times (14M passwords, 3 epochs) ---")
print("\t CUDA (RTX 3090/4090): 1-3 hours")
print("\t MPS (Apple Silicon): 6-12 hours")
print("\t CPU: Very slow (not recommended)")
print("\t Use --max-lines to reduce training data for faster runs.")
training_file = ctx.select_file_with_autocomplete(
"Select training wordlist", base_dir=ctx.hcatWordlists
)
@@ -568,7 +573,16 @@ def passgpt_attack(ctx: Any) -> None:
base = input(f"\n\tBase model ({default_model}): ").strip()
if not base:
base = default_model
result = ctx.hcatPassGPTTrain(training_file, base)
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")
result = ctx.hcatPassGPTTrain(training_file, base, device=device)
if result is None:
print("\n\tTraining failed. Returning to menu.")
return
+3 -1
View File
@@ -2297,7 +2297,7 @@ def _passgpt_model_dir():
# PassGPT Attack - Fine-tune a model on a custom wordlist
def hcatPassGPTTrain(training_file, base_model=None):
def hcatPassGPTTrain(training_file, base_model=None, device=None):
training_file = os.path.abspath(training_file)
if not os.path.isfile(training_file):
print(f"Error: Training file not found: {training_file}")
@@ -2321,6 +2321,8 @@ def hcatPassGPTTrain(training_file, base_model=None):
"--output-dir",
output_dir,
]
if device:
cmd.extend(["--device", device])
print(f"[*] Running: {_format_cmd(cmd)}")
proc = subprocess.Popen(cmd)
try:
+5 -4
View File
@@ -258,8 +258,8 @@ 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 max candidates
inputs = iter(["T", "", ""])
# "T" for train, "" for default base model, "" for default device (cuda), "" for default max candidates
inputs = iter(["T", "", "", ""])
with (
patch("builtins.input", side_effect=inputs),
patch("hate_crack.attacks.os.path.isdir", return_value=False),
@@ -269,7 +269,7 @@ class TestPassGPTAttackHandler:
passgpt_attack(ctx)
ctx.hcatPassGPTTrain.assert_called_once_with(
"/tmp/wordlist.txt", "javirandor/passgpt-10characters"
"/tmp/wordlist.txt", "javirandor/passgpt-10characters", device="cuda"
)
ctx.hcatPassGPT.assert_called_once()
call_kwargs = ctx.hcatPassGPT.call_args
@@ -280,7 +280,8 @@ class TestPassGPTAttackHandler:
ctx.select_file_with_autocomplete.return_value = "/tmp/wordlist.txt"
ctx.hcatPassGPTTrain.return_value = None
inputs = iter(["T", ""])
# "T" for train, "" for default base model, "" for default device (cuda)
inputs = iter(["T", "", ""])
with (
patch("builtins.input", side_effect=inputs),
patch("hate_crack.attacks.os.path.isdir", return_value=False),