WalisonCruz/function-gemma
0
1import os2from pathlib import Path3 4import pandas as pd5 6from trackio import deploy, utils7from trackio.sqlite_storage import SQLiteStorage8 9 10def import_csv(11 csv_path: str | Path,12 project: str,13 name: str | None = None,14 space_id: str | None = None,15 dataset_id: str | None = None,16 private: bool | None = None,17 force: bool = False,18) -> None:19 """20 Imports a CSV file into a Trackio project. The CSV file must contain a `"step"`21 column, may optionally contain a `"timestamp"` column, and any other columns will be22 treated as metrics. It should also include a header row with the column names.23 24 TODO: call init() and return a Run object so that the user can continue to log metrics to it.25 26 Args:27 csv_path (`str` or `Path`):28 The str or Path to the CSV file to import.29 project (`str`):30 The name of the project to import the CSV file into. Must not be an existing31 project.32 name (`str`, *optional*):33 The name of the Run to import the CSV file into. If not provided, a default34 name will be generated.35 name (`str`, *optional*):36 The name of the run (if not provided, a default name will be generated).37 space_id (`str`, *optional*):38 If provided, the project will be logged to a Hugging Face Space instead of a39 local directory. Should be a complete Space name like `"username/reponame"`40 or `"orgname/reponame"`, or just `"reponame"` in which case the Space will41 be created in the currently-logged-in Hugging Face user's namespace. If the42 Space does not exist, it will be created. If the Space already exists, the43 project will be logged to it.44 dataset_id (`str`, *optional*):45 If provided, a persistent Hugging Face Dataset will be created and the46 metrics will be synced to it every 5 minutes. Should be a complete Dataset47 name like `"username/datasetname"` or `"orgname/datasetname"`, or just48 `"datasetname"` in which case the Dataset will be created in the49 currently-logged-in Hugging Face user's namespace. If the Dataset does not50 exist, it will be created. If the Dataset already exists, the project will51 be appended to it. If not provided, the metrics will be logged to a local52 SQLite database, unless a `space_id` is provided, in which case a Dataset53 will be automatically created with the same name as the Space but with the54 `"_dataset"` suffix.55 private (`bool`, *optional*):56 Whether to make the Space private. If None (default), the repo will be57 public unless the organization's default is private. This value is ignored58 if the repo already exists.59 """60 if SQLiteStorage.get_runs(project):61 raise ValueError(62 f"Project '{project}' already exists. Cannot import CSV into existing project."63 )64 65 csv_path = Path(csv_path)66 if not csv_path.exists():67 raise FileNotFoundError(f"CSV file not found: {csv_path}")68 69 df = pd.read_csv(csv_path)70 if df.empty:71 raise ValueError("CSV file is empty")72 73 column_mapping = utils.simplify_column_names(df.columns.tolist())74 df = df.rename(columns=column_mapping)75 76 step_column = None77 for col in df.columns:78 if col.lower() == "step":79 step_column = col80 break81 82 if step_column is None:83 raise ValueError("CSV file must contain a 'step' or 'Step' column")84 85 if name is None:86 name = csv_path.stem87 88 metrics_list = []89 steps = []90 timestamps = []91 92 numeric_columns = []93 for column in df.columns:94 if column == step_column:95 continue96 if column == "timestamp":97 continue98 99 try:100 pd.to_numeric(df[column], errors="raise")101 numeric_columns.append(column)102 except (ValueError, TypeError):103 continue104 105 for _, row in df.iterrows():106 metrics = {}107 for column in numeric_columns:108 value = row[column]109 if bool(pd.notna(value)):110 metrics[column] = float(value)111 112 if metrics:113 metrics_list.append(metrics)114 steps.append(int(row[step_column]))115 116 if "timestamp" in df.columns and bool(pd.notna(row["timestamp"])):117 timestamps.append(str(row["timestamp"]))118 else:119 timestamps.append("")120 121 if metrics_list:122 SQLiteStorage.bulk_log(123 project=project,124 run=name,125 metrics_list=metrics_list,126 steps=steps,127 timestamps=timestamps,128 )129 130 print(131 f"* Imported {len(metrics_list)} rows from {csv_path} into project '{project}' as run '{name}'"132 )133 print(f"* Metrics found: {', '.join(metrics_list[0].keys())}")134 135 space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)136 if dataset_id is not None:137 os.environ["TRACKIO_DATASET_ID"] = dataset_id138 print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")139 140 if space_id is None:141 utils.print_dashboard_instructions(project)142 else:143 deploy.create_space_if_not_exists(144 space_id=space_id, dataset_id=dataset_id, private=private145 )146 deploy.wait_until_space_exists(space_id=space_id)147 deploy.upload_db_to_space(project=project, space_id=space_id, force=force)148 print(149 f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"150 )151 152 153def import_tf_events(154 log_dir: str | Path,155 project: str,156 name: str | None = None,157 space_id: str | None = None,158 dataset_id: str | None = None,159 private: bool | None = None,160 force: bool = False,161) -> None:162 """163 Imports TensorFlow Events files from a directory into a Trackio project. Each164 subdirectory in the log directory will be imported as a separate run.165 166 Args:167 log_dir (`str` or `Path`):168 The str or Path to the directory containing TensorFlow Events files.169 project (`str`):170 The name of the project to import the TensorFlow Events files into. Must not171 be an existing project.172 name (`str`, *optional*):173 The name prefix for runs (if not provided, will use directory names). Each174 subdirectory will create a separate run.175 space_id (`str`, *optional*):176 If provided, the project will be logged to a Hugging Face Space instead of a177 local directory. Should be a complete Space name like `"username/reponame"`178 or `"orgname/reponame"`, or just `"reponame"` in which case the Space will179 be created in the currently-logged-in Hugging Face user's namespace. If the180 Space does not exist, it will be created. If the Space already exists, the181 project will be logged to it.182 dataset_id (`str`, *optional*):183 If provided, a persistent Hugging Face Dataset will be created and the184 metrics will be synced to it every 5 minutes. Should be a complete Dataset185 name like `"username/datasetname"` or `"orgname/datasetname"`, or just186 `"datasetname"` in which case the Dataset will be created in the187 currently-logged-in Hugging Face user's namespace. If the Dataset does not188 exist, it will be created. If the Dataset already exists, the project will189 be appended to it. If not provided, the metrics will be logged to a local190 SQLite database, unless a `space_id` is provided, in which case a Dataset191 will be automatically created with the same name as the Space but with the192 `"_dataset"` suffix.193 private (`bool`, *optional*):194 Whether to make the Space private. If None (default), the repo will be195 public unless the organization's default is private. This value is ignored196 if the repo already exists.197 """198 try:199 from tbparse import SummaryReader200 except ImportError:201 raise ImportError(202 "The `tbparse` package is not installed but is required for `import_tf_events`. Please install trackio with the `tensorboard` extra: `pip install trackio[tensorboard]`."203 )204 205 if SQLiteStorage.get_runs(project):206 raise ValueError(207 f"Project '{project}' already exists. Cannot import TF events into existing project."208 )209 210 path = Path(log_dir)211 if not path.exists():212 raise FileNotFoundError(f"TF events directory not found: {path}")213 214 # Use tbparse to read all tfevents files in the directory structure215 reader = SummaryReader(str(path), extra_columns={"dir_name"})216 df = reader.scalars217 218 if df.empty:219 raise ValueError(f"No TensorFlow events data found in {path}")220 221 total_imported = 0222 imported_runs = []223 224 # Group by dir_name to create separate runs225 for dir_name, group_df in df.groupby("dir_name"):226 try:227 # Determine run name based on directory name228 if dir_name == "":229 run_name = "main" # For files in the root directory230 else:231 run_name = dir_name # Use directory name232 233 if name:234 run_name = f"{name}_{run_name}"235 236 if group_df.empty:237 print(f"* Skipping directory {dir_name}: no scalar data found")238 continue239 240 metrics_list = []241 steps = []242 timestamps = []243 244 for _, row in group_df.iterrows():245 # Convert row values to appropriate types246 tag = str(row["tag"])247 value = float(row["value"])248 step = int(row["step"])249 250 metrics = {tag: value}251 metrics_list.append(metrics)252 steps.append(step)253 254 # Use wall_time if present, else fallback255 if "wall_time" in group_df.columns and not bool(256 pd.isna(row["wall_time"])257 ):258 timestamps.append(str(row["wall_time"]))259 else:260 timestamps.append("")261 262 if metrics_list:263 SQLiteStorage.bulk_log(264 project=project,265 run=str(run_name),266 metrics_list=metrics_list,267 steps=steps,268 timestamps=timestamps,269 )270 271 total_imported += len(metrics_list)272 imported_runs.append(run_name)273 274 print(275 f"* Imported {len(metrics_list)} scalar events from directory '{dir_name}' as run '{run_name}'"276 )277 print(f"* Metrics in this run: {', '.join(set(group_df['tag']))}")278 279 except Exception as e:280 print(f"* Error processing directory {dir_name}: {e}")281 continue282 283 if not imported_runs:284 raise ValueError("No valid TensorFlow events data could be imported")285 286 print(f"* Total imported events: {total_imported}")287 print(f"* Created runs: {', '.join(imported_runs)}")288 289 space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)290 if dataset_id is not None:291 os.environ["TRACKIO_DATASET_ID"] = dataset_id292 print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")293 294 if space_id is None:295 utils.print_dashboard_instructions(project)296 else:297 deploy.create_space_if_not_exists(298 space_id, dataset_id=dataset_id, private=private299 )300 deploy.wait_until_space_exists(space_id)301 deploy.upload_db_to_space(project, space_id, force=force)302 print(303 f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"304 )305