Revert "tmp revert lateer"

This reverts commit e906b0bf37789596a54cc78a37fdfbd859aa4fc5.

app.py

@@ -1,577 +1,572 @@
-# from concurrent.futures import ThreadPoolExecutor
-# import enum
-# from functools import partial
-# import json
-# import os
-# from pathlib import Path
-# import re
-# import heapq
-# import tempfile
-# from typing import Literal
-# import gradio as gr
-
-# from collections import defaultdict
-# from datatrove.io import get_datafolder
-# import plotly.graph_objects as go
-# from datatrove.utils.stats import MetricStats, MetricStatsDict
-# import plotly.express as px
-# import tenacity
-
-# import gradio as gr
-# PARTITION_OPTIONS = Literal[ "Top", "Bottom", "Most frequent (n_docs)"]
-# METRICS_LOCATION_DEFAULT = os.getenv("METRICS_LOCATION_DEFAULT", "s3://fineweb-stats/summary/")
-
-
-# def find_folders(base_folder, path):
-#     base_folder = get_datafolder(base_folder)
-#     if not base_folder.exists(path):
-#         return []
-#     return sorted(
-#         [
-#             folder["name"]
-#             for folder in base_folder.ls(path, detail=True)
-#             if folder["type"] == "directory" and not folder["name"].rstrip("/") == path
-#         ]
-#     )
-
-
-# def find_metrics_folders(base_folder: str):
-#     base_data_folder = get_datafolder(base_folder)
-#     # First find all metric.json using globing for metric.json
-#     metrics_merged = base_data_folder.glob("**/metric.json")
-
-#     # Then for each of metrics.merged take the all but last two parts of the path (grouping/metric_name)
-#     metrics_folders = [str(Path(x).parent.parent.parent) for x in metrics_merged]
-#     # Finally get the unique paths
-#     return sorted(list(set(metrics_folders)))
-
-
-# def fetch_datasets(base_folder: str):
-#     datasets = sorted(find_metrics_folders(base_folder))
-#     return datasets, gr.update(choices=datasets, value=None), fetch_groups(base_folder, datasets, None, "union")
-
-
-# def export_data(exported_data: MetricStatsDict, metric_name: str):
-#     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
-#     temp_dir = tempfile.mkdtemp()
-#     temp_path = os.path.join(temp_dir, metric_name + ".json")
-#     with open(temp_path, "w") as temp_file:
-#         json.dump({
-#             name: dt.to_dict()
-#             for name, dt in exported_data.items()
-#         }, temp_file)
-#     return gr.update(visible=True, value=temp_path)
-
-
-# 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)
-
-#     if type == "intersection":
-#         new_choices = set.intersection(*(set(g) for g in GROUPS)) 
-#     else:
-#         new_choices = set.union(*(set(g) for g in GROUPS))
-#     value = None
-#     if old_groups:
-#         value = list(set.intersection(new_choices, {old_groups}))
-#         value = value[0] if value else None
-
-#     # now take the intersection of all grups
-#     return gr.update(choices=sorted(list(new_choices)), value=value)
-
-
-# def fetch_metrics(base_folder, datasets, group, old_metrics, type="intersection"):
-#     with ThreadPoolExecutor() as executor:
-#         metrics = list(executor.map(lambda run: [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")], datasets))
-#     if len(metrics) == 0:
-#         return gr.update(choices=[], value=None)
-
-#     if type == "intersection":
-#         new_possibles_choices = set.intersection(*(set(s) for s in metrics))
-#     else:
-#         new_possibles_choices = set.union(*(set(s) for s in metrics))
-#     value = None
-#     if old_metrics:
-#         value = list(set.intersection(new_possibles_choices, {old_metrics}))
-#         value = value[0] if value else None
-
-#     return gr.update(choices=sorted(list(new_possibles_choices)), value=value)
-
-
-# def reverse_search(base_folder, possible_datasets, grouping, metric_name):
-#     with ThreadPoolExecutor() as executor:
-#         found_datasets = list(executor.map(lambda dataset: dataset if metric_exists(base_folder, dataset, metric_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 metric_exists(base_folder, path, metric_name, group_by):
-#     base_folder = get_datafolder(base_folder)
-#     return base_folder.exists(f"{path}/{group_by}/{metric_name}/metric.json")
-
-# @tenacity.retry(stop=tenacity.stop_after_attempt(5))
-# def load_metrics(base_folder, path, metric_name, group_by):
-#     base_folder = get_datafolder(base_folder)
-#     with base_folder.open(
-#         f"{path}/{group_by}/{metric_name}/metric.json",
-#     ) as f:
-#         json_metric = json.load(f)
-#         # No idea why this is necessary, but it is, otheriwse the Metric StatsDict is malformed
-#         return MetricStatsDict.from_dict(json_metric)
-
-
-# def prepare_for_non_grouped_plotting(metric, normalization, rounding):
-#     metrics_rounded = defaultdict(lambda: 0)
-#     for key, value in metric.items():
-#         metrics_rounded[round(float(key), rounding)] += value.total
-#     if normalization:
-#         normalizer = sum(metrics_rounded.values())
-#         metrics_rounded = {k: v / normalizer for k, v in metrics_rounded.items()}
-#         # check that the sum of the values is 1
-#         summed = sum(metrics_rounded.values())
-#         assert abs(summed - 1) < 0.01, summed
-#     return metrics_rounded
-
-
-# def load_data(dataset_path, base_folder, grouping, metric_name):
-#     metrics = load_metrics(base_folder, dataset_path, metric_name, grouping)
-#     return metrics
-
-# def prepare_for_group_plotting(metric, top_k, direction: PARTITION_OPTIONS, regex: str | None, rounding: int):
-#     regex_compiled = re.compile(regex) if regex else None
-#     metric = {key: value for key, value in metric.items() if not regex or regex_compiled.match(key)}
-#     means = {key: round(float(value.mean), rounding) for key, value in metric.items()}
-#     # Use heap to get top_k keys
-#     if direction == "Top":
-#         keys = heapq.nlargest(top_k, means, key=means.get)
-#     elif direction == "Most frequent (n_docs)":
-#         totals = {key: int(value.n) for key, value in metric.items()}
-#         keys = heapq.nlargest(top_k, totals, key=totals.get)
-#     else:
-#         keys = heapq.nsmallest(top_k, means, key=means.get)
+from concurrent.futures import ThreadPoolExecutor
+import enum
+from functools import partial
+import json
+import os
+from pathlib import Path
+import re
+import heapq
+import tempfile
+from typing import Literal
+import gradio as gr
+
+from collections import defaultdict
+from datatrove.io import get_datafolder
+import plotly.graph_objects as go
+from datatrove.utils.stats import MetricStats, MetricStatsDict
+import plotly.express as px
+import tenacity
+
+import gradio as gr
+PARTITION_OPTIONS = Literal[ "Top", "Bottom", "Most frequent (n_docs)"]
+METRICS_LOCATION_DEFAULT = os.getenv("METRICS_LOCATION_DEFAULT", "s3://fineweb-stats/summary/")
+
+
+def find_folders(base_folder, path):
+    base_folder = get_datafolder(base_folder)
+    if not base_folder.exists(path):
+        return []
+    return sorted(
+        [
+            folder["name"]
+            for folder in base_folder.ls(path, detail=True)
+            if folder["type"] == "directory" and not folder["name"].rstrip("/") == path
+        ]
+    )
+
+
+def find_metrics_folders(base_folder: str):
+    base_data_folder = get_datafolder(base_folder)
+    # First find all metric.json using globing for metric.json
+    metrics_merged = base_data_folder.glob("**/metric.json")
+
+    # Then for each of metrics.merged take the all but last two parts of the path (grouping/metric_name)
+    metrics_folders = [str(Path(x).parent.parent.parent) for x in metrics_merged]
+    # Finally get the unique paths
+    return sorted(list(set(metrics_folders)))
+
+
+def fetch_datasets(base_folder: str):
+    datasets = sorted(find_metrics_folders(base_folder))
+    return datasets, gr.update(choices=datasets, value=None), fetch_groups(base_folder, datasets, None, "union")
+
+
+def export_data(exported_data: MetricStatsDict, metric_name: str):
+    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, prefix=metric_name, suffix=".json") as temp:
+        json.dump({
+            name: dt.to_dict()
+            for name, dt in exported_data.items()
+        }, temp)
+        temp_path = temp.name
+    return gr.update(visible=True, value=temp_path)
+
+
+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)
+
+    if type == "intersection":
+        new_choices = set.intersection(*(set(g) for g in GROUPS)) 
+    else:
+        new_choices = set.union(*(set(g) for g in GROUPS))
+    value = None
+    if old_groups:
+        value = list(set.intersection(new_choices, {old_groups}))
+        value = value[0] if value else None
+
+    # now take the intersection of all grups
+    return gr.update(choices=sorted(list(new_choices)), value=value)
+
+
+def fetch_metrics(base_folder, datasets, group, old_metrics, type="intersection"):
+    with ThreadPoolExecutor() as executor:
+        metrics = list(executor.map(lambda run: [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")], datasets))
+    if len(metrics) == 0:
+        return gr.update(choices=[], value=None)
+
+    if type == "intersection":
+        new_possibles_choices = set.intersection(*(set(s) for s in metrics))
+    else:
+        new_possibles_choices = set.union(*(set(s) for s in metrics))
+    value = None
+    if old_metrics:
+        value = list(set.intersection(new_possibles_choices, {old_metrics}))
+        value = value[0] if value else None
+
+    return gr.update(choices=sorted(list(new_possibles_choices)), value=value)
+
+
+def reverse_search(base_folder, possible_datasets, grouping, metric_name):
+    with ThreadPoolExecutor() as executor:
+        found_datasets = list(executor.map(lambda dataset: dataset if metric_exists(base_folder, dataset, metric_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 metric_exists(base_folder, path, metric_name, group_by):
+    base_folder = get_datafolder(base_folder)
+    return base_folder.exists(f"{path}/{group_by}/{metric_name}/metric.json")
+
+@tenacity.retry(stop=tenacity.stop_after_attempt(5))
+def load_metrics(base_folder, path, metric_name, group_by):
+    base_folder = get_datafolder(base_folder)
+    with base_folder.open(
+        f"{path}/{group_by}/{metric_name}/metric.json",
+    ) as f:
+        json_metric = json.load(f)
+        # No idea why this is necessary, but it is, otheriwse the Metric StatsDict is malformed
+        return MetricStatsDict.from_dict(json_metric)
+
+
+def prepare_for_non_grouped_plotting(metric, normalization, rounding):
+    metrics_rounded = defaultdict(lambda: 0)
+    for key, value in metric.items():
+        metrics_rounded[round(float(key), rounding)] += value.total
+    if normalization:
+        normalizer = sum(metrics_rounded.values())
+        metrics_rounded = {k: v / normalizer for k, v in metrics_rounded.items()}
+        # check that the sum of the values is 1
+        summed = sum(metrics_rounded.values())
+        assert abs(summed - 1) < 0.01, summed
+    return metrics_rounded
+
+
+def load_data(dataset_path, base_folder, grouping, metric_name):
+    metrics = load_metrics(base_folder, dataset_path, metric_name, grouping)
+    return metrics
+
+def prepare_for_group_plotting(metric, top_k, direction: PARTITION_OPTIONS, regex: str | None, rounding: int):
+    regex_compiled = re.compile(regex) if regex else None
+    metric = {key: value for key, value in metric.items() if not regex or regex_compiled.match(key)}
+    means = {key: round(float(value.mean), rounding) for key, value in metric.items()}
+    # Use heap to get top_k keys
+    if direction == "Top":
+        keys = heapq.nlargest(top_k, means, key=means.get)
+    elif direction == "Most frequent (n_docs)":
+        totals = {key: int(value.n) for key, value in metric.items()}
+        keys = heapq.nlargest(top_k, totals, key=totals.get)
+    else:
+        keys = heapq.nsmallest(top_k, means, key=means.get)
 
     
-#     means = [means[key] for key in keys]
-#     stds = [metric[key].standard_deviation for key in keys]
-#     return keys, means, stds
-
-
-# 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(
-#     data: dict[str, dict[float, float]],
-#     metric_name: str,
-#     log_scale_x: bool,
-#     log_scale_y: bool,
-#     normalization: bool,
-#     rounding: int,
-#     progress: gr.Progress,
-# ):
-#     fig = go.Figure()
-
-#     # First sort the histograms, by their name
-#     data = {name: histogram for name, histogram in sorted(data.items())}
-#     for i, (name, histogram) in enumerate(progress.tqdm(data.items(), total=len(data), desc="Plotting...")):
-#         histogram_prepared = prepare_for_non_grouped_plotting(histogram, normalization, rounding)
-#         x = sorted(histogram_prepared.keys())
-#         y = [histogram_prepared[k] for k in x]
-
-#         fig.add_trace(
-#             go.Scatter(
-#                 x=x,
-#                 y=y,
-#                 mode="lines",
-#                 name=name,
-#                 marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5)),
-#             )
-#         )
-
-#     yaxis_title = "Frequency" if normalization else "Total"
-
-#     fig.update_layout(
-#         title=f"Line Plots for {metric_name}",
-#         xaxis_title=metric_name,
-#         yaxis_title=yaxis_title,
-#         xaxis_type="log" if log_scale_x and len(x) > 1 else None,
-#         yaxis_type="log" if log_scale_y and len(y) > 1 else None,
-#         width=1200,
-#         height=600,
-#         showlegend=True,
-#     )
-
-#     return fig
-
-
-# def plot_bars(
-#     data: dict[str, list[dict[str, float]]],
-#     metric_name: str,
-#     top_k: int,
-#     direction: PARTITION_OPTIONS,
-#     regex: str | None,
-#     rounding: int,
-#     log_scale_x: bool,
-#     log_scale_y: bool,
-#     progress: gr.Progress,
-# ):
-#     fig = go.Figure()
-#     x = []
-#     y = []
-
-#     for i, (name, histogram) in enumerate(progress.tqdm(data.items(), total=len(data), desc="Plotting...")):
-#         x, y, stds = prepare_for_group_plotting(histogram, top_k, direction, regex, rounding)
-
-#         fig.add_trace(go.Bar(
-#             x=x, 
-#             y=y, 
-#             name=f"{name} Mean", 
-#             marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5)),
-#             # error_y=dict(type='data', array=stds, visible=True)
-#         ))
-
-#     fig.update_layout(
-#         title=f"Bar Plots for {metric_name}",
-#         xaxis_title=metric_name,
-#         yaxis_title="Avg. value",
-#         xaxis_type="log" if log_scale_x and len(x) > 1 else None,
-#         yaxis_type="log" if log_scale_y and len(y) > 1 else None,
-#         autosize=True,
-#         width=1200,
-#         height=600,
-#         showlegend=True,
-#     )
-
-#     return fig
-
-
-# def update_graph(
-#     base_folder,
-#     datasets,
-#     metric_name,
-#     grouping,
-#     log_scale_x,
-#     log_scale_y,
-#     rounding,
-#     normalization,
-#     top_k,
-#     direction,
-#     regex,
-#     progress=gr.Progress(),
-# ):
-#     if len(datasets) <= 0 or not metric_name or not grouping:
-#         return None
-#     # Placeholder for logic to rerender the graph based on the inputs
-
-#     with ThreadPoolExecutor() as pool:
-#         data = list(
-#             progress.tqdm(
-#                 pool.map(
-#                     partial(load_data, base_folder=base_folder, metric_name=metric_name, grouping=grouping),
-#                     datasets,
-#                 ),
-#                 total=len(datasets),
-#                 desc="Loading data...",
-#             )
-#         )
-
-#     data = {path: result for path, result in zip(datasets, data)}
-#     return plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y, progress), data, export_data(data, metric_name)
-
-# def plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y, progress=gr.Progress()):
-#     if rounding is None or top_k is None:
-#         return None
-#     graph_fc = (
-#         partial(plot_scatter, normalization=normalization, rounding=rounding)
-#         if grouping == "histogram"
-#         else partial(plot_bars, top_k=top_k, direction=direction, regex=regex, rounding=rounding)
-#     )
-#     return graph_fc(data=data, metric_name=metric_name, progress=progress, log_scale_x=log_scale_x, log_scale_y=log_scale_y)
-
-
-
-# # Create the Gradio interface
-# with gr.Blocks() as demo:
-#     datasets = gr.State([])
-#     exported_data = gr.State([])
-#     metrics_headline = gr.Markdown(value="# Metrics Exploration")
-#     with gr.Row():
-#         with gr.Column(scale=2):
-#             with gr.Row():
-#                 with gr.Column(scale=1):
-#                     base_folder = gr.Textbox(
-#                         label="Metrics Location",
-#                         value=METRICS_LOCATION_DEFAULT,
-#                     )
-#                     datasets_refetch = gr.Button("Fetch Datasets")
-
-#                 with gr.Column(scale=1):
-#                     regex_select = gr.Text(label="Regex filter", value=".*")
-#                     regex_button = gr.Button("Search")
-#             with gr.Row():
-#                 datasets_selected = gr.Dropdown(
-#                     choices=[],
-#                     label="Datasets",
-#                     multiselect=True,
-#                 )
-
-#             # add a readme description
-#             readme_description = gr.Markdown(
-#                 label="Readme",
-#                 value="""
-# ## How to use:
-# 1) Specify Metrics location (Stats block `output_folder` without the last path segment) and click "Fetch Datasets"
-# 2) Select datasets you are interested in using the dropdown or regex filter
-# 3) Specify Grouping (global average/value/fqdn/suffix) and Metric name
-# 4) Click "Update Graph" 
-
-
-# ## Groupings:
-# - **histogram**: Creates a line plot of values with their frequencies. If normalization is on, the frequencies sum to 1.
-#     * normalize:
-# - **(fqdn/suffix)**: Creates a bar plot of the avg. values of the metric for full qualifed domain name/suffix of domain.
-#     * k: the number of groups to show
-#     * Top/Bottom/Most frequent (n_docs): Groups with the top/bottom k values/most prevalant docs are shown
-# - **none**: Shows the average value of given metric
-
-# ## Reverse search:
-# To search for datasets containing a grouping and certain metric, use the Reverse search section.
-# Specify the search parameters and click "Search". This will show you found datasets in the "Found datasets" textbox. You can modify the selection after search by removing unwanted lines and clicking "Add to selection".
-
-# ## Note:
-# The data might not be 100% representative, due to the sampling and optimistic merging of the metrics (fqdn/suffix).
-# """,
-#             )
-#         with gr.Column(scale=1):
-#             # Define the dropdown for grouping
-#             grouping_dropdown = gr.Dropdown(
-#                 choices=[],
-#                 label="Grouping",
-#                 multiselect=False,
-#             )
-#             # Define the dropdown for metric_name
-#             metric_name_dropdown = gr.Dropdown(
-#                 choices=[],
-#                 label="Metric name",
-#                 multiselect=False,
-#             )
-
-
-#             update_button = gr.Button("Update Graph", variant="primary")
-
-#     with gr.Row():
-#         with gr.Column(scale=1):
-#             log_scale_x_checkbox = gr.Checkbox(
-#                 label="Log scale x",
-#                 value=False,
-#             )
-#             log_scale_y_checkbox = gr.Checkbox(
-#                 label="Log scale y",
-#                 value=False,
-#             )
-#             rounding = gr.Number(
-#                 label="Rounding",
-#                 value=2,
-#             )
-#             normalization_checkbox = gr.Checkbox(
-#                 label="Normalize",
-#                 value=True,  # Default value
-#                 visible=False
-#             )
-#             with gr.Row():
-#                 # export_data_button = gr.Button("Export data", visible=True, link=export_data_json)
-#                 export_data_json = gr.File(visible=False)
-#         with gr.Column(scale=4):
-#             with gr.Row(visible=False) as group_choices:
-#                 with gr.Column(scale=2):
-#                     group_regex = gr.Text(
-#                         label="Group Regex",
-#                         value=None,
-#                     )
-#                     with gr.Row():
-#                         top_select = gr.Number(
-#                             label="N Groups",
-#                             value=100,
-#                             interactive=True,
-#                         )
+    means = [means[key] for key in keys]
+    stds = [metric[key].standard_deviation for key in keys]
+    return keys, means, stds
+
+
+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(
+    data: dict[str, dict[float, float]],
+    metric_name: str,
+    log_scale_x: bool,
+    log_scale_y: bool,
+    normalization: bool,
+    rounding: int,
+    progress: gr.Progress,
+):
+    fig = go.Figure()
+
+    # First sort the histograms, by their name
+    data = {name: histogram for name, histogram in sorted(data.items())}
+    for i, (name, histogram) in enumerate(progress.tqdm(data.items(), total=len(data), desc="Plotting...")):
+        histogram_prepared = prepare_for_non_grouped_plotting(histogram, normalization, rounding)
+        x = sorted(histogram_prepared.keys())
+        y = [histogram_prepared[k] for k in x]
+
+        fig.add_trace(
+            go.Scatter(
+                x=x,
+                y=y,
+                mode="lines",
+                name=name,
+                marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5)),
+            )
+        )
+
+    yaxis_title = "Frequency" if normalization else "Total"
+
+    fig.update_layout(
+        title=f"Line Plots for {metric_name}",
+        xaxis_title=metric_name,
+        yaxis_title=yaxis_title,
+        xaxis_type="log" if log_scale_x and len(x) > 1 else None,
+        yaxis_type="log" if log_scale_y and len(y) > 1 else None,
+        width=1200,
+        height=600,
+        showlegend=True,
+    )
+
+    return fig
+
+
+def plot_bars(
+    data: dict[str, list[dict[str, float]]],
+    metric_name: str,
+    top_k: int,
+    direction: PARTITION_OPTIONS,
+    regex: str | None,
+    rounding: int,
+    log_scale_x: bool,
+    log_scale_y: bool,
+    progress: gr.Progress,
+):
+    fig = go.Figure()
+    x = []
+    y = []
+
+    for i, (name, histogram) in enumerate(progress.tqdm(data.items(), total=len(data), desc="Plotting...")):
+        x, y, stds = prepare_for_group_plotting(histogram, top_k, direction, regex, rounding)
+
+        fig.add_trace(go.Bar(
+            x=x, 
+            y=y, 
+            name=f"{name} Mean", 
+            marker=dict(color=set_alpha(px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)], 0.5)),
+            error_y=dict(type='data', array=stds, visible=True)
+        ))
+
+    fig.update_layout(
+        title=f"Bar Plots for {metric_name}",
+        xaxis_title=metric_name,
+        yaxis_title="Avg. value",
+        xaxis_type="log" if log_scale_x and len(x) > 1 else None,
+        yaxis_type="log" if log_scale_y and len(y) > 1 else None,
+        autosize=True,
+        width=1200,
+        height=600,
+        showlegend=True,
+    )
+
+    return fig
+
+
+def update_graph(
+    base_folder,
+    datasets,
+    metric_name,
+    grouping,
+    log_scale_x,
+    log_scale_y,
+    rounding,
+    normalization,
+    top_k,
+    direction,
+    regex,
+    progress=gr.Progress(),
+):
+    if len(datasets) <= 0 or not metric_name or not grouping:
+        return None
+    # Placeholder for logic to rerender the graph based on the inputs
+
+    with ThreadPoolExecutor() as pool:
+        data = list(
+            progress.tqdm(
+                pool.map(
+                    partial(load_data, base_folder=base_folder, metric_name=metric_name, grouping=grouping),
+                    datasets,
+                ),
+                total=len(datasets),
+                desc="Loading data...",
+            )
+        )
+
+    data = {path: result for path, result in zip(datasets, data)}
+    return plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y, progress), data, export_data(data, metric_name)
+
+def plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y, progress=gr.Progress()):
+    if rounding is None or top_k is None:
+        return None
+    graph_fc = (
+        partial(plot_scatter, normalization=normalization, rounding=rounding)
+        if grouping == "histogram"
+        else partial(plot_bars, top_k=top_k, direction=direction, regex=regex, rounding=rounding)
+    )
+    return graph_fc(data=data, metric_name=metric_name, progress=progress, log_scale_x=log_scale_x, log_scale_y=log_scale_y)
+
+
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    datasets = gr.State([])
+    exported_data = gr.State([])
+    metrics_headline = gr.Markdown(value="# Metrics Exploration")
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    base_folder = gr.Textbox(
+                        label="Metrics Location",
+                        value=METRICS_LOCATION_DEFAULT,
+                    )
+                    datasets_refetch = gr.Button("Fetch Datasets")
+
+                with gr.Column(scale=1):
+                    regex_select = gr.Text(label="Regex filter", value=".*")
+                    regex_button = gr.Button("Search")
+            with gr.Row():
+                datasets_selected = gr.Dropdown(
+                    choices=[],
+                    label="Datasets",
+                    multiselect=True,
+                )
+
+            # add a readme description
+            readme_description = gr.Markdown(
+                label="Readme",
+                value="""
+## How to use:
+1) Specify Metrics location (Stats block `output_folder` without the last path segment) and click "Fetch Datasets"
+2) Select datasets you are interested in using the dropdown or regex filter
+3) Specify Grouping (global average/value/fqdn/suffix) and Metric name
+4) Click "Update Graph" 
+
+
+## Groupings:
+- **histogram**: Creates a line plot of values with their frequencies. If normalization is on, the frequencies sum to 1.
+    * normalize:
+- **(fqdn/suffix)**: Creates a bar plot of the avg. values of the metric for full qualifed domain name/suffix of domain.
+    * k: the number of groups to show
+    * Top/Bottom/Most frequent (n_docs): Groups with the top/bottom k values/most prevalant docs are shown
+- **none**: Shows the average value of given metric
+
+## Reverse search:
+To search for datasets containing a grouping and certain metric, use the Reverse search section.
+Specify the search parameters and click "Search". This will show you found datasets in the "Found datasets" textbox. You can modify the selection after search by removing unwanted lines and clicking "Add to selection".
+
+## Note:
+The data might not be 100% representative, due to the sampling and optimistic merging of the metrics (fqdn/suffix).
+""",
+            )
+        with gr.Column(scale=1):
+            # Define the dropdown for grouping
+            grouping_dropdown = gr.Dropdown(
+                choices=[],
+                label="Grouping",
+                multiselect=False,
+            )
+            # Define the dropdown for metric_name
+            metric_name_dropdown = gr.Dropdown(
+                choices=[],
+                label="Metric name",
+                multiselect=False,
+            )
+
+
+            update_button = gr.Button("Update Graph", variant="primary")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            log_scale_x_checkbox = gr.Checkbox(
+                label="Log scale x",
+                value=False,
+            )
+            log_scale_y_checkbox = gr.Checkbox(
+                label="Log scale y",
+                value=False,
+            )
+            rounding = gr.Number(
+                label="Rounding",
+                value=2,
+            )
+            normalization_checkbox = gr.Checkbox(
+                label="Normalize",
+                value=True,  # Default value
+                visible=False
+            )
+            with gr.Row():
+                # export_data_button = gr.Button("Export data", visible=True, link=export_data_json)
+                export_data_json = gr.File(visible=False)
+        with gr.Column(scale=4):
+            with gr.Row(visible=False) as group_choices:
+                with gr.Column(scale=2):
+                    group_regex = gr.Text(
+                        label="Group Regex",
+                        value=None,
+                    )
+                    with gr.Row():
+                        top_select = gr.Number(
+                            label="N Groups",
+                            value=100,
+                            interactive=True,
+                        )
                     
