Fix Zeno Visualizer (#2227)

* fix: arguments data

* fix based on comment

* Update zeno_visualize.py

updated all output types

---------

Co-authored-by: Baber Abbasi <92168766+baberabb@users.noreply.github.com>

Author: namtranase

scripts/zeno_visualize.py

@@ -173,19 +173,19 @@ def generate_dataset(
     instance = [""] * len(ids)
 
     if config["output_type"] == "loglikelihood":
-        instance = [x["arguments"][0][0] for x in data]
-        labels = [x["arguments"][0][1] for x in data]
+        instance = [x["arguments"]["gen_args_0"]["arg_0"] for x in data]
+        labels = [x["arguments"]["gen_args_0"]["arg_1"] for x in data]
     elif config["output_type"] == "multiple_choice":
         instance = [
-            x["arguments"][0][0]
+            x["arguments"]["gen_args_0"]["arg_0"]
             + "\n\n"
             + "\n".join([f"- {y[1]}" for y in x["arguments"]])
             for x in data
         ]
     elif config["output_type"] == "loglikelihood_rolling":
-        instance = [x["arguments"][0][0] for x in data]
+        instance = [x["arguments"]["gen_args_0"]["arg_0"] for x in data]
     elif config["output_type"] == "generate_until":
-        instance = [x["arguments"][0][0] for x in data]
+        instance = [x["arguments"]["gen_args_0"]["arg_0"] for x in data]
 
     return pd.DataFrame(
         {