Add ConversationalPipeline for huggingface models

garak/generators/huggingface.py

@@ -179,6 +179,7 @@ def __init__(self, name, do_sample=True, generations=10, device=0):
             do_sample=do_sample,
             device=device,
         )
+        self.device = device
         self.conversation = Conversation()
         self.deprefix_prompt = name in models_to_deprefix
         if _config.loaded:

garak/resources/tap/generator_utils.py

@@ -4,7 +4,7 @@
 from typing import Union
 
 from garak.generators.openai import chat_models, OpenAIGenerator
-from garak.generators.huggingface import Model
+from garak.generators.huggingface import ConversationalPipeline
 
 supported_openai = chat_models
 supported_huggingface = [
@@ -51,15 +51,15 @@ def load_generator(
             generations=generations,
         )
     elif model_name in supported_huggingface:
-        generator = Model(model_name, generations=generations, device=device)
+        generator = ConversationalPipeline(model_name, generations=generations, device=device)
     else:
         msg = (
             f"{model_name} is not currently supported for TAP generation. Support is available for the following "
             f"OpenAI and HuggingFace models:\nOpenAI: {supported_openai}\nHuggingFace: {supported_huggingface}\n"
             f"Your jailbreaks will *NOT* be saved."
         )
         print(msg)
-        generator = Model(model_name, generations=generations, device=device)
+        generator = ConversationalPipeline(model_name, generations=generations, device=device)
 
     generator.max_tokens = max_tokens
     if temperature is not None:

garak/resources/tap/tap_main.py

@@ -121,7 +121,7 @@ def get_attack(self, convs, prompts):
                     continue
 
                 for full_prompt in full_prompts_subset[left:right]:
-                    outputs_list.append(self.attack_generator.generate(full_prompt)[0])
+                    outputs_list.append(self.attack_generator.generate(full_prompt))
 
             # Check for valid outputs and update the list
             new_indices_to_regenerate = []