reverse search

app.py

@@ -29,6 +29,8 @@ LOG_SCALE_STATS = {
 
 def find_folders(base_folder, path):
     base_folder = get_datafolder(base_folder)
+    if not base_folder.exists(path):
+        return []
     return sorted(
         [
             folder["name"]
@@ -49,9 +51,9 @@ def find_stats_folders(base_folder: str):
     return sorted(list(set(stats_folders)))
 
 
-def fetch_runs(base_folder: str):
-    runs = sorted(find_stats_folders(base_folder))
-    return runs, gr.update(choices=runs, value=None)
+def fetch_datasets(base_folder: str):
+    datasets = sorted(find_stats_folders(base_folder))
+    return datasets, gr.update(choices=datasets, value=None), fetch_groups(base_folder, datasets, None, "union")
 
 
 def export_data(exported_data):
@@ -64,14 +66,19 @@ def export_data(exported_data):
     return gr.update(visible=True, value=temp_path)
 
 
-def fetch_groups(base_folder, datasets, old_groups):
-    GROUPS = [
-        [Path(x).name for x in find_folders(base_folder, run)] for run in datasets
-    ]
+def fetch_groups(base_folder, datasets, old_groups, type="intersection"):
+    if not datasets:
+        return gr.update(choices=[], value=None)
+
+    with ThreadPoolExecutor() as executor:
+        GROUPS = list(executor.map(lambda run: [Path(x).name for x in find_folders(base_folder, run)], datasets))
     if len(GROUPS) == 0:
         return gr.update(choices=[], value=None)
 
-    new_choices = set.intersection(*(set(g) for g in GROUPS))
+    if type == "intersection":
+        new_choices = set.intersection(*(set(g) for g in GROUPS)) 
+    elif type == "union":
+        new_choices = set.union(*(set(g) for g in GROUPS))
     value = None
     if old_groups:
         value = list(set.intersection(new_choices, {old_groups}))
@@ -81,15 +88,17 @@ def fetch_groups(base_folder, datasets, old_groups):
     return gr.update(choices=sorted(list(new_choices)), value=value)
 
 
-def fetch_stats(base_folder, datasets, group, old_stats):
-    STATS = [
-        [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")]
-        for run in datasets
-    ]
+def fetch_stats(base_folder, datasets, group, old_stats, type="intersection"):
+    print("Fetching stats")
+    with ThreadPoolExecutor() as executor:
+        STATS = list(executor.map(lambda run: [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")], datasets))
     if len(STATS) == 0:
         return gr.update(choices=[], value=None)
 
-    new_possibles_choices = set.intersection(*(set(s) for s in STATS))
+    if type == "intersection":
+        new_possibles_choices = set.intersection(*(set(s) for s in STATS))
+    elif type == "union":
+        new_possibles_choices = set.union(*(set(s) for s in STATS))
     value = None
     if old_stats:
         value = list(set.intersection(new_possibles_choices, {old_stats}))
@@ -98,6 +107,23 @@ def fetch_stats(base_folder, datasets, group, old_stats):
     return gr.update(choices=sorted(list(new_possibles_choices)), value=value)
 
 
+def reverse_search(base_folder, possible_datasets, grouping, stat_name):
+    with ThreadPoolExecutor() as executor:
+        found_datasets = list(executor.map(lambda dataset: dataset if stat_exists(base_folder, dataset, stat_name, grouping) else None, possible_datasets))
+    found_datasets = [dataset for dataset in found_datasets if dataset is not None]
+    return "\n".join(found_datasets)
+
+
+def reverse_search_add(datasets, reverse_search_results):
+    datasets = datasets or []
+    return sorted(list(set(datasets + reverse_search_results.strip().split("\n"))))
+
+
+
+def stat_exists(base_folder, path, stat_name, group_by):
+    base_folder = get_datafolder(base_folder)
+    return base_folder.exists(f"{path}/{group_by}/{stat_name}/stats-merged.json")
+
 def load_stats(base_folder, path, stat_name, group_by):
     base_folder = get_datafolder(base_folder)
     with base_folder.open(
@@ -265,12 +291,13 @@ def update_graph(
 with gr.Blocks() as demo:
     datasets = gr.State([])
     exported_data = gr.State([])
+    stats_headline = gr.Markdown(value="# Stats Exploration")
     with gr.Row():
         with gr.Column(scale=2):
             # Define the multiselect for crawls
             with gr.Row():
                 with gr.Column(scale=1):
-                    stats_folder = gr.Textbox(
+                    base_folder = gr.Textbox(
                         label="Stats Location",
                         value="s3://fineweb-stats/summary/",
                     )
@@ -298,10 +325,6 @@ Groupings:
     * k: the number of groups to show
     * Top/Bottom: the top/bottom k groups are shown
 - summary: simply shows the average value of given stat for selected crawls
-
-
-                
-                
                 """,
             )
         with gr.Column(scale=1):
@@ -348,11 +371,41 @@ Groupings:
     with gr.Row():
         # Define the graph output
         graph_output = gr.Plot(label="Graph")
+    
+    with gr.Row():
+        reverse_search_headline = gr.Markdown(value="# Reverse stats search")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Define the dropdown for grouping
+            reverse_grouping_dropdown = gr.Dropdown(
+                choices=[],
+                label="Grouping",
+                multiselect=False,
+            )
+            # Define the dropdown for stat_name
+            reverse_stat_name_dropdown = gr.Dropdown(
+                choices=[],
+                label="Stat name",
+                multiselect=False,
+            )
+        
+        with gr.Column(scale=1):
+            reverse_search_button = gr.Button("Search")
+            reverse_search_add_button = gr.Button("Add to selection")
+
+        with gr.Column(scale=2):
+            reverse_search_results = gr.Textbox(
+                label="Found datasets",
+                lines=10,
+            )
+        
+
 
     update_button.click(
         fn=update_graph,
         inputs=[
-            stats_folder,
+            base_folder,
             datasets_selected,
             stat_name_dropdown,
             grouping_dropdown,
@@ -369,22 +422,41 @@ Groupings:
         outputs=export_data_json,
     )
 
-    datasets_selected.select(
+    datasets_selected.change(
         fn=fetch_groups,
-        inputs=[stats_folder, datasets_selected, grouping_dropdown],
+        inputs=[base_folder, datasets_selected, grouping_dropdown],
         outputs=grouping_dropdown,
     )
 
     grouping_dropdown.select(
         fn=fetch_stats,
-        inputs=[stats_folder, datasets_selected, grouping_dropdown, stat_name_dropdown],
+        inputs=[base_folder, datasets_selected, grouping_dropdown, stat_name_dropdown],
         outputs=stat_name_dropdown,
     )
 
+    reverse_grouping_dropdown.select(
+        fn=partial(fetch_stats, type="union"),
+        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_stat_name_dropdown],
+        outputs=reverse_stat_name_dropdown,
+    )
+
+    reverse_search_button.click(
+        fn=reverse_search,
+        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_stat_name_dropdown],
+        outputs=reverse_search_results,
+    )
+
+    reverse_search_add_button.click(
+        fn=reverse_search_add,
+        inputs=[datasets_selected, reverse_search_results],
+        outputs=datasets_selected,
+    )
+
+
     datasets_refetch.click(
-        fn=fetch_runs,
-        inputs=[stats_folder],
-        outputs=[datasets, datasets_selected],
+        fn=fetch_datasets,
+        inputs=[base_folder],
+        outputs=[datasets, datasets_selected, reverse_grouping_dropdown],
     )
 
     def update_datasets_with_regex(regex, selected_runs, all_runs):