WalisonCruz/function-gemma
0
1import os2from typing import Any, Literal3 4from pandas import DataFrame5 6try:7 from trackio.media.media import TrackioMedia8 from trackio.utils import MEDIA_DIR9except ImportError:10 from media.media import TrackioMedia11 from utils import MEDIA_DIR12 13 14class Table:15 """16 Initializes a Table object.17 18 Tables can be used to log tabular data including images, numbers, and text.19 20 Args:21 columns (`list[str]`, *optional*):22 Names of the columns in the table. Optional if `data` is provided. Not23 expected if `dataframe` is provided. Currently ignored.24 data (`list[list[Any]]`, *optional*):25 2D row-oriented array of values. Each value can be a number, a string26 (treated as Markdown and truncated if too long), or a `Trackio.Image` or27 list of `Trackio.Image` objects.28 dataframe (`pandas.DataFrame`, *optional*):29 DataFrame used to create the table. When set, `data` and `columns`30 arguments are ignored.31 rows (`list[list[Any]]`, *optional*):32 Currently ignored.33 optional (`bool` or `list[bool]`, *optional*, defaults to `True`):34 Currently ignored.35 allow_mixed_types (`bool`, *optional*, defaults to `False`):36 Currently ignored.37 log_mode: (`Literal["IMMUTABLE", "MUTABLE", "INCREMENTAL"]` or `None`, *optional*, defaults to `"IMMUTABLE"`):38 Currently ignored.39 """40 41 TYPE = "trackio.table"42 43 def __init__(44 self,45 columns: list[str] | None = None,46 data: list[list[Any]] | None = None,47 dataframe: DataFrame | None = None,48 rows: list[list[Any]] | None = None,49 optional: bool | list[bool] = True,50 allow_mixed_types: bool = False,51 log_mode: Literal["IMMUTABLE", "MUTABLE", "INCREMENTAL"] | None = "IMMUTABLE",52 ):53 # TODO: implement support for columns, dtype, optional, allow_mixed_types, and log_mode.54 # for now (like `rows`) they are included for API compat but don't do anything.55 if dataframe is None:56 self.data = DataFrame(data) if data is not None else DataFrame()57 else:58 self.data = dataframe59 60 def _has_media_objects(self, dataframe: DataFrame) -> bool:61 """Check if dataframe contains any TrackioMedia objects or lists of TrackioMedia objects."""62 for col in dataframe.columns:63 if dataframe[col].apply(lambda x: isinstance(x, TrackioMedia)).any():64 return True65 if (66 dataframe[col]67 .apply(68 lambda x: isinstance(x, list)69 and len(x) > 070 and isinstance(x[0], TrackioMedia)71 )72 .any()73 ):74 return True75 return False76 77 def _process_data(self, project: str, run: str, step: int = 0):78 """Convert dataframe to dict format, processing any TrackioMedia objects if present."""79 df = self.data80 if not self._has_media_objects(df):81 return df.to_dict(orient="records")82 83 processed_df = df.copy()84 for col in processed_df.columns:85 for idx in processed_df.index:86 value = processed_df.at[idx, col]87 if isinstance(value, TrackioMedia):88 value._save(project, run, step)89 processed_df.at[idx, col] = value._to_dict()90 if (91 isinstance(value, list)92 and len(value) > 093 and isinstance(value[0], TrackioMedia)94 ):95 [v._save(project, run, step) for v in value]96 processed_df.at[idx, col] = [v._to_dict() for v in value]97 98 return processed_df.to_dict(orient="records")99 100 @staticmethod101 def to_display_format(table_data: list[dict]) -> list[dict]:102 """103 Converts stored table data to display format for UI rendering.104 105 Note:106 This does not use the `self.data` attribute, but instead uses the107 `table_data` parameter, which is what the UI receives.108 109 Args:110 table_data (`list[dict]`):111 List of dictionaries representing table rows (from stored `_value`).112 113 Returns:114 `list[dict]`: Table data with images converted to markdown syntax and long115 text truncated.116 """117 truncate_length = int(os.getenv("TRACKIO_TABLE_TRUNCATE_LENGTH", "250"))118 119 def convert_image_to_markdown(image_data: dict) -> str:120 relative_path = image_data.get("file_path", "")121 caption = image_data.get("caption", "")122 absolute_path = MEDIA_DIR / relative_path123 return f'<img src="/gradio_api/file={absolute_path}" alt="{caption}" />'124 125 processed_data = []126 for row in table_data:127 processed_row = {}128 for key, value in row.items():129 if isinstance(value, dict) and value.get("_type") == "trackio.image":130 processed_row[key] = convert_image_to_markdown(value)131 elif (132 isinstance(value, list)133 and len(value) > 0134 and isinstance(value[0], dict)135 and value[0].get("_type") == "trackio.image"136 ):137 # This assumes that if the first item is an image, all items are images. Ok for now since we don't support mixed types in a single cell.138 processed_row[key] = (139 '<div style="display: flex; gap: 10px;">'140 + "".join([convert_image_to_markdown(item) for item in value])141 + "</div>"142 )143 elif isinstance(value, str) and len(value) > truncate_length:144 truncated = value[:truncate_length]145 full_text = value.replace("<", "<").replace(">", ">")146 processed_row[key] = (147 f'<details style="display: inline;">'148 f'<summary style="display: inline; cursor: pointer;">{truncated}…<span><em>(truncated, click to expand)</em></span></summary>'149 f'<div style="margin-top: 10px; padding: 10px; background: #f5f5f5; border-radius: 4px; max-height: 400px; overflow: auto;">'150 f'<pre style="white-space: pre-wrap; word-wrap: break-word; margin: 0;">{full_text}</pre>'151 f"</div>"152 f"</details>"153 )154 else:155 processed_row[key] = value156 processed_data.append(processed_row)157 return processed_data158 159 def _to_dict(self, project: str, run: str, step: int = 0):160 """161 Converts the table to a dictionary representation.162 163 Args:164 project (`str`):165 Project name for saving media files.166 run (`str`):167 Run name for saving media files.168 step (`int`, *optional*, defaults to `0`):169 Step number for saving media files.170 """171 data = self._process_data(project, run, step)172 return {173 "_type": self.TYPE,174 "_value": data,175 }176