Adding "test type" for redteaming notes.

app_dialogue.py

@@ -204,6 +204,7 @@ def model_inference(
     max_new_tokens,
     repetition_penalty,
     top_p,
+    test_type
 ):
     if user_prompt["text"].strip() == "" and not user_prompt["files"]:
         gr.Error("Please input a query and optionally image(s).")
@@ -288,6 +289,7 @@ FEATURES = datasets.Features(
         "max_new_tokens": datasets.Value("int32"),
         "repetition_penalty": datasets.Value("float32"),
         "top_p": datasets.Value("int32"),
+        "test_type": datasets.Value("string")
         }
     )
 
@@ -299,6 +301,7 @@ def flag_dope(
     max_new_tokens,
     repetition_penalty,
     top_p,
+    test_type=""
 ):
     images = []
     conversation = []
@@ -318,6 +321,7 @@ def flag_dope(
         "max_new_tokens": [max_new_tokens],
         "repetition_penalty": [repetition_penalty],
         "top_p": [top_p],
+        "test_type": [test_type]
     }
     try:
         ds = datasets.load_dataset("HuggingFaceM4/dope-dataset-red-teaming", split="train", token=HF_WRITE_TOKEN)
@@ -336,6 +340,7 @@ def flag_problematic(
     max_new_tokens,
     repetition_penalty,
     top_p,
+    test_type=""
 ):
     images = []
     conversation = []
@@ -355,6 +360,7 @@ def flag_problematic(
         "max_new_tokens": [max_new_tokens],
         "repetition_penalty": [repetition_penalty],
         "top_p": [top_p],
+        "test_type": [test_type]
     }
     try:
         ds = datasets.load_dataset("HuggingFaceM4/problematic-dataset-red-teaming", split="train", token=HF_WRITE_TOKEN)
@@ -364,6 +370,8 @@ def flag_problematic(
         hf_dataset = datasets.Dataset.from_dict(data, features=FEATURES)
     hf_dataset.push_to_hub( "HuggingFaceM4/problematic-dataset-red-teaming", split="train", token=HF_WRITE_TOKEN, private=True)
 
+# Change this to say whatever in order to have notes on the specific redteaming test added to the datasets.
+test_type = gr.Textbox(label="Optional: Type of redteaming test.")
 
 # Hyper-parameters for generation
 max_new_tokens = gr.Slider(
@@ -517,6 +525,7 @@ with gr.Blocks(
             max_new_tokens,
             repetition_penalty,
             top_p,
+            test_type,
         ],
         outputs=None,
         preprocess=False,
@@ -531,6 +540,7 @@ with gr.Blocks(
             max_new_tokens,
             repetition_penalty,
             top_p,
+            test_type,
         ],
         outputs=None,
         preprocess=False,