Update app.py

app.py

@@ -142,44 +142,42 @@ def generate_text(seed_text, next_words=30, temperature=0.5):
 
     return decoded #, gr.update(interactive=True)
 
-flagged_outputs = []
-
-def flag_output(prompt, generated_text, next_words, temperature):
-    if not generated_text.strip():
-        return "Cannot flag an empty output."
-
-    timestamp = datetime.now().isoformat()
-
-    dataset_id = "InfinitodeLTD/DungenDev-FlaggedOutputs"
-    
-    # Load the existing dataset (if it exists) or create an empty DatasetDict
-    try:
-        dataset = load_dataset(dataset_id)
-    except:
-        dataset = DatasetDict()
-
-    # Prepare new data to append
-    new_data = [{
-        "Prompt": prompt,
-        "Generated Text": generated_text,
-        "Next Words": next_words,
-        "Temperature": temperature,
-        "Timestamp": timestamp
-    }]
-
-    new_dataset = Dataset.from_list(new_data)
-
-    if "train" in dataset:
-        dataset["train"] = concatenate_datasets([dataset["train"], new_dataset])  # Append to existing train
-    else:
-        dataset["train"] = new_dataset  # Create the train split
-
-    dataset.push_to_hub(dataset_id)
-
-    return "Output flagged successfully."
-
-def enable_flag_button(state):
-    return gr.update(interactive=state)
+# Custom flagging callback to use the existing flagging logic
+class CustomFlaggingCallback(gr.FlaggingCallback):
+    def __init__(self, dataset_id):
+        self.dataset_id = dataset_id
+
+    def flag(self, flag_data, flag_option=None, username=None):
+        prompt, generated_text, next_words, temperature = flag_data
+        timestamp = datetime.now().isoformat()
+        
+        # Custom flagging logic
+        try:
+            dataset = load_dataset(self.dataset_id)
+        except:
+            dataset = DatasetDict()
+
+        new_data = [{
+            "Prompt": prompt,
+            "Generated Text": generated_text,
+            "Next Words": next_words,
+            "Temperature": temperature,
+            "Timestamp": timestamp
+        }]
+
+        new_dataset = Dataset.from_list(new_data)
+
+        if "train" in dataset:
+            dataset["train"] = concatenate_datasets([dataset["train"], new_dataset])  # Append to existing train
+        else:
+            dataset["train"] = new_dataset  # Create the train split
+
+        dataset.push_to_hub(self.dataset_id)
+        return "Output flagged successfully."
+
+# Initialize the custom flagging callback
+dataset_id = "InfinitodeLTD/DungenDev-FlaggedOutputs"
+custom_flag_callback = CustomFlaggingCallback(dataset_id)
 
 demo = gr.Interface(
     fn=generate_text,
@@ -203,9 +201,9 @@ demo = gr.Interface(
         ["an evil character name", 30, 0.5]
     ],
     theme="default",
-    flagging_mode='manual',
-    flagging_options=["Incorrect", "Offensive/Hateful", "Harmful/Violent", "Illegal", "NSFW", "Other"],
-    flagging_dir='flagged_outputs'
+    flagging_mode="manual",
+    flagging_callback=custom_flag_callback
+)
 )
 
 demo.queue()