From f29680ecfe63dfb2bb307a425aa93f8eeca524cf Mon Sep 17 00:00:00 2001
From: Boris Feld <boris@comet.com>
Date: Tue, 10 Sep 2024 16:27:00 +0200
Subject: [PATCH] Start updating Optuna example (#182)

* Start updating Optuna example

* Add Optuna to test matrix

* Remove duplicate
---
 .../optuna-hello-world/optuna-hello-world.py  | 22 ++++++++-----------
 .../optuna-hello-world/requirements.txt       |  3 ++-
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/integrations/model-optimization/optuna/optuna-hello-world/optuna-hello-world.py b/integrations/model-optimization/optuna/optuna-hello-world/optuna-hello-world.py
index 8e9d9e7d..226d340d 100644
--- a/integrations/model-optimization/optuna/optuna-hello-world/optuna-hello-world.py
+++ b/integrations/model-optimization/optuna/optuna-hello-world/optuna-hello-world.py
@@ -1,31 +1,27 @@
 # coding: utf-8
-from comet_ml import Experiment, login
+from comet_ml import login
 
 import optuna
+from optuna_integration.comet import CometCallback
 
 # Login to Comet if needed
 login()
 
+study = optuna.create_study()
+comet = CometCallback(
+    study, project_name="comet-example-optuna-hello-world", metric_names=["score"]
+)
+
 
+@comet.track_in_comet()
 def objective(trial):
     x = trial.suggest_float("x", -10, 10)
     objective = (x - 2) ** 2
 
-    experiment = Experiment(project_name="comet-example-optuna-hello-world")
-
-    experiment.log_optimization(
-        optimization_id=trial.study.study_name,
-        metric_name="objective",
-        metric_value=objective,
-        parameters={"x": x},
-        objective="minimize",
-    )
-
     return objective
 
 
-study = optuna.create_study()
-study.optimize(objective, n_trials=20)
+study.optimize(objective, n_trials=20, callbacks=[comet])
 
 best_params = study.best_params
 found_x = best_params["x"]
diff --git a/integrations/model-optimization/optuna/optuna-hello-world/requirements.txt b/integrations/model-optimization/optuna/optuna-hello-world/requirements.txt
index 2612fcbe..296a0bec 100644
--- a/integrations/model-optimization/optuna/optuna-hello-world/requirements.txt
+++ b/integrations/model-optimization/optuna/optuna-hello-world/requirements.txt
@@ -1,2 +1,3 @@
 comet_ml>=3.33.10
-optuna
+optuna>=4.0.0
+optuna-integration>=4.0.0