Save first predicted label for better display in HF Dataset Viewer

app.py

@@ -332,6 +332,8 @@ with gr.Blocks(theme=my_theme, css=custom_css, head=font_awesome_html) as demo:
         """
         Append input/outputs and user feedback to a JSON Lines file.
         """
+        # Get the first label (prediction with highest probability)
+        prediction = {k: v for k, v in [next(iter(label.items()))]}
         with USER_FEEDBACK_PATH.open("a") as f:
             f.write(
                 json.dumps(
@@ -339,7 +341,7 @@ with gr.Blocks(theme=my_theme, css=custom_css, head=font_awesome_html) as demo:
                         "claim": claim,
                         "evidence": evidence,
                         "model": model,
-                        "prediction": label,
+                        "prediction": prediction,
                         "user_label": user_label,
                         "datetime": datetime.now().isoformat(),
                     }