UI overhaul + seapration of concerns

app.py

@@ -1,618 +1,4 @@
-import heapq
-import json
-import os
-import re
-import tempfile
-from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor
-from functools import partial
-from pathlib import Path
-from typing import Literal
+from src.view.view import create_interface
 
-import gradio as gr
-import numpy as np
-import plotly.express as px
-import plotly.graph_objects as go
-import tenacity
-from datatrove.io import get_datafolder
-from datatrove.utils.stats import MetricStatsDict
-
-PARTITION_OPTIONS = Literal["Top", "Bottom", "Most frequent (n_docs)"]
-METRICS_LOCATION_DEFAULT = os.getenv("METRICS_LOCATION_DEFAULT", "hf://datasets/HuggingFaceFW-Dev/summary-stats-files")
-
-
-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: sorted([{"value": key, **value} for key, value in dt.to_dict().items()], key=lambda x: x["value"])
-            for name, dt in exported_data.items()
-        }, temp, indent=2)
-        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
-
-    if not value and len(new_choices) == 1:
-        value = list(new_choices)[0]
-
-    # 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"):
-    if not group:
-        return gr.update(choices=[], value=None)
-
-    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
-
-    if not value and len(new_possibles_choices) == 1:
-        value = list(new_possibles_choices)[0]
-
-    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,
-        cumsum: bool,
-        perc: bool,
-        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]
-        if cumsum:
-            y = np.cumsum(y).tolist()
-        if perc:
-            y = (np.array(y) * 100).tolist()
-
-        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 get_desc(data):
-    res = {name: list(dt.to_dict().keys()) for name, dt in data.items()}
-    return "\n".join([
-        f"**{name}**: min={min(dt)}, max={max(dt)}" for name, dt in res.items()
-    ])
-
-
-def update_graph(
-        base_folder,
-        datasets,
-        metric_name,
-        grouping,
-        log_scale_x,
-        log_scale_y,
-        rounding,
-        normalization,
-        top_k,
-        direction,
-        regex,
-        cumsum,
-        perc,
-        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, cumsum, perc, progress), data, export_data(data, metric_name), get_desc(data)
-
-
-def plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y,
-              cumsum, perc,
-              progress=gr.Progress()):
-    if rounding is None or top_k is None:
-        return None
-    graph_fc = (
-        partial(plot_scatter, normalization=normalization, rounding=rounding, cumsum=cumsum, perc=perc)
-        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):
-            cumsum_checkbox = gr.Checkbox(
-                label="Cumsum",
-                value=False,
-            )
-            perc_checkbox = gr.Checkbox(
-                label="%",
-                value=False,
-            )
-            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(visible=False) as min_max_hist:
-            min_max_hist_data = gr.Markdown()
-        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")
-
-    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.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,
-            cumsum_checkbox,
-            perc_checkbox
-        ],
-        outputs=[graph_output, exported_data, export_data_json, min_max_hist_data],
-    )
-
-    gr.on(
-        triggers=[normalization_checkbox.change, rounding.change, group_regex.change, direction_checkbox.change,
-                  top_select.change, log_scale_x_checkbox.change,
-                  log_scale_y_checkbox.change, cumsum_checkbox.change, perc_checkbox.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,
-            cumsum_checkbox,
-            perc_checkbox
-        ],
-        outputs=[graph_output],
-    )
-
-    datasets_selected.change(
-        fn=fetch_groups,
-        inputs=[base_folder, datasets_selected, grouping_dropdown],
-        outputs=grouping_dropdown,
-    )
-
-    grouping_dropdown.change(
-        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),
-                min_max_hist: gr.Column(visible=True),
-            }
-        else:
-            return {
-                normalization_checkbox: gr.Column(visible=False),
-                group_choices: gr.Column(visible=True),
-                min_max_hist: gr.Column(visible=False),
-            }
-
-
-    grouping_dropdown.change(
-        fn=update_grouping_options,
-        inputs=[grouping_dropdown],
-        outputs=[normalization_checkbox, group_choices, min_max_hist],
-    )
-
-# Launch the application
-if __name__ == "__main__":
-    demo.launch()
+demo = create_interface()
+demo.launch()

src/logic/data_fetching.py

