regex + folder definition + export

app.py

@@ -1,15 +1,22 @@
+from concurrent.futures import ThreadPoolExecutor
+import enum
 from functools import partial
 import json
 from pathlib import Path
+import re
+import tempfile
+from typing import Literal
 import gradio as gr
 
 from collections import defaultdict
-import fsspec.config
-import math
 from datatrove.io import DataFolder, get_datafolder
+import plotly.graph_objects as go
 from datatrove.utils.stats import MetricStatsDict
+import plotly.express as px
+
+import gradio as gr
+PARTITION_OPTIONS = Literal[ "Top", "Bottom", "Most frequent (n_docs)"]
 
-BASE_DATA_FOLDER = get_datafolder("s3://fineweb-stats/summary/")
 LOG_SCALE_STATS = {
     "length",
     "n_lines",
@@ -18,35 +25,10 @@ LOG_SCALE_STATS = {
     "avg_words_per_line",
     "pages_with_lorem_ipsum",
 }
-colors = list(
-    [
-        "rgba(31, 119, 180, 0.5)",
-        "rgba(255, 127, 14, 0.5)",
-        "rgba(44, 160, 44, 0.5)",
-        "rgba(214, 39, 40, 0.5)",
-        "rgba(148, 103, 189, 0.5)",
-        "rgba(227, 119, 194, 0.5)",
-        "rgba(127, 127, 127, 0.5)",
-        "rgba(188, 189, 34, 0.5)",
-        "rgba(23, 190, 207, 0.5)",
-        "rgba(255, 193, 7, 0.5)",
-        "rgba(40, 167, 69, 0.5)",
-        "rgba(23, 162, 184, 0.5)",
-        "rgba(108, 117, 125, 0.5)",
-        "rgba(0, 123, 255, 0.5)",
-        "rgba(220, 53, 69, 0.5)",
-        "rgba(255, 159, 67, 0.5)",
-        "rgba(255, 87, 34, 0.5)",
-        "rgba(41, 182, 246, 0.5)",
-        "rgba(142, 36, 170, 0.5)",
-        "rgba(0, 188, 212, 0.5)",
-        "rgba(255, 235, 59, 0.5)",
-        "rgba(156, 39, 176, 0.5)",
-    ]
-)
 
 
 def find_folders(base_folder, path):
+    base_folder = get_datafolder(base_folder)
     return sorted(
         [
             folder["name"]
@@ -56,9 +38,10 @@ def find_folders(base_folder, path):
     )
 
 
-def find_stats_folders(base_folder: DataFolder):
+def find_stats_folders(base_folder: str):
+    base_data_folder = get_datafolder(base_folder)
     # First find all stats-merged.json using globing for stats-merged.json
-    stats_merged = base_folder.glob("**/stats-merged.json")
+    stats_merged = base_data_folder.glob("**/stats-merged.json")
 
     # Then for each of stats.merged take the all but last two parts of the path (grouping/stat_name)
     stats_folders = [str(Path(x).parent.parent.parent) for x in stats_merged]
@@ -66,14 +49,25 @@ def find_stats_folders(base_folder: DataFolder):
     return sorted(list(set(stats_folders)))
 
 
-RUNS = sorted(find_stats_folders(BASE_DATA_FOLDER))
+def fetch_runs(base_folder: str):
+    runs = sorted(find_stats_folders(base_folder))
+    return runs, gr.update(choices=runs, value=None)
 
 
-def fetch_groups(runs, old_groups):
+def export_data(exported_data):
+    if not exported_data:
+        return None
+    # Assuming exported_data is a dictionary where the key is the dataset name and the value is the data to be exported
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as temp:
+        json.dump(exported_data, temp)
+        temp_path = temp.name
+    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_DATA_FOLDER, run)] for run in runs
+        [Path(x).name for x in find_folders(base_folder, run)] for run in datasets
     ]
-    # DO the intersection
     if len(GROUPS) == 0:
         return gr.update(choices=[], value=None)
 
@@ -84,13 +78,13 @@ def fetch_groups(runs, old_groups):
         value = value[0] if value else None
 
     # now take the intersection of all grups
