Update Comet experiment tracker docs

docs/book/component-guide/experiment-trackers/comet.md

@@ -6,6 +6,8 @@ description: Logging and visualizing experiments with Comet.
 
 The Comet Experiment Tracker is an [Experiment Tracker](./experiment-trackers.md) flavor provided with the Comet ZenML integration that uses [the Comet experiment tracking platform](https://www.comet.com/site/products/ml-experiment-tracking/) to log and visualize information from your pipeline steps (e.g., models, parameters, metrics).
 
+<figure><img src="../../.gitbook/assets/comet_pipeline.png" alt=""><figcaption><p>A pipeline with a Comet experiment tracker url as metadata</p></figcaption></figure>
+
 ### When would you want to use it?
 
 [Comet](https://www.comet.com/site/products/ml-experiment-tracking/) is a popular platform that you would normally use in the iterative ML experimentation phase to track and visualize experiment results. That doesn't mean that it cannot be repurposed to track and visualize the results produced by your automated pipeline runs, as you make the transition towards a more production-oriented workflow.
@@ -37,22 +39,6 @@ You need to configure the following credentials for authentication to the Comet
 * `workspace`: Optional. The name of the workspace where your project is located. If not specified, the default workspace associated with your API key will be used.
 
 {% tabs %}
-{% tab title="Basic Authentication" %}
-This option configures the credentials for the Comet platform directly as stack component attributes.
-
-{% hint style="warning" %}
-This is not recommended for production settings as the credentials won't be stored securely and will be clearly visible in the stack configuration.
-{% endhint %}
-
-```bash
-# Register the Comet experiment tracker
-zenml experiment-tracker register comet_experiment_tracker --flavor=comet \
-    --workspace=<workspace> --project_name=<project_name> --api_key=<key>
-
-# Register and set a stack with the new experiment tracker
-zenml stack register custom_stack -e comet_experiment_tracker ... --set
-```
-{% endtab %}
 
 {% tab title="ZenML Secret (Recommended)" %}
 This method requires you to [configure a ZenML secret](../../getting-started/deploying-zenml/secret-management.md) to store the Comet tracking service credentials securely.
@@ -76,15 +62,38 @@ zenml experiment-tracker register comet_tracker \
     --project_name={{comet_secret.project_name}} \
     --api_key={{comet_secret.api_key}}
     ...
+
+# Register and set a stack with the new experiment tracker
+zenml stack register custom_stack -e comet_experiment_tracker ... --set
 ```
 
 {% hint style="info" %}
 Read more about [ZenML Secrets](../../getting-started/deploying-zenml/secret-management.md) in the ZenML documentation.
 {% endhint %}
 {% endtab %}
+
+{% tab title="Basic Authentication" %}
+This option configures the credentials for the Comet platform directly as stack component attributes.
+
+{% hint style="warning" %}
+This is not recommended for production settings as the credentials won't be stored securely and will be clearly visible in the stack configuration.
+{% endhint %}
+
+```bash
+# Register the Comet experiment tracker
+zenml experiment-tracker register comet_experiment_tracker --flavor=comet \
+    --workspace=<workspace> --project_name=<project_name> --api_key=<key>
+
+# Register and set a stack with the new experiment tracker
+zenml stack register custom_stack -e comet_experiment_tracker ... --set
+```
+{% endtab %}
+
 {% endtabs %}
 
-For more up-to-date information on the Comet Experiment Tracker implementation and its configuration, you can have a look at [the SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-comet/#zenml.integrations.comet.experiment\_trackers.comet\_experiment\_tracker).
+<figure><img src="../../.gitbook/assets/comet_stack.png" alt=""><figcaption><p>A stack with the Comet experiment tracker</p></figcaption></figure>
+
+For more up-to-date information on the Comet Experiment Tracker implementation and its configuration, you can have a look at [the SDK docs for our Comet integration](https://sdkdocs.zenml.io/0.66.0/integration_code_docs/integrations-comet/#zenml.integrations.comet.flavors.comet_experiment_tracker_flavor.CometExperimentTrackerConfig).
 
 ### How do you use it?
 
@@ -98,8 +107,20 @@ experiment_tracker = Client().active_stack.experiment_tracker
 @step(experiment_tracker=experiment_tracker.name)
 def my_step():
     ...
+    # go through some experiment tracker methods
     experiment_tracker.log_metrics({"my_metric": 42})
     experiment_tracker.log_params({"my_param": "hello"})
+
+    # or use the Experiment object directly
+    experiment_tracker.experiment.log_model(...)
+
+    # or pass the Comet Experiment object into helper methods
+    from comet_ml.integration.sklearn import log_model
+    log_model(
+        experiment=experiment_tracker.experiment,
+        model_name="SVC",
+        model=model,
+    )
     ...
 ```
 
@@ -113,6 +134,10 @@ Comet comes with a web-based UI that you can use to find further details about y
 
 Every ZenML step that uses Comet should create a separate experiment which you can inspect in the Comet UI.
 
+<figure><img src="../../.gitbook/assets/comet_experiment_confusion_matrix.png" alt=""><figcaption><p>A confusion matrix logged in the Comet UI</p></figcaption></figure>
+
+<figure><img src="../../.gitbook/assets/comet_experiment_model.png" alt=""><figcaption><p>A model tracked in the Comet UI</p></figcaption></figure>
+
 You can find the URL of the Comet experiment linked to a specific ZenML run via the metadata of the step in which the experiment tracker was used:
 
 ```python
@@ -124,21 +149,131 @@ tracking_url = trainer_step.run_metadata["experiment_tracker_url"].value
 print(tracking_url)
 ```
 
+<figure><img src="../../.gitbook/assets/comet_pipeline.png" alt=""><figcaption><p>A pipeline with a Comet experiment tracker url as metadata</p></figcaption></figure>
+
 Alternatively, you can see an overview of all experiments at `https://www.comet.com/{WORKSPACE_NAME}/{PROJECT_NAME}/experiments/`.
 
 {% hint style="info" %}
 The naming convention of each Comet experiment is `{pipeline_run_name}_{step_name}` (e.g., `comet_example_pipeline-25_Apr_22-20_06_33_535737_my_step`), and each experiment will be tagged with both `pipeline_name` and `pipeline_run_name`, which you can use to group and filter experiments.
 {% endhint %}
 
+## Full Code Example
+
+This section combines all the code from this section into one simple script that you can use to run easily:
+
+<details>
+
+<summary>Code Example of this Section</summary>
+
+```python
+from comet_ml.integration.sklearn import log_model
+
+import numpy as np
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.svm import SVC
+from sklearn.metrics import accuracy_score
+from typing import Tuple
+
+from zenml import pipeline, step
+from zenml.client import Client
+from zenml.integrations.comet.flavors.comet_experiment_tracker_flavor import (
+    CometExperimentTrackerSettings,
+)
+from zenml.integrations.comet.experiment_trackers import CometExperimentTracker
+
+# Get the experiment tracker from the active stack
+experiment_tracker: CometExperimentTracker = Client().active_stack.experiment_tracker
+
+
+@step
+def load_data() -> Tuple[np.ndarray, np.ndarray]:
+    iris = load_iris()
+    X = iris.data
+    y = iris.target
+    return X, y
+
+
+@step
+def preprocess_data(
+    X: np.ndarray, y: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=0.2, random_state=42
+    )
+    scaler = StandardScaler()
+    X_train_scaled = scaler.fit_transform(X_train)
+    X_test_scaled = scaler.transform(X_test)
+    return X_train_scaled, X_test_scaled, y_train, y_test
+
+
+@step(experiment_tracker=experiment_tracker.name)
+def train_model(X_train: np.ndarray, y_train: np.ndarray) -> SVC:
+    model = SVC(kernel="rbf", C=1.0)
+    model.fit(X_train, y_train)
+    log_model(
+        experiment=experiment_tracker.experiment,
+        model_name="SVC",
+        model=model,
+    )
+    return model
+
+
+@step(experiment_tracker=experiment_tracker.name)
+def evaluate_model(model: SVC, X_test: np.ndarray, y_test: np.ndarray) -> float:
+    y_pred = model.predict(X_test)
+    accuracy = accuracy_score(y_test, y_pred)
+
+    # Log metrics using Comet
+    experiment_tracker.log_metrics({"accuracy": accuracy})
+    experiment_tracker.experiment.log_confusion_matrix(y_test, y_pred)
+    return accuracy
+
+
+@pipeline(enable_cache=False)
+def iris_classification_pipeline():
+    X, y = load_data()
+    X_train, X_test, y_train, y_test = preprocess_data(X, y)
+    model = train_model(X_train, y_train)
+    accuracy = evaluate_model(model, X_test, y_test)
+
+
+if __name__ == "__main__":
+    # Configure Comet settings
+    comet_settings = CometExperimentTrackerSettings(tags=["iris_classification", "svm"])
+
+    # Run the pipeline
+    last_run = iris_classification_pipeline.with_options(
+        settings={"experiment_tracker": comet_settings}
+    )()
+
+    # Get the URLs for the trainer and evaluator steps
+    trainer_step, evaluator_step = (
+        last_run.steps["train_model"],
+        last_run.steps["evaluate_model"],
+    )
+    trainer_url = trainer_step.run_metadata["experiment_tracker_url"].value
+    evaluator_url = evaluator_step.run_metadata["experiment_tracker_url"].value
+    print(f"URL for trainer step: {trainer_url}")
+    print(f"URL for evaluator step: {evaluator_url}")
+```
+
+</details>
+
 #### Additional configuration
 
 For additional configuration of the Comet experiment tracker, you can pass `CometExperimentTrackerSettings` to provide additional tags for your experiments:
 
-```
-from zenml.integrations.comet.flavors.comet_experiment_tracker_flavor import CometExperimentTrackerSettings
+```python
+from zenml.integrations.comet.flavors.comet_experiment_tracker_flavor import (
+    CometExperimentTrackerSettings,
+)
 
 comet_settings = CometExperimentTrackerSettings(
-    tags=["some_tag"]
+    tags=["some_tag"],
+    run_name="",
+    settings={},
 )
 
 @step(
@@ -151,6 +286,6 @@ def my_step():
     ...
 ```
 
-Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-comet/#zenml.integrations.comet.flavors.comet\_experiment\_tracker\_flavor.CometExperimentTrackerSettings) for a full list of available attributes and [this docs page](../../how-to/use-configuration-files/runtime-configuration.md) for more information on how to specify settings.
+Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration_code_docs/integrations-comet/#zenml.integrations.comet.flavors.comet_experiment_tracker_flavor.CometExperimentTrackerSettings) for a full list of available attributes and [this docs page](../../how-to/use-configuration-files/runtime-configuration.md) for more information on how to specify settings.
 
-<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>
+<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>