@@ -0,0 +1,108 @@
+import os
+import json
+import tempfile
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Dict
+from datatrove.io import get_datafolder
+from datatrove.utils.stats import MetricStatsDict
+import gradio as gr
+import tenacity
+
+def find_folders(base_folder: str, path: str) -> List[str]:
+    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) -> List[str]:
+    base_data_df = get_datafolder(base_folder)
+    dirs = sorted(
+        folder
+        for folder, info in base_data_df.find("", detail=True, maxdepth=1, withdirs=True).items()
+            if info["type"] == "directory"
+        )
+    return sorted(list(set(dirs)))
+
+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 fetch_groups(base_folder: str, datasets: List[str], old_groups: str, type: str = "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
+
+    if not value and len(new_choices) == 1:
+        value = list(new_choices)[0]
+
+    return gr.update(choices=sorted(list(new_choices)), value=value)
+
+def fetch_metrics(base_folder: str, datasets: List[str], group: str, old_metrics: str, type: str = "intersection"):
+    if not group:
+        return gr.update(choices=[], value=None)
+
+    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
+
+    if not value and len(new_possibles_choices) == 1:
+        value = list(new_possibles_choices)[0]
+
+    return gr.update(choices=sorted(list(new_possibles_choices)), value=value)
+
+def reverse_search(base_folder: str, possible_datasets: List[str], grouping: str, metric_name: str) -> str:
+    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: List[str], reverse_search_results: str) -> List[str]:
+    datasets = datasets or []
+    return sorted(list(set(datasets + reverse_search_results.strip().split("\n"))))
+
+def metric_exists(base_folder: str, path: str, metric_name: str, group_by: str) -> bool:
+    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: str, path: str, metric_name: str, group_by: str) -> MetricStatsDict:
+    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)
+        return MetricStatsDict.from_dict(json_metric)
+
+def load_data(dataset_path: str, base_folder: str, grouping: str, metric_name: str) -> MetricStatsDict:
+    return load_metrics(base_folder, dataset_path, metric_name, grouping)

src/logic/data_processing.py

@@ -0,0 +1,47 @@
+import json
+import re
+import heapq
+from collections import defaultdict
+import tempfile
+from typing import Dict, Tuple, List, Literal
+import gradio as gr
+from datatrove.utils.stats import MetricStatsDict
+
+PARTITION_OPTIONS = Literal["Top", "Bottom", "Most frequent (n_docs)"]
+
+def prepare_for_non_grouped_plotting(metric: Dict[str, MetricStatsDict], normalization: bool, rounding: int) -> Dict[float, float]:
+    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()}
+        assert abs(sum(metrics_rounded.values()) - 1) < 0.01
+    return metrics_rounded
+
+def prepare_for_group_plotting(metric: Dict[str, MetricStatsDict], top_k: int, direction: PARTITION_OPTIONS, regex: str | None, rounding: int) -> Tuple[List[str], List[float], List[float]]:
+    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()}
+    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 export_data(exported_data: Dict[str, MetricStatsDict], metric_name: str):
+    if not exported_data:
+        return None
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, prefix=metric_name, suffix=".json") as temp:
+        json.dump({
+            name: sorted([{"value": key, **value} for key, value in dt.to_dict().items()], key=lambda x: x["value"])
+            for name, dt in exported_data.items()
+        }, temp, indent=2)
+        temp_path = temp.name
+    return gr.update(visible=True, value=temp_path)

src/logic/plotting.py

@@ -0,0 +1,107 @@
+from functools import partial
+import plotly.express as px
+import plotly.graph_objects as go
+import numpy as np
+import gradio as gr
+from typing import Dict, List
+from .data_processing import prepare_for_non_grouped_plotting, prepare_for_group_plotting
+from .utils import set_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,
+        cumsum: bool,
+        perc: bool,
+        progress: gr.Progress,
+):
+    fig = go.Figure()
+    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]
+        if cumsum:
+            y = np.cumsum(y).tolist()
+        if perc:
+            y = (np.array(y) * 100).tolist()
+
+        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: str,
+        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 plot_data(data, metric_name, normalization, rounding, grouping, top_k, direction, regex, log_scale_x, log_scale_y,
+              cumsum, perc, progress=gr.Progress()):
+    if rounding is None or top_k is None:
+        return None
+    graph_fc = (
+        partial(plot_scatter, normalization=normalization, rounding=rounding, cumsum=cumsum, perc=perc)
+        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)

src/logic/utils.py

@@ -0,0 +1,12 @@
+def set_alpha(color: str, alpha: float) -> str:
+    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
+    return f"rgba({r}, {g}, {b}, {alpha})"
+
+def get_desc(data):
+    res = {name: list(dt.to_dict().keys()) for name, dt in data.items()}
+    return "\n".join([
+        f"**{name}**: min={min(dt)}, max={max(dt)}" for name, dt in res.items()
+    ])

src/view/view.py

@@ -0,0 +1,325 @@
+import gradio as gr
+from src.logic.data_fetching import fetch_datasets, fetch_groups, fetch_metrics, load_data, reverse_search, reverse_search_add
+from src.logic.data_processing import export_data
+from src.logic.plotting import plot_data
+from src.logic.utils import get_desc
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import os
+import re
+
+
+
+METRICS_LOCATION_DEFAULT = os.getenv("METRICS_LOCATION_DEFAULT", "hf://datasets/HuggingFaceFW-Dev/summary-stats-files")
+
+def update_graph(
+        base_folder,
+        datasets,
+        metric_name,
+        grouping,
+        log_scale_x,
+        log_scale_y,
+        rounding,
+        normalization,
+        top_k,
+        direction,
+        regex,
+        cumsum,
+        perc,
+        progress=gr.Progress(),
+):
+    if len(datasets) <= 0 or not metric_name or not grouping:
+        return None
+
+    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, cumsum, perc, progress), data, export_data(data, metric_name), get_desc(data)
+
+def create_interface():
+    with gr.Blocks() as demo:
+        datasets = gr.State([])
+        exported_data = gr.State([])
+        metrics_headline = gr.Markdown(value="# Metrics Exploration")
+
+        with gr.Tabs():
+            with gr.TabItem("Help"):
+                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 "Render Metric" 
+
+
+## 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.TabItem("Metric View"):
+                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,
+                            )
+
+                    with gr.Column(scale=1):
+                        grouping_dropdown = gr.Dropdown(
+                            choices=[],
+                            label="Grouping",
+                            multiselect=False,
+                        )
+                        metric_name_dropdown = gr.Dropdown(
+                            choices=[],
+                            label="Metric name",
+                            multiselect=False,
+                        )
+
+                        render_button = gr.Button("Render Metric", variant="primary")
+
+                with gr.Tabs():
+                    with gr.TabItem("Graph Settings"):
+                        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,
+                            visible=False
+                        )
+                        with gr.Row():
+                            export_data_json = gr.File(visible=False)
+
+                    with gr.TabItem("Grouping Settings"):
+                        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)",
+                                    )
+
+                    with gr.TabItem("Histogram Settings") as histogram_settings:
+                        cdf_checkbox = gr.Checkbox(
+                            label="CDF",
+                            value=False,
+                        )
+                        perc_checkbox = gr.Checkbox(
+                            label="%",
+                            value=False,
+                        )
+                        with gr.Column(visible=False) as min_max_hist:
+                            min_max_hist_data = gr.Markdown()
+
+                with gr.Row():
+                    graph_output = gr.Plot(label="Graph")
+
+            with gr.TabItem("Reverse Metrics Search"):
+                reverse_search_headline = gr.Markdown(value="# Reverse Metrics Search")
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        reverse_grouping_dropdown = gr.Dropdown(
+                            choices=[],
+                            label="Grouping",
+                            multiselect=False,
+                        )
+                        reverse_metric_name_dropdown = gr.Dropdown(
+                            choices=[],
+                            label="Metric Name",
+                            multiselect=False,
+                        )
+                        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"
+                        )
+
+        render_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,
+                cdf_checkbox,
+                perc_checkbox
+            ],
+            outputs=[graph_output, exported_data, export_data_json, min_max_hist_data],
+        )
+
+        gr.on(
+            triggers=[normalization_checkbox.change, rounding.change, group_regex.change, direction_checkbox.change,
+                      top_select.change, log_scale_x_checkbox.change,
+                      log_scale_y_checkbox.change, cdf_checkbox.change, perc_checkbox.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,
+                cdf_checkbox,
+                perc_checkbox
+            ],
+            outputs=[graph_output],
+        )
+
+        datasets_selected.change(
+            fn=fetch_groups,
+            inputs=[base_folder, datasets_selected, grouping_dropdown],
+            outputs=grouping_dropdown,
+        )
+
+        grouping_dropdown.change(
+            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),
+                    min_max_hist: gr.Column(visible=True),
+                    histogram_settings: gr.TabItem(visible=True),
+                }
+            else:
+                return {
+                    normalization_checkbox: gr.Column(visible=False),
+                    group_choices: gr.Column(visible=True),
+                    min_max_hist: gr.Column(visible=False),
+                    histogram_settings: gr.TabItem(visible=False),
+                }
+
+
+        grouping_dropdown.change(
+            fn=update_grouping_options,
+            inputs=[grouping_dropdown],
+            outputs=[normalization_checkbox, group_choices, min_max_hist, histogram_settings],
+        )
+
+    return demo