Update load_data.py

load_data.py

@@ -2,21 +2,33 @@ import sys
 import time
 import os
 
-import argilla as rg
 import pandas as pd
 import requests
 from datasets import load_dataset, concatenate_datasets
 
+import argilla as rg
 from argilla.listeners import listener
 
+### Configuration section ###
+
+# needed for pushing the validated data to HUB_DATASET_NAME
 HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# The source dataset to read Alpaca translated examples
 SOURCE_DATASET = "LEL-A/translated_german_alpaca"
+
+# The name of the dataset in Argilla
 RG_DATASET_NAME = "translated-german-alpaca"
+
+# The name of the Hub dataset to push the validations every 20 min and keep the dataset synced
 HUB_DATASET_NAME = os.environ.get('HUB_DATASET_NAME', f"{SOURCE_DATASET}_validation")
 
+# The labels for the tasks (they can be extended if needed)
+LABELS = ["BAD INSTRUCTION", "BAD INPUT", "BAD OUTPUT", "INAPPROPRIATE", "BIASED", "ALL GOOD"]
+
 @listener(
     dataset=RG_DATASET_NAME, 
-    query="status:Validated", # https://docs.argilla.io/en/latest/guides/features/queries.html
+    query="status:Validated", 
     execution_interval_in_seconds=1200, # interval to check the execution of `save_validated_to_hub`
 )
 def save_validated_to_hub(records, ctx):
@@ -60,7 +72,7 @@ class LoadDatasets:
         records = rg.DatasetForTextClassification.from_datasets(dataset)
 
         settings = rg.TextClassificationSettings(
-            label_schema=["BAD INSTRUCTION", "BAD INPUT", "BAD OUTPUT", "INAPPROPRIATE", "BIASED", "ALL GOOD"]
+            label_schema=LABELS
         )
         
         print(f"Configuring dataset: {RG_DATASET_NAME}")