Get dataset abstraction

main.py

@@ -92,15 +92,7 @@ def run_main_continuously():
     # Calculate the start time for running the main_function every day.
     start_time = datetime.now().time()
 
-    # Load the existing dataset from the Hugging Face hub or create a new one
-    try:
-        dataset = load_dataset(dataset_name, download_mode="reuse_cache_if_exists", ignore_verifications=True)
-        logger.debug("Loading existing dataset")
-        if "__index_level_0__" in dataset["all_days"].column_names:
-            dataset = dataset.remove_columns(["__index_level_0__"])
-    except FileNotFoundError:
-        logger.warning("Creating new dataset")
-        dataset = DatasetDict()
+    dataset = get_dataset()
 
     while True:
         today = datetime.now().date()
@@ -119,5 +111,18 @@ def run_main_continuously():
             time.sleep(wait_until_tomorrow)
 
 
+def get_dataset():
+    # Load the existing dataset from the Hugging Face hub or create a new one
+    try:
+        dataset = load_dataset(dataset_name, download_mode="reuse_cache_if_exists", ignore_verifications=True)
+        logger.debug("Loading existing dataset")
+        if "__index_level_0__" in dataset["all_days"].column_names:
+            dataset = dataset.remove_columns(["__index_level_0__"])
+    except FileNotFoundError:
+        logger.warning("Creating new dataset")
+        dataset = DatasetDict()
+    return dataset
+
+
 if __name__ == '__main__':
     run_main_continuously()