-    return gr.update(choices=list(new_choices), value=value)
+    return gr.update(choices=sorted(list(new_choices)), value=value)
 
 
-def fetch_stats(runs, group, old_stats):
+def fetch_stats(base_folder, datasets, group, old_stats):
     STATS = [
-        [Path(x).name for x in find_folders(BASE_DATA_FOLDER, f"{run}/{group}")]
-        for run in runs
+        [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")]
+        for run in datasets
     ]
     if len(STATS) == 0:
         return gr.update(choices=[], value=None)
@@ -101,21 +95,21 @@ def fetch_stats(runs, group, old_stats):
         value = list(set.intersection(new_possibles_choices, {old_stats}))
         value = value[0] if value else None
 
-    return gr.update(choices=list(new_possibles_choices), value=value)
+    return gr.update(choices=sorted(list(new_possibles_choices)), value=value)
 
 
-def load_stats(path, stat_name, group_by):
-    with BASE_DATA_FOLDER.open(
+def load_stats(base_folder, path, stat_name, group_by):
+    base_folder = get_datafolder(base_folder)
+    with base_folder.open(
         f"{path}/{group_by}/{stat_name}/stats-merged.json",
-        filecache={"cache_storage": "/tmp/files"},
     ) as f:
         json_stat = json.load(f)
         # No idea why this is necessary, but it is, otheriwse the Metric StatsDict is malforme
         return MetricStatsDict() + MetricStatsDict(init=json_stat)
 
 
-def prepare_non_grouped_data(path, stat_name, grouping, normalization):
-    stats = load_stats(path, stat_name, grouping)
+def prepare_non_grouped_data(dataset_path, base_folder, grouping, stat_name, normalization):
+    stats = load_stats(base_folder, dataset_path, stat_name, grouping)
     stats_rounded = defaultdict(lambda: 0)
     for key, value in stats.items():
         stats_rounded[float(key)] += value.total
@@ -125,10 +119,10 @@ def prepare_non_grouped_data(path, stat_name, grouping, normalization):
     return stats_rounded
 
 
-def prepare_grouped_data(path, stat_name, grouping, top_k, direction):
+def prepare_grouped_data(dataset_path, base_folder, grouping, stat_name, top_k, direction: PARTITION_OPTIONS):
     import heapq
 
-    stats = load_stats(path, stat_name, grouping)
+    stats = load_stats(base_folder, dataset_path, stat_name, grouping)
 
     means = {key: value.mean for key, value in stats.items()}
 
@@ -136,13 +130,7 @@ def prepare_grouped_data(path, stat_name, grouping, top_k, direction):
     if direction == "Top":
         keys = heapq.nlargest(top_k, means, key=means.get)
     elif direction == "Most frequent (n_docs)":
-        n_docs = load_stats(path, "n_docs", grouping)
-        totals = {key: value.total for key, value in n_docs.items()}
-        keys = heapq.nlargest(top_k, totals, key=totals.get)
-
-    elif direction == "Most frequent (length)":
-        n_docs = load_stats(path, "length", grouping)
-        totals = {key: value.total for key, value in n_docs.items()}
+        totals = {key: value.n for key, value in stats.items()}
         keys = heapq.nlargest(top_k, totals, key=totals.get)
     else:
         keys = heapq.nsmallest(top_k, means, key=means.get)
@@ -150,17 +138,29 @@ def prepare_grouped_data(path, stat_name, grouping, top_k, direction):
     return [(key, means[key]) for key in keys]
 
 
-import math
-import plotly.graph_objects as go
-from plotly.offline import plot
+def set_alpha(color, alpha):
+    """
+    Takes a hex color and returns
+    rgba(r, g, b, a)
+    """
+    if color.startswith('#'):
+        r, g, b = int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16)
+    else:
+        r, g, b = 0, 0, 0  # Fallback to black if the color format is not recognized
+    return f"rgba({r}, {g}, {b}, {alpha})"
+
+
 
 
 def plot_scatter(
-    histograms: dict[str, dict[float, float]], stat_name: str, normalization: bool
+    histograms: dict[str, dict[float, float]],
+    stat_name: str,
+    normalization: bool,
+    progress: gr.Progress,
 ):
     fig = go.Figure()
 
-    for i, (name, histogram) in enumerate(histograms.items()):
+    for i, (name, histogram) in enumerate(progress.tqdm(histograms.items(), total=len(histograms), desc="Plotting...")):
         if all(isinstance(k, str) for k in histogram.keys()):
             x = [k for k, v in sorted(histogram.items(), key=lambda item: item[1])]
         else:
@@ -174,7 +174,7 @@ def plot_scatter(
                 y=y,
                 mode="lines",
                 name=name,
-                line=dict(color=colors[i % len(colors)]),
+                marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5)),
             )
         )
 
