Updated load hugging face dataset function

helper.py

@@ -47,20 +47,23 @@ def encode_query(query: Union[str, Image.Image]) -> torch.Tensor:
 
 def load_hf_datasets(dataset_name):
     """
-    Load Datasets from Hugging Face as DF
-    ---------------------------------------
+    Load all splits containing 'Main' from a Hugging Face dataset as a DataFrame
+    ---------------------------------------------------------------------------
     dataset_name: str - name of dataset on Hugging Face
-    ---------------------------------------
-
-    RETURNS: dataset as pandas dataframe
+    ---------------------------------------------------------------------------
+    RETURNS: concatenated dataset as a pandas DataFrame
     """
     dataset = load_dataset(f"quasara-io/{dataset_name}")
-    # Access only the 'Main' split
-    main_dataset = dataset['Main']
-    # Convert to Pandas DataFrame
-    df = main_dataset.to_pandas()
-    return df
-
+    
+    # Filter splits that contain the word 'Main'
+    main_splits = [split for split in dataset if 'Main' in split]
+    
+    # Load and concatenate all splits containing 'Main' into a single DataFrame
+    df_list = [dataset[split].to_pandas() for split in main_splits]
+    combined_df = pd.concat(df_list, ignore_index=True)
+    
+    return combined_df
+    
 def get_image_vectors(df):
     # Get the image vectors from the dataframe
     image_vectors = np.vstack(df['Vector'].to_numpy())