-#                         direction_checkbox = gr.Radio(
-#                             label="Partition",
-#                             choices=[
-#                                 "Top",
-#                                 "Bottom",
-#                                 "Most frequent (n_docs)",
-#                             ],
-#                             value="Most frequent (n_docs)",
-#                     )
-#         # Define the graph output
-#     with gr.Row():
-#         graph_output = gr.Plot(label="Graph")
+                        direction_checkbox = gr.Radio(
+                            label="Partition",
+                            choices=[
+                                "Top",
+                                "Bottom",
+                                "Most frequent (n_docs)",
+                            ],
+                            value="Most frequent (n_docs)",
+                    )
+        # Define the graph output
+    with gr.Row():
+        graph_output = gr.Plot(label="Graph")
     
-#     with gr.Row():
-#         reverse_search_headline = gr.Markdown(value="# Reverse metrics search")
+    with gr.Row():
+        reverse_search_headline = gr.Markdown(value="# Reverse metrics 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 metric_name
-#             reverse_metric_name_dropdown = gr.Dropdown(
-#                 choices=[],
-#                 label="Stat name",
-#                 multiselect=False,
-#             )
+    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 metric_name
+            reverse_metric_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,
-#                 placeholder="Found datasets containing the group/metric name. You can modify the selection after search by removing unwanted lines and clicking Add to selection"
-#             )
-
-
-#     update_button.click(
-#         fn=update_graph,
-#         inputs=[
-#             base_folder,
-#             datasets_selected,
-#             metric_name_dropdown,
-#             grouping_dropdown,
-#             log_scale_x_checkbox,
-#             log_scale_y_checkbox,
-#             rounding,
-#             normalization_checkbox,
-#             top_select,
-#             direction_checkbox,
-#             group_regex,
-#         ],
-#         outputs=[graph_output, exported_data, export_data_json],
-#     )
-
-#     for inp in [normalization_checkbox, rounding, group_regex, direction_checkbox, top_select, log_scale_x_checkbox, log_scale_y_checkbox]:
-#         inp.change(
-#             fn=plot_data,
-#             inputs=[
-#             exported_data,
-#             metric_name_dropdown,
-#             normalization_checkbox,
-#             rounding,
-#             grouping_dropdown,
-#             top_select,
-#             direction_checkbox,
-#             group_regex,
-#             log_scale_x_checkbox,
-#             log_scale_y_checkbox,
-#         ],
-#         outputs=[graph_output],
-#     )
-
-
-
-#     datasets_selected.change(
-#         fn=fetch_groups,
-#         inputs=[base_folder, datasets_selected, grouping_dropdown],
-#         outputs=grouping_dropdown,
-#     )
-
-#     grouping_dropdown.select(
-#         fn=fetch_metrics,
-#         inputs=[base_folder, datasets_selected, grouping_dropdown, metric_name_dropdown],
-#         outputs=metric_name_dropdown,
-#     )
-
-#     reverse_grouping_dropdown.select(
-#         fn=partial(fetch_metrics, type="union"),
-#         inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_metric_name_dropdown],
-#         outputs=reverse_metric_name_dropdown,
-#     )
-
-#     reverse_search_button.click(
-#         fn=reverse_search,
-#         inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_metric_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_datasets,
-#         inputs=[base_folder],
-#         outputs=[datasets, datasets_selected, reverse_grouping_dropdown],
-#     )
-
-#     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)}
-#         if not new_dsts:
-#             return gr.update(value=list(selected_runs))
-#         dst_union = new_dsts.union(selected_runs or [])
-#         return gr.update(value=sorted(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 {
-#                 normalization_checkbox: gr.Column(visible=True),
-#                 group_choices: gr.Column(visible=False),
-#             }
-#         else:
-#             return {
-#                 normalization_checkbox: gr.Column(visible=False),
-#                 group_choices: gr.Column(visible=True),
-#             }
-
-#     grouping_dropdown.select(
-#         fn=update_grouping_options,
-#         inputs=[grouping_dropdown],
-#         outputs=[normalization_checkbox, group_choices],
-#     )
-
-
-# # Launch the application
-# if __name__ == "__main__":
-#     demo.launch()
-
-
-import os
-print(os.environ["AWS_ACCESS_KEY_ID"])
+        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,
+                placeholder="Found datasets containing the group/metric name. You can modify the selection after search by removing unwanted lines and clicking Add to selection"
+            )
+
+
+    update_button.click(
+        fn=update_graph,
+        inputs=[
+            base_folder,
+            datasets_selected,
+            metric_name_dropdown,
+            grouping_dropdown,
+            log_scale_x_checkbox,
+            log_scale_y_checkbox,
+            rounding,
+            normalization_checkbox,
+            top_select,
+            direction_checkbox,
+            group_regex,
+        ],
+        outputs=[graph_output, exported_data, export_data_json],
+    )
+
+    for inp in [normalization_checkbox, rounding, group_regex, direction_checkbox, top_select, log_scale_x_checkbox, log_scale_y_checkbox]:
+        inp.change(
+            fn=plot_data,
+            inputs=[
+            exported_data,
+            metric_name_dropdown,
+            normalization_checkbox,
+            rounding,
+            grouping_dropdown,
+            top_select,
+            direction_checkbox,
+            group_regex,
+            log_scale_x_checkbox,
+            log_scale_y_checkbox,
+        ],
+        outputs=[graph_output],
+    )
+
+
+
+    datasets_selected.change(
+        fn=fetch_groups,
+        inputs=[base_folder, datasets_selected, grouping_dropdown],
+        outputs=grouping_dropdown,
+    )
+
+    grouping_dropdown.select(
+        fn=fetch_metrics,
+        inputs=[base_folder, datasets_selected, grouping_dropdown, metric_name_dropdown],
+        outputs=metric_name_dropdown,
+    )
+
+    reverse_grouping_dropdown.select(
+        fn=partial(fetch_metrics, type="union"),
+        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_metric_name_dropdown],
+        outputs=reverse_metric_name_dropdown,
+    )
+
+    reverse_search_button.click(
+        fn=reverse_search,
+        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_metric_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_datasets,
+        inputs=[base_folder],
+        outputs=[datasets, datasets_selected, reverse_grouping_dropdown],
+    )
+
+    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)}
+        if not new_dsts:
+            return gr.update(value=list(selected_runs))
+        dst_union = new_dsts.union(selected_runs or [])
+        return gr.update(value=sorted(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 {
+                normalization_checkbox: gr.Column(visible=True),
+                group_choices: gr.Column(visible=False),
+            }
+        else:
+            return {
+                normalization_checkbox: gr.Column(visible=False),
+                group_choices: gr.Column(visible=True),
+            }
+
+    grouping_dropdown.select(
+        fn=update_grouping_options,
+        inputs=[grouping_dropdown],
+        outputs=[normalization_checkbox, group_choices],
+    )
+
+
+# Launch the application
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+datatrove[dev] @ git+https://github.com/huggingface/datatrove.git
+plotly