app.py

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()