From 0bcbe2937fd3153b18a421a8363cb12d0cffb93e Mon Sep 17 00:00:00 2001
From: Boris Feld <boris@comet.com>
Date: Thu, 13 Jun 2024 10:56:14 +0200
Subject: [PATCH] Update for latest Comet SDK version

---
 comet_for_mlflow/cli.py                |  6 +++---
 comet_for_mlflow/comet_for_mlflow.py   |  6 +++---
 examples/keras-example/run.py          | 21 ++++++++-------------
 examples/model-registry-example/run.py | 23 +++++++++--------------
 setup.py                               |  2 +-
 5 files changed, 24 insertions(+), 34 deletions(-)

diff --git a/comet_for_mlflow/cli.py b/comet_for_mlflow/cli.py
index 75c7b32..9dac54c 100644
--- a/comet_for_mlflow/cli.py
+++ b/comet_for_mlflow/cli.py
@@ -65,10 +65,10 @@ def main():
         help="set the directory to store prepared runs; only relevant with --no-upload",
     )
     parser.add_argument(
-        "--force-reupload",
+        "--force-upload",
         action="store_true",
         default=False,
-        help="Force reupload of prepared experiments that were previously uploaded",
+        help="Force the upload of prepared experiments even if they were previously uploaded",
     )
     command_group = parser.add_mutually_exclusive_group()
     command_group.add_argument(
@@ -98,7 +98,7 @@ def main():
         args.upload,
         args.api_key,
         args.output_dir,
-        args.force_reupload,
+        args.force_upload,
         args.mlflow_store_uri,
         args.answer,
         args.email,
diff --git a/comet_for_mlflow/comet_for_mlflow.py b/comet_for_mlflow/comet_for_mlflow.py
index 52c656b..0d903e3 100644
--- a/comet_for_mlflow/comet_for_mlflow.py
+++ b/comet_for_mlflow/comet_for_mlflow.py
@@ -97,7 +97,7 @@ def __init__(
         upload_experiment,
         api_key,
         output_dir,
-        force_reupload,
+        force_upload,
         mlflow_store_uri,
         answer,
         email,
@@ -146,7 +146,7 @@ def __init__(
 
         self.upload_experiment = upload_experiment
         self.output_dir = output_dir
-        self.force_reupload = force_reupload
+        self.force_upload = force_upload
         self.mlflow_store_uri = mlflow_store_uri
 
     def prepare(self):
@@ -486,7 +486,7 @@ def upload(self, prepared_data):
                     upload_single_offline_experiment(
                         archive_path,
                         self.api_key,
-                        force_reupload=self.force_reupload,
+                        force_upload=self.force_upload,
                         display_level="debug",
                     )
 
diff --git a/examples/keras-example/run.py b/examples/keras-example/run.py
index 4fc0b73..23d6ae1 100644
--- a/examples/keras-example/run.py
+++ b/examples/keras-example/run.py
@@ -3,17 +3,15 @@
 """Trains and evaluate a simple MLP
 on the Reuters newswire topic classification task.
 """
-from __future__ import print_function
-
-import keras
-import mlflow.keras
+import mlflow
 import numpy as np
-from keras.datasets import reuters
-from keras.layers import Activation, Dense, Dropout
-from keras.models import Sequential
-from keras.preprocessing.text import Tokenizer
+from tensorflow import keras
+from tensorflow.keras.datasets import reuters
+from tensorflow.keras.layers import Activation, Dense, Dropout
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.preprocessing.text import Tokenizer
 
-mlflow.keras.autolog()
+mlflow.tensorflow.autolog()
 
 max_words = 1000
 batch_size = 32
@@ -38,8 +36,7 @@
 print("x_test shape:", x_test.shape)
 
 print(
-    "Convert class vector to binary class matrix "
-    "(for use with categorical_crossentropy)"
+    "Convert class vector to binary class matrix (for use with categorical_crossentropy)"
 )
 y_train = keras.utils.to_categorical(y_train, num_classes)
 y_test = keras.utils.to_categorical(y_test, num_classes)
@@ -67,5 +64,3 @@
 score = model.evaluate(x_test, y_test, batch_size=batch_size, verbose=1)
 print("Test score:", score[0])
 print("Test accuracy:", score[1])
-
-mlflow.keras.log_model(model, "models")
diff --git a/examples/model-registry-example/run.py b/examples/model-registry-example/run.py
index 2301dea..8eb3b54 100644
--- a/examples/model-registry-example/run.py
+++ b/examples/model-registry-example/run.py
@@ -1,24 +1,20 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+
 """Trains and evaluate a simple MLP
 on the Reuters newswire topic classification task.
 """
-from __future__ import print_function
-
-import keras
-import mlflow.keras
+import mlflow
 import numpy as np
-from keras.datasets import reuters
-from keras.layers import Activation, Dense, Dropout
-from keras.models import Sequential
-from keras.preprocessing.text import Tokenizer
+from tensorflow import keras
+from tensorflow.keras.datasets import reuters
+from tensorflow.keras.layers import Activation, Dense, Dropout
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.preprocessing.text import Tokenizer
 
 mlflow.set_tracking_uri("sqlite:///db.sqlite")
 
-
-mlflow.keras.autolog()
-
-mlflow.log_param("Test", "FOO")
+mlflow.tensorflow.autolog()
 
 max_words = 1000
 batch_size = 32
@@ -43,8 +39,7 @@
 print("x_test shape:", x_test.shape)
 
 print(
-    "Convert class vector to binary class matrix "
-    "(for use with categorical_crossentropy)"
+    "Convert class vector to binary class matrix (for use with categorical_crossentropy)"
 )
 y_train = keras.utils.to_categorical(y_train, num_classes)
 y_test = keras.utils.to_categorical(y_test, num_classes)
diff --git a/setup.py b/setup.py
index e7df3b8..16892b7 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
 with open("HISTORY.md") as history_file:
     history = history_file.read()
 
-requirements = ["mlflow", "comet_ml>=3.1.1", "tabulate", "tqdm", "typing"]
+requirements = ["mlflow", "comet_ml>=3.39.0", "tabulate", "tqdm", "typing"]
 
 
 setup(