Update load_data.py

load_data.py

@@ -9,7 +9,7 @@ from datasets import load_dataset, concatenate_datasets
 
 from argilla.listeners import listener
 
-HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN
+HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN from space env variables
 HUB_DATASET_NAME = os.environ.get('HUB_DATASET_NAME') #get dataset name
 
 
@@ -24,7 +24,7 @@ def save_validated_to_hub(records, ctx):
         if HF_TOKEN:
             print("Pushing the dataset")
             print(ds)
-            ds.push_to_hub(HUB_DATASET_NAME, token=HF_TOKEN)
+            ds.push_to_hub(HUB_DATASET_NAME, token=HF_TOKEN, private=True, revision = 'validated') #push annotations on the validated branch of the dataset
         else:
             print("SET HF_TOKEN and HUB_DATASET_NAME TO SYNC YOUR DATASET!!!")
     else:
@@ -35,17 +35,17 @@ class LoadDatasets:
         rg.init(api_key=api_key, workspace=workspace)
 
     @staticmethod
-    def load_somos():
-        # Leer el dataset del Hub
+    def load_saved_data():
+        # load data from validated branch
         try:
             print(f"Trying to sync with {HUB_DATASET_NAME}")
-            old_ds = load_dataset(HUB_DATASET_NAME, split="train")
+            old_ds = load_dataset(HUB_DATASET_NAME, split="train", use_auth_token=HF_TOKEN, revision='validated')
         except Exception as e:
             print(f"Not possible to sync with {HUB_DATASET_NAME}")
             print(e)
             old_ds = None
             
-        dataset = load_dataset("LukasGe/JOB_TITLES", split="train")
+        dataset = load_dataset({HUB_DATASET_NAME}, split="train", use_auth_token=HF_TOKEN, revision='main')  # get starting file from the main branch of the dataset
     
         
         if old_ds:
@@ -85,7 +85,7 @@ if __name__ == "__main__":
                 response = requests.get("http://0.0.0.0:6900/")
                 if response.status_code == 200:
                     ld = LoadDatasets(API_KEY)
-                    ld.load_somos()
+                    ld.load_saved_data()
                     break
 
             except requests.exceptions.ConnectionError: