Update app.py

app.py

@@ -99,12 +99,22 @@ max_seq_len = 25
 logging.basicConfig(filename="app.log", level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 # Dataset setup
-FLAGGED_DATASET_ID = "InfinitodeLTD/DungenDev-FlaggedOutputs"  # Replace with your dataset ID
-try:
-    flagged_dataset = load_dataset(FLAGGED_DATASET_ID)
-except:
-    flagged_dataset = Dataset.from_dict({"Timestamp": [], "Prompt": [], "Flagged Text": []})
-    flagged_dataset.push_to_hub(FLAGGED_DATASET_ID)
+FLAGGED_DATASET_ID = "InfinitodeLTD/DungenDev-FlaggedOutputs"
+
+def load_or_create_dataset(dataset_id):
+    try:
+        dataset = load_dataset(dataset_id)
+        if "flagged_data" not in dataset:
+            raise ValueError("Dataset does not contain the 'flagged_data' config.")
+
+        return dataset["flagged_data"]
+    except (datasets.DatasetNotFoundError, ValueError) as e:
+        logging.warning(f"Dataset not found or incorrect schema: {e}. Creating a new dataset.")
+        dataset = Dataset.from_dict({"Timestamp": [], "Prompt": [], "Flagged Text": []})
+        dataset.push_to_hub(dataset_id, config_name="flagged_data") # important: config_name
+        return dataset
+
+flagged_dataset = load_or_create_dataset(FLAGGED_DATASET_ID)
 
 def generate_text(seed_text, next_words=30, temperature=0.5):
     seed_text = seed_text.strip().lower()