add publish function to wandb (#13)

yuriihavrylko · Feb 11, 2024 · ab633c0 · ab633c0
1 parent 3f001a6
commit ab633c0
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/app/src/helpers/wandb_registry.py b/app/src/helpers/wandb_registry.py
@@ -0,0 +1,9 @@
+import wandb
+
+def publish_model(model_path, project, name, model_type="model"):
+    run = wandb.init(project=project, job_type="model-publishing")
+    artifact = wandb.Artifact(name, type=model_type)
+    artifact.add_dir(model_path)
+    run.log_artifact(artifact)
+    run.finish()
+    print(f"Published {name} to W&B")