@@ -194,14 +194,18 @@ def plot_scatter(
     return fig
 
 
-def plot_bars(histograms: dict[str, list[tuple[str, float]]], stat_name: str):
+def plot_bars(
+    histograms: dict[str, list[tuple[str, float]]],
+    stat_name: str,
+    progress: gr.Progress,
+):
     fig = go.Figure()
 
-    for i, (name, histogram) in enumerate(histograms.items()):
+    for i, (name, histogram) in enumerate(progress.tqdm(histograms.items(), total=len(histograms), desc="Plotting...")):
         x = [k for k, v in histogram]
         y = [v for k, v in histogram]
 
-        fig.add_trace(go.Bar(x=x, y=y, name=name, marker_color=colors[i % len(colors)]))
+        fig.add_trace(go.Bar(x=x, y=y, name=name, marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5))))
 
     fig.update_layout(
         title=f"Bar Plots for {stat_name}",
@@ -217,9 +221,16 @@ def plot_bars(histograms: dict[str, list[tuple[str, float]]], stat_name: str):
 
 
 def update_graph(
-    multiselect_crawls, stat_name, grouping, normalization, top_k, direction
+    base_folder,
+    datasets,
+    stat_name,
+    grouping,
+    normalization,
+    top_k,
+    direction,
+    progress=gr.Progress(),
 ):
-    if len(multiselect_crawls) <= 0 or not stat_name or not grouping:
+    if len(datasets) <= 0 or not stat_name or not grouping:
         return None
     # Placeholder for logic to rerender the graph based on the inputs
     prepare_fc = (
@@ -233,25 +244,48 @@ def update_graph(
         else plot_bars
     )
 
-    print("Loading stats")
-    histograms = {
-        path: prepare_fc(path, stat_name, grouping) for path in multiselect_crawls
-    }
+    with ThreadPoolExecutor() as pool:
+        data = list(
+            progress.tqdm(
+                pool.map(
+                    partial(prepare_fc, base_folder=base_folder, stat_name=stat_name, grouping=grouping),
+                    datasets,
+                ),
+                total=len(datasets),
+                desc="Loading data...",
+            )
+        )
+
+    histograms = {path: result for path, result in zip(datasets, data)}
 
-    print("Plotting")
-    return graph_fc(histograms, stat_name)
+    return graph_fc(histograms=histograms, stat_name=stat_name, progress=progress), histograms, gr.update(visible=True)
 
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
+    datasets = gr.State([])
+    exported_data = gr.State([])
     with gr.Row():
         with gr.Column(scale=2):
             # Define the multiselect for crawls
-            multiselect_crawls = gr.Dropdown(
-                choices=RUNS,
-                label="Multiselect for crawls",
-                multiselect=True,
-            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    stats_folder = gr.Textbox(
+                        label="Stats Location",
+                        value="s3://fineweb-stats/summary/",
+                    )
+                    datasets_refetch = gr.Button("Fetch Datasets")
+
+                with gr.Column(scale=1):
+                    regex_select = gr.Text(label="Regex select datasets", value=".*")
+                    regex_button = gr.Button("Filter")
+            with gr.Row():
+                datasets_selected = gr.Dropdown(
+                    choices=[],
+                    label="Datasets",
+                    multiselect=True,
+                )
+
             # add a readme description
             readme_description = gr.Markdown(
                 label="Readme",
@@ -283,6 +317,7 @@ Groupings:
                 label="Stat name",
                 multiselect=False,
             )
+
             with gr.Row(visible=False) as histogram_choices:
                 normalization_checkbox = gr.Checkbox(
                     label="Normalize",
@@ -301,11 +336,15 @@ Groupings:
                         "Top",
                         "Bottom",
                         "Most frequent (n_docs)",
-                        "Most frequent (length)",
                     ],
+                    value="Top",
                 )
 
             update_button = gr.Button("Update Graph", variant="primary")
+    with gr.Row():
+        export_data_button = gr.Button("Export data", visible=False)
+        export_data_json = gr.File(visible=False)
+
     with gr.Row():
         # Define the graph output
         graph_output = gr.Plot(label="Graph")
@@ -313,28 +352,54 @@ Groupings:
     update_button.click(
         fn=update_graph,
         inputs=[
-            multiselect_crawls,
+            stats_folder,
+            datasets_selected,
             stat_name_dropdown,
             grouping_dropdown,
             normalization_checkbox,
             top_select,
             direction_checkbox,
         ],
-        outputs=graph_output,
+        outputs=[graph_output, exported_data, export_data_button],
+    )
+
+    export_data_button.click(
+        fn=export_data,
+        inputs=[exported_data],
+        outputs=export_data_json,
     )
 
-    multiselect_crawls.select(
+    datasets_selected.select(
         fn=fetch_groups,
-        inputs=[multiselect_crawls, grouping_dropdown],
+        inputs=[stats_folder, datasets_selected, grouping_dropdown],
         outputs=grouping_dropdown,
     )
 
     grouping_dropdown.select(
         fn=fetch_stats,
-        inputs=[multiselect_crawls, grouping_dropdown, stat_name_dropdown],
+        inputs=[stats_folder, datasets_selected, grouping_dropdown, stat_name_dropdown],
         outputs=stat_name_dropdown,
     )
 
+    datasets_refetch.click(
+        fn=fetch_runs,
+        inputs=[stats_folder],
+        outputs=[datasets, datasets_selected],
+    )
+
+    def update_datasets_with_regex(regex, selected_runs, all_runs):
+        if not regex:
+            return
+        new_dsts = {run for run in all_runs if re.search(regex, run)}
+        dst_union = new_dsts.union(selected_runs)
+        return gr.update(value=list(dst_union))
+
+    regex_button.click(
+        fn=update_datasets_with_regex,
+        inputs=[regex_select, datasets_selected, datasets],
+        outputs=datasets_selected,
+    )
+
     def update_grouping_options(grouping):
         if grouping == "histogram":
             return {