WalisonCruz/function-gemma
0
1import math2import os3import re4import secrets5import time6from datetime import datetime, timezone7from functools import lru_cache8from pathlib import Path9from typing import TYPE_CHECKING10 11import huggingface_hub12import numpy as np13import pandas as pd14from huggingface_hub.constants import HF_HOME15 16if TYPE_CHECKING:17 from trackio.commit_scheduler import CommitScheduler18 from trackio.dummy_commit_scheduler import DummyCommitScheduler19 20RESERVED_KEYS = ["project", "run", "timestamp", "step", "time", "metrics"]21 22TRACKIO_LOGO_DIR = Path(__file__).parent / "assets"23 24 25def get_logo_urls() -> dict[str, str]:26 """Get logo URLs from environment variables or use defaults."""27 light_url = os.environ.get(28 "TRACKIO_LOGO_LIGHT_URL",29 f"/gradio_api/file={TRACKIO_LOGO_DIR}/trackio_logo_type_light_transparent.png",30 )31 dark_url = os.environ.get(32 "TRACKIO_LOGO_DARK_URL",33 f"/gradio_api/file={TRACKIO_LOGO_DIR}/trackio_logo_type_dark_transparent.png",34 )35 return {"light": light_url, "dark": dark_url}36 37 38def order_metrics_by_plot_preference(metrics: list[str]) -> tuple[list[str], dict]:39 """40 Order metrics based on TRACKIO_PLOT_ORDER environment variable and group them.41 42 Args:43 metrics: List of metric names to order and group44 45 Returns:46 Tuple of (ordered_group_names, grouped_metrics_dict)47 """48 plot_order_env = os.environ.get("TRACKIO_PLOT_ORDER", "")49 if not plot_order_env.strip():50 plot_order = []51 else:52 plot_order = [53 item.strip() for item in plot_order_env.split(",") if item.strip()54 ]55 56 def get_metric_priority(metric: str) -> tuple[int, int, str]:57 if not plot_order:58 return (float("inf"), float("inf"), metric)59 60 group_prefix = metric.split("/")[0] if "/" in metric else "charts"61 no_match_priority = len(plot_order)62 63 group_priority = no_match_priority64 for i, pattern in enumerate(plot_order):65 pattern_group = pattern.split("/")[0] if "/" in pattern else "charts"66 if pattern_group == group_prefix:67 group_priority = i68 break69 70 within_group_priority = no_match_priority71 for i, pattern in enumerate(plot_order):72 if pattern == metric:73 within_group_priority = i74 break75 elif pattern.endswith("/*") and within_group_priority == no_match_priority:76 pattern_prefix = pattern[:-2]77 if metric.startswith(pattern_prefix + "/"):78 within_group_priority = i + len(plot_order)79 80 return (group_priority, within_group_priority, metric)81 82 result = {}83 for metric in metrics:84 if "/" not in metric:85 if "charts" not in result:86 result["charts"] = {"direct_metrics": [], "subgroups": {}}87 result["charts"]["direct_metrics"].append(metric)88 else:89 parts = metric.split("/")90 main_prefix = parts[0]91 if main_prefix not in result:92 result[main_prefix] = {"direct_metrics": [], "subgroups": {}}93 if len(parts) == 2:94 result[main_prefix]["direct_metrics"].append(metric)95 else:96 subprefix = parts[1]97 if subprefix not in result[main_prefix]["subgroups"]:98 result[main_prefix]["subgroups"][subprefix] = []99 result[main_prefix]["subgroups"][subprefix].append(metric)100 101 for group_data in result.values():102 group_data["direct_metrics"].sort(key=get_metric_priority)103 for subgroup_name in group_data["subgroups"]:104 group_data["subgroups"][subgroup_name].sort(key=get_metric_priority)105 106 if "charts" in result and not result["charts"]["direct_metrics"]:107 del result["charts"]108 109 def get_group_priority(group_name: str) -> tuple[int, str]:110 if not plot_order:111 return (float("inf"), group_name)112 113 min_priority = len(plot_order)114 for i, pattern in enumerate(plot_order):115 pattern_group = pattern.split("/")[0] if "/" in pattern else "charts"116 if pattern_group == group_name:117 min_priority = min(min_priority, i)118 return (min_priority, group_name)119 120 ordered_groups = sorted(result.keys(), key=get_group_priority)121 122 return ordered_groups, result123 124 125def persistent_storage_enabled() -> bool:126 return (127 os.environ.get("PERSISTANT_STORAGE_ENABLED") == "true"128 ) # typo in the name of the environment variable129 130 131def _get_trackio_dir() -> Path:132 if persistent_storage_enabled():133 return Path("/data/trackio")134 elif os.environ.get("TRACKIO_DIR"):135 return Path(os.environ.get("TRACKIO_DIR"))136 return Path(HF_HOME) / "trackio"137 138 139TRACKIO_DIR = _get_trackio_dir()140MEDIA_DIR = TRACKIO_DIR / "media"141FILES_DIR = TRACKIO_DIR / "files"142 143 144def get_or_create_project_hash(project: str) -> str:145 hash_path = TRACKIO_DIR / f"{project}.hash"146 if hash_path.exists():147 return hash_path.read_text().strip()148 hash_value = secrets.token_urlsafe(8)149 TRACKIO_DIR.mkdir(parents=True, exist_ok=True)150 hash_path.write_text(hash_value)151 return hash_value152 153 154def generate_readable_name(used_names: list[str], space_id: str | None = None) -> str:155 """156 Generates a random, readable name like "dainty-sunset-0".157 If space_id is provided, generates username-timestamp format instead.158 """159 if space_id is not None:160 username = _get_default_namespace()161 timestamp = int(time.time())162 return f"{username}-{timestamp}"163 adjectives = [164 "dainty",165 "brave",166 "calm",167 "eager",168 "fancy",169 "gentle",170 "happy",171 "jolly",172 "kind",173 "lively",174 "merry",175 "nice",176 "proud",177 "quick",178 "hugging",179 "silly",180 "tidy",181 "witty",182 "zealous",183 "bright",184 "shy",185 "bold",186 "clever",187 "daring",188 "elegant",189 "faithful",190 "graceful",191 "honest",192 "inventive",193 "jovial",194 "keen",195 "lucky",196 "modest",197 "noble",198 "optimistic",199 "patient",200 "quirky",201 "resourceful",202 "sincere",203 "thoughtful",204 "upbeat",205 "valiant",206 "warm",207 "youthful",208 "zesty",209 "adventurous",210 "breezy",211 "cheerful",212 "delightful",213 "energetic",214 "fearless",215 "glad",216 "hopeful",217 "imaginative",218 "joyful",219 "kindly",220 "luminous",221 "mysterious",222 "neat",223 "outgoing",224 "playful",225 "radiant",226 "spirited",227 "tranquil",228 "unique",229 "vivid",230 "wise",231 "zany",232 "artful",233 "bubbly",234 "charming",235 "dazzling",236 "earnest",237 "festive",238 "gentlemanly",239 "hearty",240 "intrepid",241 "jubilant",242 "knightly",243 "lively",244 "magnetic",245 "nimble",246 "orderly",247 "peaceful",248 "quick-witted",249 "robust",250 "sturdy",251 "trusty",252 "upstanding",253 "vibrant",254 "whimsical",255 ]256 nouns = [257 "sunset",258 "forest",259 "river",260 "mountain",261 "breeze",262 "meadow",263 "ocean",264 "valley",265 "sky",266 "field",267 "cloud",268 "star",269 "rain",270 "leaf",271 "stone",272 "flower",273 "bird",274 "tree",275 "wave",276 "trail",277 "island",278 "desert",279 "hill",280 "lake",281 "pond",282 "grove",283 "canyon",284 "reef",285 "bay",286 "peak",287 "glade",288 "marsh",289 "cliff",290 "dune",291 "spring",292 "brook",293 "cave",294 "plain",295 "ridge",296 "wood",297 "blossom",298 "petal",299 "root",300 "branch",301 "seed",302 "acorn",303 "pine",304 "willow",305 "cedar",306 "elm",307 "falcon",308 "eagle",309 "sparrow",310 "robin",311 "owl",312 "finch",313 "heron",314 "crane",315 "duck",316 "swan",317 "fox",318 "wolf",319 "bear",320 "deer",321 "moose",322 "otter",323 "beaver",324 "lynx",325 "hare",326 "badger",327 "butterfly",328 "bee",329 "ant",330 "beetle",331 "dragonfly",332 "firefly",333 "ladybug",334 "moth",335 "spider",336 "worm",337 "coral",338 "kelp",339 "shell",340 "pebble",341 "face",342 "boulder",343 "cobble",344 "sand",345 "wavelet",346 "tide",347 "current",348 "mist",349 ]350 number = 0351 name = f"{adjectives[0]}-{nouns[0]}-{number}"352 while name in used_names:353 number += 1354 adjective = adjectives[number % len(adjectives)]355 noun = nouns[number % len(nouns)]356 name = f"{adjective}-{noun}-{number}"357 return name358 359 360def is_in_notebook():361 """362 Detect if code is running in a notebook environment (Jupyter, Colab, etc.).363 """364 try:365 from IPython import get_ipython366 367 if get_ipython() is not None:368 return get_ipython().__class__.__name__ in [369 "ZMQInteractiveShell", # Jupyter notebook/lab370 "Shell", # IPython terminal371 ] or "google.colab" in str(get_ipython())372 except ImportError:373 pass374 return False375 376 377def block_main_thread_until_keyboard_interrupt():378 try:379 while True:380 time.sleep(0.1)381 except (KeyboardInterrupt, OSError):382 print("Keyboard interruption in main thread... closing dashboard.")383 384 385def simplify_column_names(columns: list[str]) -> dict[str, str]:386 """387 Simplifies column names to first 10 alphanumeric or "/" characters with unique suffixes.388 389 Args:390 columns: List of original column names391 392 Returns:393 Dictionary mapping original column names to simplified names394 """395 simplified_names = {}396 used_names = set()397 398 for col in columns:399 alphanumeric = re.sub(r"[^a-zA-Z0-9/]", "", col)400 base_name = alphanumeric[:10] if alphanumeric else f"col_{len(used_names)}"401 402 final_name = base_name403 suffix = 1404 while final_name in used_names:405 final_name = f"{base_name}_{suffix}"406 suffix += 1407 408 simplified_names[col] = final_name409 used_names.add(final_name)410 411 return simplified_names412 413 414def print_dashboard_instructions(project: str) -> None:415 """416 Prints instructions for viewing the Trackio dashboard.417 418 Args:419 project: The name of the project to show dashboard for.420 """421 ORANGE = "\033[38;5;208m"422 BOLD = "\033[1m"423 RESET = "\033[0m"424 425 print("* View dashboard by running in your terminal:")426 print(f'{BOLD}{ORANGE}trackio show --project "{project}"{RESET}')427 print(f'* or by running in Python: trackio.show(project="{project}")')428 429 430def preprocess_space_and_dataset_ids(431 space_id: str | None, dataset_id: str | None432) -> tuple[str | None, str | None]:433 """434 Preprocesses the Space and Dataset names to ensure they are valid "username/space_id" or "username/dataset_id" format.435 """436 if space_id is not None and "/" not in space_id:437 username = _get_default_namespace()438 space_id = f"{username}/{space_id}"439 if dataset_id is not None and "/" not in dataset_id:440 username = _get_default_namespace()441 dataset_id = f"{username}/{dataset_id}"442 if space_id is not None and dataset_id is None:443 dataset_id = f"{space_id}-dataset"444 return space_id, dataset_id445 446 447def fibo():448 """Generator for Fibonacci backoff: 1, 1, 2, 3, 5, 8, ..."""449 a, b = 1, 1450 while True:451 yield a452 a, b = b, a + b453 454 455def format_timestamp(timestamp_str):456 """Convert ISO timestamp to human-readable format like '3 minutes ago'."""457 if not timestamp_str or pd.isna(timestamp_str):458 return "Unknown"459 460 try:461 created_time = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))462 if created_time.tzinfo is None:463 created_time = created_time.replace(tzinfo=timezone.utc)464 465 now = datetime.now(timezone.utc)466 diff = now - created_time467 468 seconds = int(diff.total_seconds())469 if seconds < 60:470 return "Just now"471 elif seconds < 3600:472 minutes = seconds // 60473 return f"{minutes} minute{'s' if minutes != 1 else ''} ago"474 elif seconds < 86400:475 hours = seconds // 3600476 return f"{hours} hour{'s' if hours != 1 else ''} ago"477 else:478 days = seconds // 86400479 return f"{days} day{'s' if days != 1 else ''} ago"480 except Exception:481 return "Unknown"482 483 484DEFAULT_COLOR_PALETTE = [485 "#A8769B",486 "#E89957",487 "#3B82F6",488 "#10B981",489 "#EF4444",490 "#8B5CF6",491 "#14B8A6",492 "#F59E0B",493 "#EC4899",494 "#06B6D4",495]496 497 498def get_color_palette() -> list[str]:499 """Get the color palette from environment variable or use default."""500 env_palette = os.environ.get("TRACKIO_COLOR_PALETTE")501 if env_palette:502 return [color.strip() for color in env_palette.split(",")]503 return DEFAULT_COLOR_PALETTE504 505 506def get_color_mapping(507 runs: list[str], smoothing: bool, color_palette: list[str] | None = None508) -> dict[str, str]:509 """Generate color mapping for runs, with transparency for original data when smoothing is enabled."""510 if color_palette is None:511 color_palette = get_color_palette()512 513 color_map = {}514 515 for i, run in enumerate(runs):516 base_color = color_palette[i % len(color_palette)]517 518 if smoothing:519 color_map[run] = base_color + "4D"520 color_map[f"{run}_smoothed"] = base_color521 else:522 color_map[run] = base_color523 524 return color_map525 526 527def downsample(528 df: pd.DataFrame,529 x: str,530 y: str,531 color: str | None,532 x_lim: tuple[float | None, float | None] | None = None,533) -> tuple[pd.DataFrame, tuple[float, float] | None]:534 """535 Downsample the dataframe to reduce the number of points plotted.536 Also updates the x-axis limits to the data min/max if either of the x-axis limits are None.537 538 Args:539 df: The dataframe to downsample.540 x: The column name to use for the x-axis.541 y: The column name to use for the y-axis.542 color: The column name to use for the color.543 x_lim: The x-axis limits to use.544 545 Returns:546 A tuple containing the downsampled dataframe and the updated x-axis limits.547 """548 if df.empty:549 if x_lim is not None:550 x_lim = (x_lim[0] or 0, x_lim[1] or 0)551 return df, x_lim552 553 columns_to_keep = [x, y]554 if color is not None and color in df.columns:555 columns_to_keep.append(color)556 df = df[columns_to_keep].copy()557 558 data_x_min = df[x].min()559 data_x_max = df[x].max()560 561 if x_lim is not None:562 x_min, x_max = x_lim563 if x_min is None:564 x_min = data_x_min565 if x_max is None:566 x_max = data_x_max567 updated_x_lim = (x_min, x_max)568 else:569 updated_x_lim = None570 571 n_bins = 100572 573 if color is not None and color in df.columns:574 groups = df.groupby(color)575 else:576 groups = [(None, df)]577 578 downsampled_indices = []579 580 for _, group_df in groups:581 if group_df.empty:582 continue583 584 group_df = group_df.sort_values(x)585 586 if updated_x_lim is not None:587 x_min, x_max = updated_x_lim588 before_point = group_df[group_df[x] < x_min].tail(1)589 after_point = group_df[group_df[x] > x_max].head(1)590 group_df = group_df[(group_df[x] >= x_min) & (group_df[x] <= x_max)]591 else:592 before_point = after_point = None593 x_min = group_df[x].min()594 x_max = group_df[x].max()595 596 if before_point is not None and not before_point.empty:597 downsampled_indices.extend(before_point.index.tolist())598 if after_point is not None and not after_point.empty:599 downsampled_indices.extend(after_point.index.tolist())600 601 if group_df.empty:602 continue603 604 if x_min == x_max:605 min_y_idx = group_df[y].idxmin()606 max_y_idx = group_df[y].idxmax()607 if min_y_idx != max_y_idx:608 downsampled_indices.extend([min_y_idx, max_y_idx])609 else:610 downsampled_indices.append(min_y_idx)611 continue612 613 if len(group_df) < 500:614 downsampled_indices.extend(group_df.index.tolist())615 continue616 617 bins = np.linspace(x_min, x_max, n_bins + 1)618 group_df["bin"] = pd.cut(619 group_df[x], bins=bins, labels=False, include_lowest=True620 )621 622 for bin_idx in group_df["bin"].dropna().unique():623 bin_data = group_df[group_df["bin"] == bin_idx]624 if bin_data.empty:625 continue626 627 min_y_idx = bin_data[y].idxmin()628 max_y_idx = bin_data[y].idxmax()629 630 downsampled_indices.append(min_y_idx)631 if min_y_idx != max_y_idx:632 downsampled_indices.append(max_y_idx)633 634 unique_indices = list(set(downsampled_indices))635 636 downsampled_df = df.loc[unique_indices].copy()637 638 if color is not None:639 downsampled_df = (640 downsampled_df.groupby(color, sort=False)[downsampled_df.columns]641 .apply(lambda group: group.sort_values(x))642 .reset_index(drop=True)643 )644 else:645 downsampled_df = downsampled_df.sort_values(x).reset_index(drop=True)646 647 downsampled_df = downsampled_df.drop(columns=["bin"], errors="ignore")648 649 return downsampled_df, updated_x_lim650 651 652def sort_metrics_by_prefix(metrics: list[str]) -> list[str]:653 """654 Sort metrics by grouping prefixes together for dropdown/list display.655 Metrics without prefixes come first, then grouped by prefix.656 657 Args:658 metrics: List of metric names659 660 Returns:661 List of metric names sorted by prefix662 663 Example:664 Input: ["train/loss", "loss", "train/acc", "val/loss"]665 Output: ["loss", "train/acc", "train/loss", "val/loss"]666 """667 groups = group_metrics_by_prefix(metrics)668 result = []669 670 if "charts" in groups:671 result.extend(groups["charts"])672 673 for group_name in sorted(groups.keys()):674 if group_name != "charts":675 result.extend(groups[group_name])676 677 return result678 679 680def group_metrics_by_prefix(metrics: list[str]) -> dict[str, list[str]]:681 """682 Group metrics by their prefix. Metrics without prefix go to 'charts' group.683 684 Args:685 metrics: List of metric names686 687 Returns:688 Dictionary with prefix names as keys and lists of metrics as values689 690 Example:691 Input: ["loss", "accuracy", "train/loss", "train/acc", "val/loss"]692 Output: {693 "charts": ["loss", "accuracy"],694 "train": ["train/loss", "train/acc"],695 "val": ["val/loss"]696 }697 """698 no_prefix = []699 with_prefix = []700 701 for metric in metrics:702 if "/" in metric:703 with_prefix.append(metric)704 else:705 no_prefix.append(metric)706 707 no_prefix.sort()708 709 prefix_groups = {}710 for metric in with_prefix:711 prefix = metric.split("/")[0]712 if prefix not in prefix_groups:713 prefix_groups[prefix] = []714 prefix_groups[prefix].append(metric)715 716 for prefix in prefix_groups:717 prefix_groups[prefix].sort()718 719 groups = {}720 if no_prefix:721 groups["charts"] = no_prefix722 723 for prefix in sorted(prefix_groups.keys()):724 groups[prefix] = prefix_groups[prefix]725 726 return groups727 728 729def get_sync_status(scheduler: "CommitScheduler | DummyCommitScheduler") -> int | None:730 """Get the sync status from the CommitScheduler in an integer number of minutes, or None if not synced yet."""731 if getattr(732 scheduler, "last_push_time", None733 ): # DummyCommitScheduler doesn't have last_push_time734 time_diff = time.time() - scheduler.last_push_time735 return int(time_diff / 60)736 else:737 return None738 739 740def generate_embed_code(project: str, metrics: str, selected_runs: list = None) -> str:741 """Generate the embed iframe code based on current settings."""742 space_host = os.environ.get("SPACE_HOST", "")743 if not space_host:744 return ""745 746 params = []747 748 if project:749 params.append(f"project={project}")750 751 if metrics and metrics.strip():752 params.append(f"metrics={metrics}")753 754 if selected_runs:755 runs_param = ",".join(selected_runs)756 params.append(f"runs={runs_param}")757 758 params.append("sidebar=hidden")759 params.append("navbar=hidden")760 761 query_string = "&".join(params)762 embed_url = f"https://{space_host}?{query_string}"763 764 return f'<iframe src="{embed_url}" style="width:1600px; height:500px; border:0;"></iframe>'765 766 767def serialize_values(metrics):768 """769 Serialize infinity and NaN values in metrics dict to make it JSON-compliant.770 Only handles top-level float values.771 772 Converts:773 - float('inf') -> "Infinity"774 - float('-inf') -> "-Infinity"775 - float('nan') -> "NaN"776 777 Example:778 {"loss": float('inf'), "accuracy": 0.95} -> {"loss": "Infinity", "accuracy": 0.95}779 """780 if not isinstance(metrics, dict):781 return metrics782 783 result = {}784 for key, value in metrics.items():785 if isinstance(value, float):786 if math.isinf(value):787 result[key] = "Infinity" if value > 0 else "-Infinity"788 elif math.isnan(value):789 result[key] = "NaN"790 else:791 result[key] = value792 elif isinstance(value, np.floating):793 float_val = float(value)794 if math.isinf(float_val):795 result[key] = "Infinity" if float_val > 0 else "-Infinity"796 elif math.isnan(float_val):797 result[key] = "NaN"798 else:799 result[key] = float_val800 else:801 result[key] = value802 return result803 804 805def deserialize_values(metrics):806 """807 Deserialize infinity and NaN string values back to their numeric forms.808 Only handles top-level string values.809 810 Converts:811 - "Infinity" -> float('inf')812 - "-Infinity" -> float('-inf')813 - "NaN" -> float('nan')814 815 Example:816 {"loss": "Infinity", "accuracy": 0.95} -> {"loss": float('inf'), "accuracy": 0.95}817 """818 if not isinstance(metrics, dict):819 return metrics820 821 result = {}822 for key, value in metrics.items():823 if value == "Infinity":824 result[key] = float("inf")825 elif value == "-Infinity":826 result[key] = float("-inf")827 elif value == "NaN":828 result[key] = float("nan")829 else:830 result[key] = value831 return result832 833 834def get_full_url(835 base_url: str, project: str | None, write_token: str, footer: bool = True836) -> str:837 params = []838 if project:839 params.append(f"project={project}")840 params.append(f"write_token={write_token}")841 if not footer:842 params.append("footer=false")843 return base_url + "?" + "&".join(params)844 845 846def embed_url_in_notebook(url: str) -> None:847 try:848 from IPython.display import HTML, display849 850 embed_code = HTML(851 f'<div><iframe src="{url}" width="100%" height="1000px" allow="autoplay; camera; microphone; clipboard-read; clipboard-write;" frameborder="0" allowfullscreen></iframe></div>'852 )853 display(embed_code)854 except ImportError:855 pass856 857 858def to_json_safe(obj):859 if isinstance(obj, (str, int, float, bool, type(None))):860 return obj861 if isinstance(obj, np.generic):862 return obj.item()863 if isinstance(obj, dict):864 return {str(k): to_json_safe(v) for k, v in obj.items()}865 if isinstance(obj, (list, tuple, set)):866 return [to_json_safe(v) for v in obj]867 if hasattr(obj, "to_dict") and callable(obj.to_dict):868 return to_json_safe(obj.to_dict())869 if hasattr(obj, "__dict__"):870 return {871 str(k): to_json_safe(v)872 for k, v in vars(obj).items()873 if not k.startswith("_")874 }875 return str(obj)876 877 878def get_space() -> str | None:879 """880 Get the space ID ("user/space") if Trackio is running in a Space, or None if not.881 """882 return os.environ.get("SPACE_ID")883 884 885def ordered_subset(items: list[str], subset: list[str] | None) -> list[str]:886 subset_set = set(subset or [])887 return [item for item in items if item in subset_set]888 889 890def _get_default_namespace() -> str:891 """Get the default namespace (username).892 893 This function uses caching to avoid repeated API calls to /whoami-v2.894 """895 token = huggingface_hub.get_token()896 return _cached_whoami(token)["name"]897 898 899@lru_cache(maxsize=32)900def _cached_whoami(token: str | None) -> dict:901 return huggingface_hub.whoami(token=token)902