convert mutation_generator to use load_plugin (#795)

* convert mutation_generator to use `load_plugin` Signed-off-by: Jeffrey Martin <[email protected]> * use generator device for get_score_autodan Signed-off-by: Jeffrey Martin <[email protected]> --------- Signed-off-by: Jeffrey Martin <[email protected]>
NVIDIA · Jul 22, 2024 · 8f7eaa4 · 8f7eaa4
1 parent e8bb254
commit 8f7eaa4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/garak/resources/autodan/autodan.py b/garak/resources/autodan/autodan.py
@@ -9,7 +9,8 @@
 
 import gc
 
-from garak.generators import Generator, load_generator
+from garak._plugins import load_plugin
+from garak.generators import Generator
 from garak.generators.huggingface import Model
 import garak._config
 from garak.resources.autodan.genetic import (
@@ -133,10 +134,12 @@ def autodan_generate(
     crit = nn.CrossEntropyLoss(reduction="mean")
 
     config_root = {
-        {mutation_generator_type: {"name": mutation_generator_name, "generations": 1}}
+        "generators": {
+            mutation_generator_type: {"name": mutation_generator_name, "generations": 1}
+        }
     }
-    mutation_generator = load_generator(
-        model_type=mutation_generator_type, config=config_root
+    mutation_generator = load_plugin(
+        "generators." + mutation_generator_type, config_root=config_root
     )
 
     # Feel like this could just be text instead of storing it as tensors.

diff --git a/garak/resources/autodan/genetic.py b/garak/resources/autodan/genetic.py
@@ -594,7 +594,7 @@ def get_score_autodan(
     losses = []
     input_ids_list = []
     target_slices = []
-    device = generator.device if generator.device >= 0 else "cpu"
+    device = generator.device
     for item in test_controls:
         prefix_manager = AutoDanPrefixManager(
             generator=generator,