katira/function-gemma

1import glob2import json3import logging4import os5import warnings6import webbrowser7from pathlib import Path8from typing import Any9 10import huggingface_hub11from gradio.themes import ThemeClass12from gradio.utils import TupleNoPrint13from gradio_client import Client, handle_file14from huggingface_hub import SpaceStorage15from huggingface_hub.errors import LocalTokenNotFoundError16 17from trackio import context_vars, deploy, utils18from trackio.deploy import sync19from trackio.histogram import Histogram20from trackio.imports import import_csv, import_tf_events21from trackio.media import TrackioAudio, TrackioImage, TrackioVideo22from trackio.run import Run23from trackio.sqlite_storage import SQLiteStorage24from trackio.table import Table25from trackio.typehints import UploadEntry26from trackio.ui.main import CSS, HEAD, demo27from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR28 29logging.getLogger("httpx").setLevel(logging.WARNING)30 31warnings.filterwarnings(32    "ignore",33    message="Empty session being created. Install gradio\\[oauth\\]",34    category=UserWarning,35    module="gradio.helpers",36)37 38__version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[39    "version"40]41 42__all__ = [43    "init",44    "log",45    "finish",46    "show",47    "sync",48    "delete_project",49    "import_csv",50    "import_tf_events",51    "save",52    "Image",53    "Video",54    "Audio",55    "Table",56    "Histogram",57]58 59Image = TrackioImage60Video = TrackioVideo61Audio = TrackioAudio62 63 64config = {}65 66 67def init(68    project: str,69    name: str | None = None,70    group: str | None = None,71    space_id: str | None = None,72    space_storage: SpaceStorage | None = None,73    dataset_id: str | None = None,74    config: dict | None = None,75    resume: str = "never",76    settings: Any = None,77    private: bool | None = None,78    embed: bool = True,79) -> Run:80    """81    Creates a new Trackio project and returns a [`Run`] object.82 83    Args:84        project (`str`):85            The name of the project (can be an existing project to continue tracking or86            a new project to start tracking from scratch).87        name (`str`, *optional*):88            The name of the run (if not provided, a default name will be generated).89        group (`str`, *optional*):90            The name of the group which this run belongs to in order to help organize91            related runs together. You can toggle the entire group's visibilitiy in the92            dashboard.93        space_id (`str`, *optional*):94            If provided, the project will be logged to a Hugging Face Space instead of95            a local directory. Should be a complete Space name like96            `"username/reponame"` or `"orgname/reponame"`, or just `"reponame"` in which97            case the Space will be created in the currently-logged-in Hugging Face98            user's namespace. If the Space does not exist, it will be created. If the99            Space already exists, the project will be logged to it.100        space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):101            Choice of persistent storage tier.102        dataset_id (`str`, *optional*):103            If a `space_id` is provided, a persistent Hugging Face Dataset will be104            created and the metrics will be synced to it every 5 minutes. Specify a105            Dataset with name like `"username/datasetname"` or `"orgname/datasetname"`,106            or `"datasetname"` (uses currently-logged-in Hugging Face user's namespace),107            or `None` (uses the same name as the Space but with the `"_dataset"`108            suffix). If the Dataset does not exist, it will be created. If the Dataset109            already exists, the project will be appended to it.110        config (`dict`, *optional*):111            A dictionary of configuration options. Provided for compatibility with112            `wandb.init()`.113        resume (`str`, *optional*, defaults to `"never"`):114            Controls how to handle resuming a run. Can be one of:115 116            - `"must"`: Must resume the run with the given name, raises error if run117              doesn't exist118            - `"allow"`: Resume the run if it exists, otherwise create a new run119            - `"never"`: Never resume a run, always create a new one120        private (`bool`, *optional*):121            Whether to make the Space private. If None (default), the repo will be122            public unless the organization's default is private. This value is ignored123            if the repo already exists.124        settings (`Any`, *optional*):125            Not used. Provided for compatibility with `wandb.init()`.126        embed (`bool`, *optional*, defaults to `True`):127            If running inside a jupyter/Colab notebook, whether the dashboard should128            automatically be embedded in the cell when trackio.init() is called.129 130    Returns:131        `Run`: A [`Run`] object that can be used to log metrics and finish the run.132    """133    if settings is not None:134        warnings.warn(135            "* Warning: settings is not used. Provided for compatibility with wandb.init(). Please create an issue at: https://github.com/gradio-app/trackio/issues if you need a specific feature implemented."136        )137 138    if space_id is None and dataset_id is not None:139        raise ValueError("Must provide a `space_id` when `dataset_id` is provided.")140    try:141        space_id, dataset_id = utils.preprocess_space_and_dataset_ids(142            space_id, dataset_id143        )144    except LocalTokenNotFoundError as e:145        raise LocalTokenNotFoundError(146            f"You must be logged in to Hugging Face locally when `space_id` is provided to deploy to a Space. {e}"147        ) from e148    url = context_vars.current_server.get()149    share_url = context_vars.current_share_server.get()150 151    if url is None:152        if space_id is None:153            _, url, share_url = demo.launch(154                css=CSS,155                head=HEAD,156                footer_links=["gradio", "settings"],157                inline=False,158                quiet=True,159                prevent_thread_lock=True,160                show_error=True,161                favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",162                allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],163            )164            context_vars.current_space_id.set(None)165        else:166            url = space_id167            share_url = None168            context_vars.current_space_id.set(space_id)169 170        context_vars.current_server.set(url)171        context_vars.current_share_server.set(share_url)172    if (173        context_vars.current_project.get() is None174        or context_vars.current_project.get() != project175    ):176        print(f"* Trackio project initialized: {project}")177 178        if dataset_id is not None:179            os.environ["TRACKIO_DATASET_ID"] = dataset_id180            print(181                f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}"182            )183        if space_id is None:184            print(f"* Trackio metrics logged to: {TRACKIO_DIR}")185            if utils.is_in_notebook() and embed:186                base_url = share_url + "/" if share_url else url187                full_url = utils.get_full_url(188                    base_url, project=project, write_token=demo.write_token, footer=True189                )190                utils.embed_url_in_notebook(full_url)191            else:192                utils.print_dashboard_instructions(project)193        else:194            deploy.create_space_if_not_exists(195                space_id, space_storage, dataset_id, private196            )197            user_name, space_name = space_id.split("/")198            space_url = deploy.SPACE_HOST_URL.format(199                user_name=user_name, space_name=space_name200            )201            print(f"* View dashboard by going to: {space_url}")202            if utils.is_in_notebook() and embed:203                utils.embed_url_in_notebook(space_url)204    context_vars.current_project.set(project)205 206    client = None207    if not space_id:208        client = Client(url, verbose=False)209 210    if resume == "must":211        if name is None:212            raise ValueError("Must provide a run name when resume='must'")213        if name not in SQLiteStorage.get_runs(project):214            raise ValueError(f"Run '{name}' does not exist in project '{project}'")215        resumed = True216    elif resume == "allow":217        resumed = name is not None and name in SQLiteStorage.get_runs(project)218    elif resume == "never":219        if name is not None and name in SQLiteStorage.get_runs(project):220            warnings.warn(221                f"* Warning: resume='never' but a run '{name}' already exists in "222                f"project '{project}'. Generating a new name and instead. If you want "223                "to resume this run, call init() with resume='must' or resume='allow'."224            )225            name = None226        resumed = False227    else:228        raise ValueError("resume must be one of: 'must', 'allow', or 'never'")229 230    run = Run(231        url=url,232        project=project,233        client=client,234        name=name,235        group=group,236        config=config,237        space_id=space_id,238    )239 240    if resumed:241        print(f"* Resumed existing run: {run.name}")242    else:243        print(f"* Created new run: {run.name}")244 245    context_vars.current_run.set(run)246    globals()["config"] = run.config247    return run248 249 250def log(metrics: dict, step: int | None = None) -> None:251    """252    Logs metrics to the current run.253 254    Args:255        metrics (`dict`):256            A dictionary of metrics to log.257        step (`int`, *optional*):258            The step number. If not provided, the step will be incremented259            automatically.260    """261    run = context_vars.current_run.get()262    if run is None:263        raise RuntimeError("Call trackio.init() before trackio.log().")264    run.log(265        metrics=metrics,266        step=step,267    )268 269 270def finish():271    """272    Finishes the current run.273    """274    run = context_vars.current_run.get()275    if run is None:276        raise RuntimeError("Call trackio.init() before trackio.finish().")277    run.finish()278 279 280def delete_project(project: str, force: bool = False) -> bool:281    """282    Deletes a project by removing its local SQLite database.283 284    Args:285        project (`str`):286            The name of the project to delete.287        force (`bool`, *optional*, defaults to `False`):288            If `True`, deletes the project without prompting for confirmation.289            If `False`, prompts the user to confirm before deleting.290 291    Returns:292        `bool`: `True` if the project was deleted, `False` otherwise.293    """294    db_path = SQLiteStorage.get_project_db_path(project)295 296    if not db_path.exists():297        print(f"* Project '{project}' does not exist.")298        return False299 300    if not force:301        response = input(302            f"Are you sure you want to delete project '{project}'? "303            f"This will permanently delete all runs and metrics. (y/N): "304        )305        if response.lower() not in ["y", "yes"]:306            print("* Deletion cancelled.")307            return False308 309    try:310        db_path.unlink()311 312        for suffix in ("-wal", "-shm"):313            sidecar = Path(str(db_path) + suffix)314            if sidecar.exists():315                sidecar.unlink()316 317        print(f"* Project '{project}' has been deleted.")318        return True319    except Exception as e:320        print(f"* Error deleting project '{project}': {e}")321        return False322 323 324def save(325    glob_str: str | Path,326    project: str | None = None,327) -> str:328    """329    Saves files to a project (not linked to a specific run). If Trackio is running330    locally, the file(s) will be moved to the project's files directory. If Trackio is331    running in a Space, the file(s) will be uploaded to the Space's files directory.332 333    Args:334        glob_str (`str` or `Path`):335            The file path or glob pattern to save. Can be a single file or a pattern336            matching multiple files (e.g., `"*.py"`, `"models/**/*.pth"`).337        project (`str`, *optional*):338            The name of the project to save files to. If not provided, uses the current339            project from `trackio.init()`. If no project is initialized, raises an340            error.341 342    Returns:343        `str`: The path where the file(s) were saved (project's files directory).344 345    Example:346        ```python347        import trackio348 349        trackio.init(project="my-project")350        trackio.save("config.yaml")351        trackio.save("models/*.pth")352        ```353    """354    if project is None:355        project = context_vars.current_project.get()356        if project is None:357            raise RuntimeError(358                "No project specified. Either call trackio.init() first or provide a "359                "project parameter to trackio.save()."360            )361 362    glob_str = Path(glob_str)363    base_path = Path.cwd().resolve()364 365    matched_files = []366    if glob_str.is_file():367        matched_files = [glob_str.resolve()]368    else:369        pattern = str(glob_str)370        if not glob_str.is_absolute():371            pattern = str((Path.cwd() / glob_str).resolve())372        matched_files = [373            Path(f).resolve()374            for f in glob.glob(pattern, recursive=True)375            if Path(f).is_file()376        ]377 378    if not matched_files:379        raise ValueError(f"No files found matching pattern: {glob_str}")380 381    url = context_vars.current_server.get()382    current_run = context_vars.current_run.get()383 384    upload_entries = []385 386    for file_path in matched_files:387        try:388            relative_to_base = file_path.relative_to(base_path)389        except ValueError:390            relative_to_base = Path(file_path.name)391 392        if current_run is not None:393            # If a run is active, use its queue to upload the file to the project's files directory394            # as it's more efficent than uploading files one by one. But we should not use the run name395            # as the files should be stored in the project's files directory, not the run's, hence396            # the use_run_name flag is set to False.397            current_run._queue_upload(398                file_path,399                step=None,400                relative_path=str(relative_to_base.parent),401                use_run_name=False,402            )403        else:404            upload_entry: UploadEntry = {405                "project": project,406                "run": None,407                "step": None,408                "relative_path": str(relative_to_base),409                "uploaded_file": handle_file(file_path),410            }411            upload_entries.append(upload_entry)412 413    if upload_entries:414        if url is None:415            raise RuntimeError(416                "No server available. Call trackio.init() before trackio.save() to start the server."417            )418 419        try:420            client = Client(url, verbose=False, httpx_kwargs={"timeout": 90})421            client.predict(422                api_name="/bulk_upload_media",423                uploads=upload_entries,424                hf_token=huggingface_hub.utils.get_token(),425            )426        except Exception as e:427            warnings.warn(428                f"Failed to upload files: {e}. "429                "Files may not be available in the dashboard."430            )431 432    return str(utils.MEDIA_DIR / project / "files")433 434 435def show(436    project: str | None = None,437    *,438    theme: str | ThemeClass | None = None,439    mcp_server: bool | None = None,440    footer: bool = True,441    color_palette: list[str] | None = None,442    open_browser: bool = True,443    block_thread: bool | None = None,444):445    """446    Launches the Trackio dashboard.447 448    Args:449        project (`str`, *optional*):450            The name of the project whose runs to show. If not provided, all projects451            will be shown and the user can select one.452        theme (`str` or `ThemeClass`, *optional*):453            A Gradio Theme to use for the dashboard instead of the default Gradio theme,454            can be a built-in theme (e.g. `'soft'`, `'citrus'`), a theme from the Hub455            (e.g. `"gstaff/xkcd"`), or a custom Theme class. If not provided, the456            `TRACKIO_THEME` environment variable will be used, or if that is not set,457            the default Gradio theme will be used.458        mcp_server (`bool`, *optional*):459            If `True`, the Trackio dashboard will be set up as an MCP server and certain460            functions will be added as MCP tools. If `None` (default behavior), then the461            `GRADIO_MCP_SERVER` environment variable will be used to determine if the462            MCP server should be enabled (which is `"True"` on Hugging Face Spaces).463        footer (`bool`, *optional*, defaults to `True`):464            Whether to show the Gradio footer. When `False`, the footer will be hidden.465            This can also be controlled via the `footer` query parameter in the URL.466        color_palette (`list[str]`, *optional*):467            A list of hex color codes to use for plot lines. If not provided, the468            `TRACKIO_COLOR_PALETTE` environment variable will be used (comma-separated469            hex codes), or if that is not set, the default color palette will be used.470            Example: `['#FF0000', '#00FF00', '#0000FF']`471        open_browser (`bool`, *optional*, defaults to `True`):472            If `True` and not in a notebook, a new browser tab will be opened with the473            dashboard. If `False`, the browser will not be opened.474        block_thread (`bool`, *optional*):475            If `True`, the main thread will be blocked until the dashboard is closed.476            If `None` (default behavior), then the main thread will not be blocked if the477            dashboard is launched in a notebook, otherwise the main thread will be blocked.478 479        Returns:480            `app`: The Gradio app object corresponding to the dashboard launched by Trackio.481            `url`: The local URL of the dashboard.482            `share_url`: The public share URL of the dashboard.483            `full_url`: The full URL of the dashboard including the write token (will use the public share URL if launched publicly, otherwise the local URL).484    """485    if color_palette is not None:486        os.environ["TRACKIO_COLOR_PALETTE"] = ",".join(color_palette)487 488    theme = theme or os.environ.get("TRACKIO_THEME")489 490    _mcp_server = (491        mcp_server492        if mcp_server is not None493        else os.environ.get("GRADIO_MCP_SERVER", "False") == "True"494    )495 496    app, url, share_url = demo.launch(497        css=CSS,498        head=HEAD,499        footer_links=["gradio", "settings"] + (["api"] if _mcp_server else []),500        quiet=True,501        inline=False,502        prevent_thread_lock=True,503        favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",504        allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],505        mcp_server=_mcp_server,506        theme=theme,507    )508 509    base_url = share_url + "/" if share_url else url510    full_url = utils.get_full_url(511        base_url, project=project, write_token=demo.write_token, footer=footer512    )513 514    if not utils.is_in_notebook():515        print(f"* Trackio UI launched at: {full_url}")516        if open_browser:517            webbrowser.open(full_url)518        block_thread = block_thread if block_thread is not None else True519    else:520        utils.embed_url_in_notebook(full_url)521        block_thread = block_thread if block_thread is not None else False522 523    if block_thread:524        utils.block_main_thread_until_keyboard_interrupt()525    return TupleNoPrint((demo, url, share_url, full_url))526