first commit

app.py

@@ -0,0 +1,331 @@
+import ast
+from collections import defaultdict
+from functools import partial
+import itertools
+import os
+import re
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+from datetime import datetime
+
+import gradio as gr
+import huggingface_hub
+import pandas as pd
+import plotly.graph_objects as go
+from huggingface_hub.file_download import repo_folder_name
+from huggingface_hub.hf_api import RepoFile
+from huggingface_hub.utils import EntryNotFoundError
+
+FALLBACK_TOKEN_NAME = "HF_TOKEN"
+
+def is_arary_like(x):
+    return isinstance(x, list) or isinstance(x, tuple) or isinstance(x, np.ndarray)
+
+def get_task_type(df):
+    if all(isinstance(pred, str) for pred in df['predictions'].iloc[0]):
+        return "generative"
+    if all(is_arary_like(pred) and all(isinstance(item, float) for item in pred) for pred in df['predictions'].iloc[0]):
+        return "multiple_choice"
+    return "mixed"
+
+def fix_df(df):
+    # For some reason some metrics and predictions are stored as strings
+    for col in ["predictions", "metrics", "choices", "gold", "gold_index"]:
+        df[col] = [ast.literal_eval(x) if isinstance(x, str) else x for x in df[col].values]
+    return df
+
+def get_run_name_seed(run_name):
+    if "-seed-" not in run_name:
+        return run_name, 5
+    run_name, seed = run_name.split("-seed-")
+    return run_name, int(seed)
+
+def fetch_repo_structure(repo_name, oauth_token: gr.OAuthToken | None = None):
+    token = os.environ.get(FALLBACK_TOKEN_NAME)
+    if oauth_token:
+        token = oauth_token.token
+
+    files = list(huggingface_hub.list_repo_tree(repo_name, "details", recursive=False, token=token))
+    
+    runs = {file.path.split('/')[-1] for file in files if isinstance(file, huggingface_hub.hf_api.RepoFolder)}
+    if not runs:
+        return {}, gr.update(choices=[], value=None)
+
+    def process_run(run):
+        run_files = list(huggingface_hub.list_repo_tree(repo_name, f"details/{run}", recursive=False, token=token))
+        return run, [file.path.split('/')[-1] for file in run_files if isinstance(file, huggingface_hub.hf_api.RepoFolder)]
+
+    with ThreadPoolExecutor() as executor:
+        results = list(executor.map(process_run, runs))
+
+    checkpoints_dict = dict(results)
+
+    return checkpoints_dict, gr.update(choices=list(checkpoints_dict), value=None)
+
+def update_checkpoints(selected_runs, checkpoints):
+    if not selected_runs:
+        return gr.update(choices=[], value=None)
+    
+    common_checkpoints = set(checkpoints[selected_runs[0]])
+    for run in selected_runs[1:]:
+        common_checkpoints.intersection_update(set(checkpoints[run]))
+    
+    common_checkpoints = sorted(list(common_checkpoints))
+    
+    return gr.update(choices=common_checkpoints, value=common_checkpoints[0] if common_checkpoints else None)
+
+
+def select_runs_by_regex(runs, current_selected, regex_to_select):
+    comp_re = re.compile(regex_to_select)
+    return list(sorted(set((current_selected if current_selected else []) +
+                           [run for run in runs if comp_re.fullmatch(run)])))
+
+def select_runs_by_language(runs, current_selected, language):
+    if language:
+        return select_runs_by_regex(runs, current_selected, f".*-{language}-.*")
+    return current_selected
+
+def fetch_available_tasks(repo_name, runs_to_fetch, checkpoint) -> dict[str, dict[str, str]]:
+    token = os.environ.get(FALLBACK_TOKEN_NAME)
+    
+    all_tasks = defaultdict(lambda: defaultdict(dict))
+    for run in runs_to_fetch:
+        try:
+            files = huggingface_hub.list_repo_tree(repo_name, f"details/{run}/{checkpoint}", token=token)
+            parquet_files = [f.path.split('/')[-1] for f in files if f.path.endswith('.parquet')]
+            
+            for full_filename in parquet_files:
+                task_name, date_str = full_filename.replace('.parquet', '').rsplit('_', 1)
+                date = datetime.strptime(date_str, '%Y-%m-%dT%H-%M-%S.%f')
+                
+                if run not in all_tasks[task_name] or date > all_tasks[task_name][run]['date']:
+                    all_tasks[task_name][run] = {'filename': full_filename, 'date': date}
+        except EntryNotFoundError:
+            print(f"Checkpoint not found for run: {run}")
+    
+    available_tasks = {
+        task: {run: info['filename'] for run, info in runs.items()}
+        for task, runs in all_tasks.items()
+        if set(runs.keys()) == set(runs_to_fetch)
+    }
+    
+    return available_tasks
+
+def fetch_run_results(repo_name, runs_to_fetch, checkpoint,
+                      oauth_token: gr.OAuthToken | None = None, progress=gr.Progress()):
+
+    task_runs_dict = fetch_available_tasks(repo_name, runs_to_fetch, checkpoint)
+    task_names = list(task_runs_dict.keys())
+    return gr.update(choices=task_names, value=task_names[0] if task_names else None), task_runs_dict
+
+
+def filter_with_metric(df, selected_runs, metric_name):
+    if df is None or not selected_runs or not metric_name:
+        return None
+    kept_metrics = [f"metric_{metric_name}_{run_name}" for run_name in selected_runs]
+    other_metrics = [col for col in df.columns if col.startswith(f"metric_") and col not in kept_metrics]
+    df = df.drop(columns=other_metrics)
+    widths = get_column_widths(df)
+    df = consize_runname_metric(df, selected_runs, metric_name)
+    return gr.update(value=df, column_widths=widths)
+
+def get_column_widths(df):
+    column_widths = []
+    for col in df.columns:
+        if col == "full_prompt":
+            column_widths.append("300px")
+        elif col in ["choices", "gold"]:
+            column_widths.append("250px")
+        elif col.startswith("metric_"):
+            column_widths.append("100px")
+        else:
+            column_widths.append("200px")  # Default width for other columns
+    return column_widths
+
+
+def consize_runname_metric(df, run_names, metric_name):
+    """
+    Turns metric columns (metric_{metric}_{run_name}) into {metric}_i
+    """
+    # Initialize the new column with empty strings
+    for idx, run_name in enumerate(run_names):
+        original_column = f"metric_{metric_name}_{run_name}"
+        if original_column in df.columns:
+            # Append the run name and metric value to the concise column
+            df[f"{metric_name}_{idx}"] = df[original_column]
+            df = df.drop(columns=[original_column])
+    return df
+
+
+def load_task_data(repo_name, runs_to_fetch, checkpoint, task_name, tasks_files, progress=gr.Progress()):
+    token = os.environ.get(FALLBACK_TOKEN_NAME)
+    if not runs_to_fetch or not task_name:
+        return None, None, None
+    
+    def fetch_run_file(run_to_fetch):
+        file_path = f"details/{run_to_fetch}/{checkpoint}/{tasks_files[task_name][run_to_fetch]}"
+        try:
+            cached_path = huggingface_hub.hf_hub_download(repo_name, file_path, token=token)
+            df = pd.read_parquet(cached_path)
+            return df, run_to_fetch
+        except EntryNotFoundError:
+            print(f"File not found: {file_path}")
+            return None, run_to_fetch
+
+    with ThreadPoolExecutor() as pool:
+        results = list(progress.tqdm(pool.map(fetch_run_file, runs_to_fetch), total=len(runs_to_fetch),
+                                     desc="Fetching run data..."))
+    
+    dfs = [fix_df(df) for df, _ in results if df is not None]
+    run_names = [run for _, run in results if run is not None]
+
+    if not dfs:
+        return None, None, gr.update(choices=[], value=None)
+    
+    task_type = get_task_type(dfs[0])
+    def prepare_df(df, run_name, task_type):
+        def get_choice_predictions(df, task_type):
+            # For some evals it's string for other it's list
+            predictions = df['predictions']
+            if task_type == "generative":
+                return predictions
+            
+            if task_type == "multiple_choice":
+                n_choices = len(df['choices'])
+                return df['choices'][np.argmax([pred[0] for pred in predictions[:n_choices]])]
+            
+            if task_type == "mixed":
+                return predictions[0]
+            
+            return predictions
+
+        prepared_df = pd.DataFrame({
+            'full_prompt': df['full_prompt'],
+            f'{run_name}': df.apply(partial(get_choice_predictions, task_type=task_type), axis=1)
+        })
+        # For some reason some metrics are stored as strings
+        metrics = df['metrics']
+        # Assume all metrics are the same
+        for metric_key in metrics[0].keys():
+            prepared_df[f'metric_{metric_key}_{run_name}'] = [metric[metric_key] for metric in metrics]
+        return prepared_df.set_index('full_prompt')
+    
+    def get_gold_label(df, task_type):
+        if task_type == "generative":
+            return df['gold']
+        return [df['choices'][idx] for idx in df['gold_index']]
+
+    # Prepare the first DataFrame with choices and gold
+    combined_df = dfs[0][['full_prompt', 'choices']].set_index('full_prompt')
+    combined_df['gold'] = dfs[0].apply(lambda row: get_gold_label(row, task_type), axis=1).values
+    
+    # Join all prepared DataFrames
+    for df, run_name in zip(dfs, run_names):
+        prepared_df = prepare_df(df, run_name, task_type)
+        combined_df = combined_df.join(prepared_df, how='outer', )
+        
+
+    available_metrics = list(set("_".join(col.split('_')[1:-1]) for col in combined_df.columns if col.startswith("metric_")))
+    combined_df = combined_df.reset_index()
+
+    return combined_df, filter_with_metric(combined_df, runs_to_fetch, available_metrics[0]), gr.update(choices=available_metrics, value=available_metrics[0])
+
+def render_results_table(df: pd.DataFrame):
+    if df is None or df.empty:
+        return None
+    
+    # Select a subset of 100 examples
+    df_subset = df.sample(n=min(100, len(df)), random_state=42)
+    
+    # Prepare the data for display
+    display_data = []
+    for _, row in df_subset.iterrows():
+        example_data = {
+            'text': row['example'],
+            'choices': row['choices'],
+            'gold_index': row['gold_index'],
+        }
+        for run in df['run'].unique():
+            run_data = df[(df['run'] == run) & (df['example'] == row['example'])]
+            if not run_data.empty:
+                example_data[f'{run}_prediction'] = run_data['predictions'].values[0]
+                example_data[f'{run}_score'] = run_data['metrics'].values[0]
+        display_data.append(example_data)
+    
+    return pd.DataFrame(display_data)
+
+with gr.Blocks() as demo:
+    runs_checkpoints = gr.State({})
+    results_df_full = gr.State(None)
+    tasks_files = gr.State({})
+    login_button = gr.LoginButton(visible=False)
+    repo = gr.Textbox(label="HF Repo", value="HuggingFaceFW-Dev/multiligual-ablation-logs-dev", visible=True)
+    with gr.Column():
+        gr.Markdown("# FineWeb experiments results explorer")
+        with gr.Row():
+            with gr.Column():
+                select_by_regex_text = gr.Textbox(label="Regex to select runs",
+                                                  value="ind_minhash(-CC-MAIN-|_)\\d{4}-\\d{2}-seed.*")
+                select_by_regex_button = gr.Button("Select matching runs")
+            with gr.Column():
+                select_by_language = gr.Dropdown(choices=["ar", "fr", "ru", "hi", "th", "tr", "zh", "sw", "te"], 
+                                                 interactive=True, label="Select by language", 
+                                                 info="Choose a language to prefill the regex")
+        selected_runs = gr.Dropdown(choices=[], interactive=True, multiselect=True, label="Selected runs")
+        checkpoint = gr.Dropdown(choices=[], interactive=True, label="Checkpoint")
+        fetch_res = gr.Button("Fetch results")
+        task_name = gr.Dropdown(choices=[], interactive=True, label="Task name")
+        metric_name = gr.Dropdown(choices=[], interactive=True, label="Metric")
+        results_df = gr.Dataframe(interactive=False, wrap=True)
+
+    # Run selection
+    gr.on(
+        triggers=[repo.change],
+        fn=fetch_repo_structure, inputs=[repo], outputs=[runs_checkpoints, selected_runs],
+    )
+    gr.on(
+        triggers=[select_by_regex_button.click],
+        fn=select_runs_by_regex,
+        inputs=[runs_checkpoints, selected_runs, select_by_regex_text], outputs=[selected_runs]
+    )
+    gr.on(
+        triggers=[select_by_language.change],
+        fn=select_runs_by_language,
+        inputs=[runs_checkpoints, selected_runs, select_by_language], outputs=[selected_runs]
+    )
+    
+    # Update checkpoints based on selected runs
+    gr.on(
+        triggers=[selected_runs.change],
+        fn=update_checkpoints,
+        inputs=[selected_runs, runs_checkpoints],
+        outputs=[checkpoint]
+    )
+    
+    # Fetch available tasks
+    gr.on(
+        triggers=[fetch_res.click],
+        fn=fetch_run_results,
+        inputs=[repo, selected_runs, checkpoint],
+        outputs=[task_name, tasks_files]
+    )
+    
+
+    # Update results when task name or metric changes
+    gr.on(
+        triggers=[task_name.change],
+        fn=load_task_data,
+        inputs=[repo, selected_runs, checkpoint, task_name, tasks_files],
+        outputs=[results_df_full, results_df, metric_name]
+    )
+    
+    gr.on(
+        triggers=[metric_name.change],
+        fn=filter_with_metric,
+        inputs=[results_df_full, selected_runs, metric_name],
+        outputs=[results_df]
+    )
+    
+    demo.load(fn=fetch_repo_structure, inputs=[repo], outputs=[runs_checkpoints, selected_runs])
+
+demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+pandas
+plotly