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 import gradio as gr 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: 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, ) 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, ], 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()