reinstate single generation injection in CLI, before run is started

NVIDIA · Feb 19, 2025 · dc39223 · dc39223
1 parent 0fc7c84
commit dc39223
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/garak/cli.py b/garak/cli.py
@@ -554,6 +554,19 @@ def main(arguments=None) -> None:
                         msg_list = ",".join(rejected)
                         raise ValueError(f"❌Unknown {spec_namespace}❌: {msg_list}")
 
+            # configure generations counts for main run
+            for probe in parsed_specs["probe"]:
+                # distribute `generations` to the probes
+                p_type, p_module, p_klass = probe.split(".")
+                if (
+                    hasattr(_config.run, "generations")
+                    and _config.run.generations
+                    is not None  # garak.core.yaml always provides run.generations
+                ):
+                    _config.plugins.probes[p_module][p_klass][
+                        "generations"
+                    ] = _config.run.generations
+
             # generator init
             from garak import _plugins
 
@@ -579,9 +592,6 @@ def main(arguments=None) -> None:
             if _config.run.policy_scan:
                 command.run_policy_scan(generator, _config)
 
-            # configure generations counts for main run
-            command.distribute_generations_config(parsed_specs["probe"], _config)
-
             # set up plugins for main run
             # instantiate evaluator
             evaluator = garak.evaluators.ThresholdEvaluator(_config.run.eval_threshold)