Fetch a table from wandb based on the artifact ID. Args: artifact_id (str): You can find this on the Artifacts tab in the column "Name". For example "run-fsa64p5b-evalpredictionsiter1000:v0". project_name (str): Name of the wandb project. Defaults to "zg-neox".
(
artifact_id: str,
project_name: str = "jrt",
entity: str = "hazy-research"
)
| 213 | wandb.log_artifact(artifact) |
| 214 | |
| 215 | def load_table( |
| 216 | artifact_id: str, |
| 217 | project_name: str = "jrt", |
| 218 | entity: str = "hazy-research" |
| 219 | ) -> pd.DataFrame: |
| 220 | """ |
| 221 | Fetch a table from wandb based on the artifact ID. |
| 222 | Args: |
| 223 | artifact_id (str): You can find this on the Artifacts tab in the column "Name". |
| 224 | For example "run-fsa64p5b-evalpredictionsiter1000:v0". |
| 225 | project_name (str): Name of the wandb project. Defaults to "zg-neox". |
| 226 | entity (str): The entity (usually a user or team) in wandb. Defaults to "hazy-research". |
| 227 | |
| 228 | Returns: |
| 229 | pd.DataFrame: Fetched table in DataFrame format. |
| 230 | """ |
| 231 | import wandb |
| 232 | import tempfile |
| 233 | api = wandb.Api() |
| 234 | artifact = api.artifact(f"{entity}/{project_name}/{artifact_id}") |
| 235 | |
| 236 | # download to temporary directory |
| 237 | with tempfile.TemporaryDirectory() as tmp_dir: |
| 238 | artifact_dir = artifact.download(root=tmp_dir) |
| 239 | |
| 240 | # find the table files |
| 241 | tables = [file for file in os.listdir(artifact_dir) if file.endswith(".table.json") or file.endswith(".feather")] |
| 242 | |
| 243 | # validate number of tables found |
| 244 | assert len(tables) == 1, f"Expected 1 table, found {len(tables)} tables." |
| 245 | table_path = os.path.join(artifact_dir, tables[0]) |
| 246 | if table_path.endswith(".feather"): |
| 247 | df = pd.read_feather(table_path) |
| 248 | else: |
| 249 | with open(table_path, 'r') as file: |
| 250 | json_dict = json.load(file) |
| 251 | df = pd.DataFrame(json_dict["data"], columns=json_dict["columns"]) |
| 252 | return df |
| 253 | |
| 254 | def load_config(run_id: int): |
| 255 | """ |