Tonic/g-android-control

1import re2import sys3import time4from pathlib import Path5 6import huggingface_hub7import numpy as np8import pandas as pd9from huggingface_hub.constants import HF_HOME10 11RESERVED_KEYS = ["project", "run", "timestamp", "step", "time", "metrics"]12TRACKIO_DIR = Path(HF_HOME) / "trackio"13 14TRACKIO_LOGO_DIR = Path(__file__).parent / "assets"15 16 17def generate_readable_name(used_names: list[str]) -> str:18    """19    Generates a random, readable name like "dainty-sunset-0"20    """21    adjectives = [22        "dainty",23        "brave",24        "calm",25        "eager",26        "fancy",27        "gentle",28        "happy",29        "jolly",30        "kind",31        "lively",32        "merry",33        "nice",34        "proud",35        "quick",36        "hugging",37        "silly",38        "tidy",39        "witty",40        "zealous",41        "bright",42        "shy",43        "bold",44        "clever",45        "daring",46        "elegant",47        "faithful",48        "graceful",49        "honest",50        "inventive",51        "jovial",52        "keen",53        "lucky",54        "modest",55        "noble",56        "optimistic",57        "patient",58        "quirky",59        "resourceful",60        "sincere",61        "thoughtful",62        "upbeat",63        "valiant",64        "warm",65        "youthful",66        "zesty",67        "adventurous",68        "breezy",69        "cheerful",70        "delightful",71        "energetic",72        "fearless",73        "glad",74        "hopeful",75        "imaginative",76        "joyful",77        "kindly",78        "luminous",79        "mysterious",80        "neat",81        "outgoing",82        "playful",83        "radiant",84        "spirited",85        "tranquil",86        "unique",87        "vivid",88        "wise",89        "zany",90        "artful",91        "bubbly",92        "charming",93        "dazzling",94        "earnest",95        "festive",96        "gentlemanly",97        "hearty",98        "intrepid",99        "jubilant",100        "knightly",101        "lively",102        "magnetic",103        "nimble",104        "orderly",105        "peaceful",106        "quick-witted",107        "robust",108        "sturdy",109        "trusty",110        "upstanding",111        "vibrant",112        "whimsical",113    ]114    nouns = [115        "sunset",116        "forest",117        "river",118        "mountain",119        "breeze",120        "meadow",121        "ocean",122        "valley",123        "sky",124        "field",125        "cloud",126        "star",127        "rain",128        "leaf",129        "stone",130        "flower",131        "bird",132        "tree",133        "wave",134        "trail",135        "island",136        "desert",137        "hill",138        "lake",139        "pond",140        "grove",141        "canyon",142        "reef",143        "bay",144        "peak",145        "glade",146        "marsh",147        "cliff",148        "dune",149        "spring",150        "brook",151        "cave",152        "plain",153        "ridge",154        "wood",155        "blossom",156        "petal",157        "root",158        "branch",159        "seed",160        "acorn",161        "pine",162        "willow",163        "cedar",164        "elm",165        "falcon",166        "eagle",167        "sparrow",168        "robin",169        "owl",170        "finch",171        "heron",172        "crane",173        "duck",174        "swan",175        "fox",176        "wolf",177        "bear",178        "deer",179        "moose",180        "otter",181        "beaver",182        "lynx",183        "hare",184        "badger",185        "butterfly",186        "bee",187        "ant",188        "beetle",189        "dragonfly",190        "firefly",191        "ladybug",192        "moth",193        "spider",194        "worm",195        "coral",196        "kelp",197        "shell",198        "pebble",199        "face",200        "boulder",201        "cobble",202        "sand",203        "wavelet",204        "tide",205        "current",206        "mist",207    ]208    number = 0209    name = f"{adjectives[0]}-{nouns[0]}-{number}"210    while name in used_names:211        number += 1212        adjective = adjectives[number % len(adjectives)]213        noun = nouns[number % len(nouns)]214        name = f"{adjective}-{noun}-{number}"215    return name216 217 218def block_except_in_notebook():219    in_notebook = bool(getattr(sys, "ps1", sys.flags.interactive))220    if in_notebook:221        return222    try:223        while True:224            time.sleep(0.1)225    except (KeyboardInterrupt, OSError):226        print("Keyboard interruption in main thread... closing dashboard.")227 228 229def simplify_column_names(columns: list[str]) -> dict[str, str]:230    """231    Simplifies column names to first 10 alphanumeric or "/" characters with unique suffixes.232 233    Args:234        columns: List of original column names235 236    Returns:237        Dictionary mapping original column names to simplified names238    """239    simplified_names = {}240    used_names = set()241 242    for col in columns:243        alphanumeric = re.sub(r"[^a-zA-Z0-9/]", "", col)244        base_name = alphanumeric[:10] if alphanumeric else f"col_{len(used_names)}"245 246        final_name = base_name247        suffix = 1248        while final_name in used_names:249            final_name = f"{base_name}_{suffix}"250            suffix += 1251 252        simplified_names[col] = final_name253        used_names.add(final_name)254 255    return simplified_names256 257 258def print_dashboard_instructions(project: str) -> None:259    """260    Prints instructions for viewing the Trackio dashboard.261 262    Args:263        project: The name of the project to show dashboard for.264    """265    YELLOW = "\033[93m"266    BOLD = "\033[1m"267    RESET = "\033[0m"268 269    print("* View dashboard by running in your terminal:")270    print(f'{BOLD}{YELLOW}trackio show --project "{project}"{RESET}')271    print(f'* or by running in Python: trackio.show(project="{project}")')272 273 274def preprocess_space_and_dataset_ids(275    space_id: str | None, dataset_id: str | None276) -> tuple[str | None, str | None]:277    if space_id is not None and "/" not in space_id:278        username = huggingface_hub.whoami()["name"]279        space_id = f"{username}/{space_id}"280    if dataset_id is not None and "/" not in dataset_id:281        username = huggingface_hub.whoami()["name"]282        dataset_id = f"{username}/{dataset_id}"283    if space_id is not None and dataset_id is None:284        dataset_id = f"{space_id}_dataset"285    return space_id, dataset_id286 287 288def fibo():289    """Generator for Fibonacci backoff: 1, 1, 2, 3, 5, 8, ..."""290    a, b = 1, 1291    while True:292        yield a293        a, b = b, a + b294 295 296COLOR_PALETTE = [297    "#3B82F6",298    "#EF4444",299    "#10B981",300    "#F59E0B",301    "#8B5CF6",302    "#EC4899",303    "#06B6D4",304    "#84CC16",305    "#F97316",306    "#6366F1",307]308 309 310def get_color_mapping(runs: list[str], smoothing: bool) -> dict[str, str]:311    """Generate color mapping for runs, with transparency for original data when smoothing is enabled."""312    color_map = {}313 314    for i, run in enumerate(runs):315        base_color = COLOR_PALETTE[i % len(COLOR_PALETTE)]316 317        if smoothing:318            color_map[f"{run}_smoothed"] = base_color319            color_map[f"{run}_original"] = base_color + "4D"320        else:321            color_map[run] = base_color322 323    return color_map324 325 326def downsample(327    df: pd.DataFrame,328    x: str,329    y: str,330    color: str | None,331    x_lim: tuple[float, float] | None = None,332) -> pd.DataFrame:333    if df.empty:334        return df335 336    columns_to_keep = [x, y]337    if color is not None and color in df.columns:338        columns_to_keep.append(color)339    df = df[columns_to_keep].copy()340 341    n_bins = 100342 343    if color is not None and color in df.columns:344        groups = df.groupby(color)345    else:346        groups = [(None, df)]347 348    downsampled_indices = []349 350    for _, group_df in groups:351        if group_df.empty:352            continue353 354        group_df = group_df.sort_values(x)355 356        if x_lim is not None:357            x_min, x_max = x_lim358            before_point = group_df[group_df[x] < x_min].tail(1)359            after_point = group_df[group_df[x] > x_max].head(1)360            group_df = group_df[(group_df[x] >= x_min) & (group_df[x] <= x_max)]361        else:362            before_point = after_point = None363            x_min = group_df[x].min()364            x_max = group_df[x].max()365 366        if before_point is not None and not before_point.empty:367            downsampled_indices.extend(before_point.index.tolist())368        if after_point is not None and not after_point.empty:369            downsampled_indices.extend(after_point.index.tolist())370 371        if group_df.empty:372            continue373 374        if x_min == x_max:375            min_y_idx = group_df[y].idxmin()376            max_y_idx = group_df[y].idxmax()377            if min_y_idx != max_y_idx:378                downsampled_indices.extend([min_y_idx, max_y_idx])379            else:380                downsampled_indices.append(min_y_idx)381            continue382 383        if len(group_df) < 500:384            downsampled_indices.extend(group_df.index.tolist())385            continue386 387        bins = np.linspace(x_min, x_max, n_bins + 1)388        group_df["bin"] = pd.cut(389            group_df[x], bins=bins, labels=False, include_lowest=True390        )391 392        for bin_idx in group_df["bin"].dropna().unique():393            bin_data = group_df[group_df["bin"] == bin_idx]394            if bin_data.empty:395                continue396 397            min_y_idx = bin_data[y].idxmin()398            max_y_idx = bin_data[y].idxmax()399 400            downsampled_indices.append(min_y_idx)401            if min_y_idx != max_y_idx:402                downsampled_indices.append(max_y_idx)403 404    unique_indices = list(set(downsampled_indices))405 406    downsampled_df = df.loc[unique_indices].copy()407    downsampled_df = downsampled_df.sort_values(x).reset_index(drop=True)408    downsampled_df = downsampled_df.drop(columns=["bin"], errors="ignore")409 410    return downsampled_df411