Update helper.py

helper.py

@@ -47,22 +47,19 @@ def encode_query(query: Union[str, Image.Image]) -> torch.Tensor:
 
 def load_hf_datasets(dataset_name):
     """
-    Load all splits containing 'Main' from a Hugging Face dataset as a DataFrame
-    ---------------------------------------------------------------------------
+    Load Datasets from Hugging Face as DF
+    ---------------------------------------
     dataset_name: str - name of dataset on Hugging Face
-    ---------------------------------------------------------------------------
-    RETURNS: concatenated dataset as a pandas DataFrame
+    ---------------------------------------
+
+    RETURNS: dataset as pandas dataframe
     """
     dataset = load_dataset(f"quasara-io/{dataset_name}")
-    
-    # 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
+    # Access only the 'Main' split
+    main_dataset = dataset['Main_1']
+    # Convert to Pandas DataFrame
+    df = main_dataset.to_pandas()
+    return df
     
 def get_image_vectors(df):
     # Get the image vectors from the dataframe