Updated helper with new loading function

helper.py

@@ -46,49 +46,24 @@ def encode_query(query: Union[str, Image.Image]) -> torch.Tensor:
     
     return query_embedding
 
-def load_hf_datasets(key,dataset):
+def load_dataset_with_limit(dataset_name, dataset_subset, search_in_small_objects,limit=1000):
     """
-    Load Datasets from Hugging Face as DF
-    ---------------------------------------
-    dataset_name: str - name of dataset on Hugging Face
-    ---------------------------------------
-    RETURNS: dataset as pandas dataframe
+    Load a dataset from Hugging Face and limit the number of rows.
     """
-    df = dataset[key].to_pandas()
-            
-    return df
-
-def parallel_load_and_combine(dataset_keys, dataset):
-    """
-    Load datasets in parallel and combine Main and Split keys
-    ----------------------------------------------------------
-    dataset_keys: list - keys of the dataset (e.g., ['Main_1', 'Split_1', ...])
-    dataset: DatasetDict - the loaded Hugging Face dataset
-    ----------------------------------------------------------
-    RETURNS: combined DataFrame from both Main and Split keys
-    """
-    # Separate keys into Main and Split lists
-    main_keys = [key for key in dataset_keys if key.startswith('Main')]
-    split_keys = [key for key in dataset_keys if key.startswith('Split')]
-
-    def process_key(key, key_type):
-        df = load_hf_datasets(key, dataset)
-        return df
-
-    # Parallel loading of Main keys
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        main_dfs = list(executor.map(lambda key: process_key(key, 'Main'), main_keys))
-
-    # Parallel loading of Split keys
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        split_dfs = list(executor.map(lambda key: process_key(key, 'Split'), split_keys))
-
-    # Combine Main DataFrames and Split DataFrames
-    main_combined_df = pd.concat(main_dfs, ignore_index=True) if main_dfs else pd.DataFrame()
-    split_combined_df = pd.concat(split_dfs, ignore_index=True) if split_dfs else pd.DataFrame()
-
-
-    return main_combined_df, split_combined_df
+    if search_in_small_objects:
+        split = f'Splits_{dataset_subset}'
+    else:
+        split = f'Main_{dataset_subset}'
+    dataset_name = f"quasara-io/{dataset_name}"
+    dataset = load_dataset(dataset_name, split=split)
+    total_rows = dataset.num_rows 
+    # Convert to DataFrame and sample if limit is provided
+    if limit is not None:
+        df = dataset.to_pandas().sample(n=limit, random_state=42)
+    else:
+        df = dataset.to_pandas()
+    
+    return df,total_rows
     
 def get_image_vectors(df):
     # Get the image vectors from the dataframe
@@ -96,7 +71,7 @@ def get_image_vectors(df):
     return torch.tensor(image_vectors, dtype=torch.float32)
 
 
-def search(query, df, limit, offset, scoring_func, search_in_images):
+def search(query, df, limit, search_in_images = True):
     if search_in_images:
         # Encode the image query
         query_vector = encode_query(query)
@@ -266,7 +241,8 @@ def main():
     dataset = load_dataset(f"quasara-io/{dataset_name}")
     print('loaded dataset')
     dataset_keys = dataset.keys()
-    main_df, split_df = parallel_load_and_combine(dataset_keys, dataset)
+    random_sample_size = 1000
+    main_df, split_df = parallel_load_and_combine(dataset_keys, dataset, n_rows=random_sample_size)
     #Now we get the coordinates and the stuff
     print('processed datasets')
     if search_in_small_objects: