metric naming

app.py

@@ -5,29 +5,21 @@ 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 DataFolder, get_datafolder
+from datatrove.io import get_datafolder
 import plotly.graph_objects as go
-from datatrove.utils.stats import MetricStatsDict
+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)"]
-
-LOG_SCALE_STATS = {
-    "length",
-    "n_lines",
-    "n_docs",
-    "n_words",
-    "avg_words_per_line",
-    "pages_with_lorem_ipsum",
-}
-
-STATS_LOCATION_DEFAULT = os.getenv("STATS_LOCATION_DEFAULT", "s3://")
+METRICS_LOCATION_DEFAULT = os.getenv("METRICS_LOCATION_DEFAULT", "s3://fineweb-stats/summary/")
 
 
 def find_folders(base_folder, path):
@@ -43,28 +35,31 @@ def find_folders(base_folder, path):
     )
 
 
-def find_stats_folders(base_folder: str):
+def find_metrics_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_data_folder.glob("**/stats-merged.json")
+    # First find all metric.json using globing for metric.json
+    metrics_merged = base_data_folder.glob("**/metric.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]
+    # 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(stats_folders)))
+    return sorted(list(set(metrics_folders)))
 
 
 def fetch_datasets(base_folder: str):
-    datasets = sorted(find_stats_folders(base_folder))
+    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, stat_name):
+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=stat_name, suffix=".json") as temp:
-        json.dump(exported_data, temp)
+    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)
 
@@ -80,7 +75,7 @@ def fetch_groups(base_folder, datasets, old_groups, type="intersection"):
 
     if type == "intersection":
         new_choices = set.intersection(*(set(g) for g in GROUPS)) 
-    elif type == "union":
+    else:
         new_choices = set.union(*(set(g) for g in GROUPS))
     value = None
     if old_groups:
@@ -91,27 +86,27 @@ def fetch_groups(base_folder, datasets, old_groups, type="intersection"):
     return gr.update(choices=sorted(list(new_choices)), value=value)
 
 
-def fetch_stats(base_folder, datasets, group, old_stats, type="intersection"):
+def fetch_metrics(base_folder, datasets, group, old_metrics, type="intersection"):
     with ThreadPoolExecutor() as executor:
-        STATS = list(executor.map(lambda run: [Path(x).name for x in find_folders(base_folder, f"{run}/{group}")], datasets))
-    if len(STATS) == 0:
+        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 STATS))
-    elif type == "union":
-        new_possibles_choices = set.union(*(set(s) for s in STATS))
+        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_stats:
-        value = list(set.intersection(new_possibles_choices, {old_stats}))
+    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, stat_name):
+def reverse_search(base_folder, possible_datasets, grouping, metric_name):
     with ThreadPoolExecutor() as executor:
-        found_datasets = list(executor.map(lambda dataset: dataset if stat_exists(base_folder, dataset, stat_name, grouping) else None, possible_datasets))
+        found_datasets = 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)
 
@@ -122,46 +117,47 @@ def reverse_search_add(datasets, reverse_search_results):
 
 
 
-def stat_exists(base_folder, path, stat_name, group_by):
+def metric_exists(base_folder, path, metric_name, group_by):
     base_folder = get_datafolder(base_folder)
-    return base_folder.exists(f"{path}/{group_by}/{stat_name}/stats-merged.json")
+    return base_folder.exists(f"{path}/{group_by}/{metric_name}/metric.json")
 
-def load_stats(base_folder, path, stat_name, group_by):
+@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}/{stat_name}/stats-merged.json",
+        f"{path}/{group_by}/{metric_name}/metric.json",
     ) 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)
+        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_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[round(float(key), 2)] += value.total
+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(stats_rounded.values())
-        stats_rounded = {k: v / normalizer for k, v in stats_rounded.items()}
+        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(stats_rounded.values())
-    return stats_rounded
+        summed = sum(metrics_rounded.values())
+        assert abs(summed - 1) < 0.01, summed
+    return metrics_rounded
 
 
-def prepare_grouped_data(dataset_path, base_folder, grouping, stat_name, top_k, direction: PARTITION_OPTIONS, regex):
-    import heapq
-    regex_compiled = re.compile(regex) if regex else None
-
-    stats = load_stats(base_folder, dataset_path, stat_name, grouping)
-    stats = {key: value for key, value in stats.items() if not regex or regex_compiled.match(key)}
-    means = {key: value.mean for key, value in stats.items()}
+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: value.n for key, value in stats.items()}
+        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)
@@ -181,23 +177,23 @@ def set_alpha(color, alpha):
     return f"rgba({r}, {g}, {b}, {alpha})"
 
 
-
-
 def plot_scatter(
-    histograms: dict[str, dict[float, float]],
-    stat_name: str,
+    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()
 
-    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:
-            x = sorted(histogram.keys())
-
-        y = [histogram[k] for k in x]
+    # 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(
@@ -209,14 +205,14 @@ def plot_scatter(
             )
         )
 
-    xaxis_scale = "log" if stat_name in LOG_SCALE_STATS else "linear"
     yaxis_title = "Frequency" if normalization else "Total"
 
     fig.update_layout(
-        title=f"Line Plots for {stat_name}",
-        xaxis_title=stat_name,
+        title=f"Line Plots for {metric_name}",
+        xaxis_title=metric_name,
         yaxis_title=yaxis_title,
-        xaxis_type=xaxis_scale,
+        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,
@@ -226,22 +222,33 @@ def plot_scatter(
 
 
 def plot_bars(
-    histograms: dict[str, list[tuple[str, float]]],
-    stat_name: str,
+    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(histograms.items(), total=len(histograms), desc="Plotting...")):
-        x = [k for k, v in histogram]
-        y = [v for k, v in histogram]
+    for i, (name, histogram) in enumerate(progress.tqdm(data.items(), total=len(data), desc="Plotting...")):
+        histogram_prepared = prepare_for_group_plotting(histogram, top_k, direction, regex, rounding)
+        x, y = zip(*histogram_prepared)
 
         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}",
-        xaxis_title=stat_name,
-        yaxis_title="Mean value",
+        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,
@@ -254,33 +261,26 @@ def plot_bars(
 def update_graph(
     base_folder,
     datasets,
-    stat_name,
+    metric_name,
     grouping,
+    log_scale_x,
+    log_scale_y,
+    rounding,
     normalization,
     top_k,
     direction,
     regex,
     progress=gr.Progress(),
 ):
-    if len(datasets) <= 0 or not stat_name or not grouping:
+    if len(datasets) <= 0 or not metric_name or not grouping:
         return None
     # Placeholder for logic to rerender the graph based on the inputs
-    prepare_fc = (
-        partial(prepare_non_grouped_data, normalization=normalization)
-        if grouping == "histogram"
-        else partial(prepare_grouped_data, top_k=top_k, direction=direction, regex=regex)
-    )
-    graph_fc = (
-        partial(plot_scatter, normalization=normalization)
-        if grouping == "histogram"
-        else plot_bars
-    )
 
     with ThreadPoolExecutor() as pool:
         data = list(
             progress.tqdm(
                 pool.map(
-                    partial(prepare_fc, base_folder=base_folder, stat_name=stat_name, grouping=grouping),
+                    partial(load_data, base_folder=base_folder, metric_name=metric_name, grouping=grouping),
                     datasets,
                 ),
                 total=len(datasets),
@@ -288,30 +288,39 @@ def update_graph(
             )
         )
 
-    histograms = {path: result for path, result in zip(datasets, 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)
 
-    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([])
-    stats_headline = gr.Markdown(value="# Stats Exploration")
+    metrics_headline = gr.Markdown(value="# Metrics Exploration")
     with gr.Row():
         with gr.Column(scale=2):
-            # Define the multiselect for crawls
             with gr.Row():
                 with gr.Column(scale=1):
                     base_folder = gr.Textbox(
-                        label="Stats Location",
-                        value="s3://fineweb-stats/summary/",
+                        label="Metrics Location",
+                        value=METRICS_LOCATION_DEFAULT,
                     )
                     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")
+                    regex_select = gr.Text(label="Regex filter", value=".*")
+                    regex_button = gr.Button("Search")
             with gr.Row():
                 datasets_selected = gr.Dropdown(
                     choices=[],
@@ -323,15 +332,28 @@ with gr.Blocks() as demo:
             readme_description = gr.Markdown(
                 label="Readme",
                 value="""
-Explaination of the tool:
+## 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 occurences. If normalization is on, the values are frequencies summing to 1.
-- (fqdn/suffix): creates a bar plot of the mean values of the stats for full qualied domain name/suffix of domain
+## 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: the top/bottom k groups are shown
-- summary: simply shows the average value of given stat for selected crawls
-                """,
+    * 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
@@ -340,19 +362,39 @@ Groupings:
                 label="Grouping",
                 multiselect=False,
             )
-            # Define the dropdown for stat_name
-            stat_name_dropdown = gr.Dropdown(
+            # Define the dropdown for metric_name
+            metric_name_dropdown = gr.Dropdown(
                 choices=[],
-                label="Stat name",
+                label="Metric name",
                 multiselect=False,
             )
 
-            with gr.Row(visible=False) as histogram_choices:
-                normalization_checkbox = gr.Checkbox(
-                    label="Normalize",
-                    value=False,  # Default value
-                )
 
+            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(
@@ -374,19 +416,13 @@ Groupings:
                                 "Most frequent (n_docs)",
                             ],
                             value="Most frequent (n_docs)",
-                        )
-
-            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
+    with gr.Row():
         graph_output = gr.Plot(label="Graph")
     
     with gr.Row():
-        reverse_search_headline = gr.Markdown(value="# Reverse stats search")
+        reverse_search_headline = gr.Markdown(value="# Reverse metrics search")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -396,8 +432,8 @@ Groupings:
                 label="Grouping",
                 multiselect=False,
             )
-            # Define the dropdown for stat_name
-            reverse_stat_name_dropdown = gr.Dropdown(
+            # Define the dropdown for metric_name
+            reverse_metric_name_dropdown = gr.Dropdown(
                 choices=[],
                 label="Stat name",
                 multiselect=False,
@@ -411,9 +447,8 @@ Groupings:
             reverse_search_results = gr.Textbox(
                 label="Found datasets",
                 lines=10,
-                placeholder="Found datasets containing the group/stat name. You can modify the selection after search by removing unwanted lines and clicking Add to selection"
+                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(
@@ -421,22 +456,39 @@ Groupings:
         inputs=[
             base_folder,
             datasets_selected,
-            stat_name_dropdown,
+            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_button],
+        outputs=[graph_output, exported_data, export_data_json],
     )
 
-    export_data_button.click(
-        fn=export_data,
-        inputs=[exported_data, stat_name_dropdown],
-        outputs=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],
@@ -444,20 +496,20 @@ Groupings:
     )
 
     grouping_dropdown.select(
-        fn=fetch_stats,
-        inputs=[base_folder, datasets_selected, grouping_dropdown, stat_name_dropdown],
-        outputs=stat_name_dropdown,
+        fn=fetch_metrics,
+        inputs=[base_folder, datasets_selected, grouping_dropdown, metric_name_dropdown],
+        outputs=metric_name_dropdown,
     )
 
     reverse_grouping_dropdown.select(
-        fn=partial(fetch_stats, type="union"),
-        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_stat_name_dropdown],
-        outputs=reverse_stat_name_dropdown,
+        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_stat_name_dropdown],
+        inputs=[base_folder, datasets, reverse_grouping_dropdown, reverse_metric_name_dropdown],
         outputs=reverse_search_results,
     )
 
@@ -478,7 +530,9 @@ Groupings:
         if not regex:
             return
         new_dsts = {run for run in all_runs if re.search(regex, run)}
-        dst_union = new_dsts.union(selected_runs)
+        if not new_dsts:
+            return gr.update(value=list(selected_runs))
+        dst_union = new_dsts.union(selected_runs or [])
         return gr.update(value=list(dst_union))
 
     regex_button.click(
@@ -490,19 +544,19 @@ Groupings:
     def update_grouping_options(grouping):
         if grouping == "histogram":
             return {
-                histogram_choices: gr.Column(visible=True),
+                normalization_checkbox: gr.Column(visible=True),
                 group_choices: gr.Column(visible=False),
             }
         else:
             return {
-                histogram_choices: gr.Column(visible=False),
+                normalization_checkbox: gr.Column(visible=False),
                 group_choices: gr.Column(visible=True),
             }
 
     grouping_dropdown.select(
         fn=update_grouping_options,
         inputs=[grouping_dropdown],
-        outputs=[histogram_choices, group_choices],
+        outputs=[normalization_checkbox, group_choices],
     )