> ## Documentation Index
> Fetch the complete documentation index at: https://wb-21fd5541-workspaces-notebooks.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# Get started (v2)

In this guide you will learn how to fetch the Zoo dataset, process it into tensors, train a simple neural network classifier, and publish the resulting model to the W\&B Registry.

In this guide you will learn how to download and use artifacts linked to the Registry. To do this, you will download both a pretrained classification model and the corresponding dataset tensors. You will then use these artifacts to perform inference and evaluate the model's performance.

## Prerequisites

Before you begin, ensure that you have a W\&B API key. Run the attached [Python training script](#run-training-script-optional) if you want to run the notebook example described on this page yourself. The script covers the first part of the workflow, including fetching the dataset, processing it into tensors, defining a neural network model, training the neural network model, before publishing the resulting model to the W\&B Registry.

### Sign up and create an API key

To authenticate your machine with W\&B, you need an API key.

To create an API key, select the **Personal API key** or **Service Account API key** tab for details.

<Tabs>
  <Tab title="Personal API key">
    To create a personal API key owned by your user ID:

    1. Log in to W\&B, then click your user profile icon **> User Settings**.
    2. Click **Create new API key**.
    3. Provide a descriptive name for your API key.
    4. Click **Create**.
    5. Copy the displayed API key immediately and store it securely.
  </Tab>

  <Tab title="Service account API key">
    To create an API key owned by a service account:

    1. In your team or organization settings, go to the **Service Accounts** tab.
    2. Find the service account in the list.
    3. Click the **action (<Icon icon="ellipsis" iconType="solid" />)** menu, then click **Create API key**.
    4. Provide a name for the API key, then click **Create**.
    5. Copy the displayed API key immediately and store it securely.
    6. Click **Done**.

    You can create multiple API keys for a single service account to support different environments or workflows.
  </Tab>
</Tabs>

<Warning>
  W\&B shows the full API key only once, when you create it. After you close the dialog, you cannot view the full API key again. Your settings display only the key ID (the first part of the key). If you lose the full API key, you must create a new one.
</Warning>

For secure storage options, see [Store API keys securely](/platform/app/settings-page/user-settings/#store-and-handle-api-keys-securely).

### Run training script (Optional)

Copy the following code into a file named `train.py` and save it locally on your machine:

```python expandable train.py theme={null}
# /// script
# requires-python = ">=3.10"
# dependencies = ["pandas", "scikit-learn", "torch", "ucimlrepo", "wandb"]
# ///

"""Publish Zoo dataset tensors and a trained model to W&B Registry.

This script is a Python conversion of ``zoo_wandb.ipynb`` through the
"Publish model to registry" section. Downloading artifacts for inference is
left for a later phase.
"""

from __future__ import annotations

import argparse
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence, TypeAlias

import pandas as pd
import torch
import wandb
from sklearn.model_selection import train_test_split
from torch import nn
from ucimlrepo import fetch_ucirepo


SCRIPT_PATH = Path(__file__).resolve()
SCRIPT_DIR = SCRIPT_PATH.parent
LOGGER = logging.getLogger(__name__)
ConfigValue: TypeAlias = bool | int | float | str

DEFAULT_ENTITY = "wandb"
DEFAULT_PROJECT = "Zoo_Demo"
DEFAULT_REGISTRY = "Zoo"

FULL_DATASET_COLLECTION = "dataset-tensors"
SPLIT_DATASET_COLLECTION = "dataset-tensors-split"
MODEL_COLLECTION = "Classifier_Models"

DATASET_FILENAME = "zoo_dataset.pt"
LABELS_FILENAME = "zoo_labels.pt"
X_TRAIN_FILENAME = "zoo_dataset_X_train.pt"
Y_TRAIN_FILENAME = "zoo_labels_y_train.pt"
X_TEST_FILENAME = "zoo_dataset_X_test.pt"
Y_TEST_FILENAME = "zoo_labels_y_test.pt"
MODEL_FILENAME = "zoo_wandb.pth"
SCRIPT_ARTIFACT_NAME = "zoo_wandb_script"

DATASET_ARTIFACT_NAME = "zoo_dataset"
SPLIT_DATASET_ARTIFACT_NAME = "split_zoo_dataset"
DATASET_ARTIFACT_FILE = "zoo_dataset"
LABELS_ARTIFACT_FILE = "zoo_labels"
X_TRAIN_ARTIFACT_FILE = "zoo_dataset_X_train"
Y_TRAIN_ARTIFACT_FILE = "zoo_labels_y_train"
X_TEST_ARTIFACT_FILE = "zoo_dataset_X_test"
Y_TEST_ARTIFACT_FILE = "zoo_labels_y_test"


@dataclass(frozen=True, slots=True)
class WandbUser:
    """W&B entity and project used for registry publishing runs."""

    entity: str
    project: str


@dataclass(frozen=True, slots=True)
class ArtifactFile:
    """A local file and its name inside a W&B artifact."""

    path: Path
    name: str


@dataclass(frozen=True, slots=True)
class WandbRegistryEntry:
    """An artifact version to link into a W&B Registry collection."""

    registry: str
    collection: str
    artifact_name: str
    artifact_type: str
    description: str
    job_type: str
    files: tuple[ArtifactFile, ...]

    @property
    def target_path(self) -> str:
        return f"wandb-registry-{self.registry}/{self.collection}"


class NeuralNetwork(nn.Module):
    """Simple neural network classifier from the Zoo registry notebook."""

    def __init__(self) -> None:
        super().__init__()
        self.linear_stack = nn.Sequential(
            nn.Linear(in_features=16, out_features=16),
            nn.Sigmoid(),
            nn.Linear(in_features=16, out_features=7),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_stack(x)


def fetch_data() -> tuple[pd.DataFrame, pd.DataFrame]:
    """Fetch the Zoo dataset from the UCI Machine Learning Repository."""

    zoo = fetch_ucirepo(id=111)
    features = zoo.data.features
    labels = zoo.data.targets

    LOGGER.info("features: %s type: %s", features.shape, type(features))
    LOGGER.info("labels: %s type: %s", labels.shape, type(labels))

    return features, labels


def process_data(
    features: pd.DataFrame,
    labels: pd.DataFrame,
    output_dir: Path,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Convert the Zoo dataset to tensors and save the processed files."""

    dataset = torch.as_tensor(features.to_numpy(copy=True), dtype=torch.float32)
    label_tensor = torch.as_tensor(labels.to_numpy(copy=True), dtype=torch.long) - 1

    LOGGER.info("dataset: %s dtype: %s", dataset.shape, dataset.dtype)
    LOGGER.info("labels: %s dtype: %s", label_tensor.shape, label_tensor.dtype)

    torch.save(dataset, output_dir / DATASET_FILENAME)
    torch.save(label_tensor, output_dir / LABELS_FILENAME)

    return dataset, label_tensor


def split_data(
    dataset: torch.Tensor,
    labels: torch.Tensor,
    output_dir: Path,
    *,
    random_state: int = 42,
    test_size: float = 0.25,
    shuffle: bool = True,
) -> dict[str, ConfigValue]:
    """Split the tensors into train/test files and return the split config.

    Args:
        dataset: The input feature tensor.
        labels: The input label tensor.
        output_dir: The directory to save the split files.
        random_state: The random seed for reproducibility.
        test_size: The proportion of the dataset to include in the test split.
        shuffle: Whether to shuffle the data before splitting.

    Returns:
        A dictionary containing the split configuration.
    """

    config: dict[str, ConfigValue] = {
        "random_state": random_state,
        "test_size": test_size,
        "shuffle": shuffle,
    }

    X_train, X_test, y_train, y_test = train_test_split(
        dataset,
        labels,
        random_state=random_state,
        test_size=test_size,
        shuffle=shuffle,
    )

    torch.save(X_train, output_dir / X_TRAIN_FILENAME)
    torch.save(y_train, output_dir / Y_TRAIN_FILENAME)
    torch.save(X_test, output_dir / X_TEST_FILENAME)
    torch.save(y_test, output_dir / Y_TEST_FILENAME)

    return config


def publish_dataset_registry(
    entry: WandbRegistryEntry,
    user: WandbUser,
    *,
    config: dict[str, ConfigValue] | None = None,
) -> None:
    """Publish a dataset artifact and link it to a W&B Registry collection.

    Args:
        entry: The W&B Registry entry describing the dataset artifact.
        user: The W&B user information.
        config: Optional configuration dictionary for the W&B run.
    """

    LOGGER.info(
        "Publishing artifact %r to registry collection %r",
        entry.artifact_name,
        entry.target_path,
    )

    with wandb.init(
        entity=user.entity,
        project=user.project,
        job_type=entry.job_type,
        config=config,
    ) as run:
        artifact = wandb.Artifact(
            name=entry.artifact_name,
            type=entry.artifact_type,
            description=entry.description,
        )

        for artifact_file in entry.files:
            artifact.add_file(
                local_path=str(artifact_file.path),
                name=artifact_file.name,
            )

        run.link_artifact(artifact=artifact, target_path=entry.target_path)


def build_registry_entries(output_dir: Path, registry: str) -> tuple[
    WandbRegistryEntry,
    WandbRegistryEntry,
]:
    """Define the dataset artifacts published by this phase of the notebook.

    Args:
        output_dir: The directory where the dataset files are stored.
        registry: The W&B registry to which the artifacts will be published.

    Returns:
        A tuple containing the full dataset entry and the split dataset entry.
    """
    full_dataset_entry = WandbRegistryEntry(
        registry=registry,
        collection=FULL_DATASET_COLLECTION,
        artifact_name=DATASET_ARTIFACT_NAME,
        artifact_type="dataset",
        description="Processed dataset and labels.",
        job_type="publish_dataset",
        files=(
            ArtifactFile(output_dir / DATASET_FILENAME, DATASET_ARTIFACT_FILE),
            ArtifactFile(output_dir / LABELS_FILENAME, LABELS_ARTIFACT_FILE),
        ),
    )

    split_dataset_entry = WandbRegistryEntry(
        registry=registry,
        collection=SPLIT_DATASET_COLLECTION,
        artifact_name=SPLIT_DATASET_ARTIFACT_NAME,
        artifact_type="dataset",
        description=(
            "Artifact contains `zoo_dataset` split into 4 datasets. "
            "For training, use `zoo_dataset_X_train` and `zoo_labels_y_train`. "
            "For testing, use `zoo_dataset_X_test` and `zoo_labels_y_test`."
        ),
        job_type="publish_split_dataset",
        files=(
            ArtifactFile(output_dir / X_TRAIN_FILENAME, X_TRAIN_ARTIFACT_FILE),
            ArtifactFile(output_dir / Y_TRAIN_FILENAME, Y_TRAIN_ARTIFACT_FILE),
            ArtifactFile(output_dir / X_TEST_FILENAME, X_TEST_ARTIFACT_FILE),
            ArtifactFile(output_dir / Y_TEST_FILENAME, Y_TEST_ARTIFACT_FILE),
        ),
    )

    return full_dataset_entry, split_dataset_entry


def build_model() -> NeuralNetwork:
    """Build the same neural network classifier used in the notebook."""

    model = NeuralNetwork()
    LOGGER.info("Model architecture:\n%s", model)
    return model


def build_hyperparameter_config(
    *,
    learning_rate: float,
    epochs: int,
) -> dict[str, ConfigValue]:
    """Define the hyperparameters logged with the model training run."""
    return {
        "learning_rate": learning_rate,
        "epochs": epochs,
        "model_type": "Multivariate_neural_network_classifier",
    }


def load_tensor(path: Path) -> torch.Tensor:
    """Load a tensor file from disk."""
    return torch.load(path, weights_only=True)


def train_model_from_registry(
    user: WandbUser,
    *,
    registry: str,
    split_collection: str,
    dataset_version: int,
    output_dir: Path,
    model_filename: str,
    hyperparameter_config: dict[str, ConfigValue],
) -> str:
    """Train a Zoo classifier using the split dataset artifact from Registry."""

    model = build_model()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=float(hyperparameter_config["learning_rate"]),
    )
    model_path = output_dir / model_filename

    with wandb.init(
        entity=user.entity,
        project=user.project,
        job_type="training",
        config=hyperparameter_config,
    ) as run:
        artifact_name = (
            f"wandb-registry-{registry.lower()}/{split_collection}:v{dataset_version}"
        )
        dataset_artifact = run.use_artifact(artifact_or_name=artifact_name)

        X_train_path = Path(
            dataset_artifact.download(path_prefix=X_TRAIN_ARTIFACT_FILE)
        )
        y_train_path = Path(
            dataset_artifact.download(path_prefix=Y_TRAIN_ARTIFACT_FILE)
        )

        X_train = load_tensor(X_train_path / X_TRAIN_ARTIFACT_FILE)
        y_train = load_tensor(y_train_path / Y_TRAIN_ARTIFACT_FILE)

        prev_best_loss = float("inf")
        model_artifact_name = f"zoo-{run.id}"

        for epoch in range(int(hyperparameter_config["epochs"]) + 1):
            pred = model(X_train)
            loss = loss_fn(pred, y_train.squeeze(1))

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            loss_value = loss.item()
            run.log(
                {
                    "train/epoch_ndx": epoch,
                    "train/train_loss": loss_value,
                }
            )

            if epoch % 100 == 0 and loss_value <= prev_best_loss:
                LOGGER.info("epoch: %s loss: %s", epoch, loss_value)
                torch.save(model.state_dict(), model_path)
                prev_best_loss = loss_value

        LOGGER.info("Saving model artifact %s", model_artifact_name)
        model_artifact = wandb.Artifact(
            name=model_artifact_name,
            type="model",
            metadata={
                "num_classes": 7,
                "model_type": hyperparameter_config["model_type"],
            },
        )
        model_artifact.add_file(str(model_path))
        logged_artifact = run.log_artifact(model_artifact)
        logged_artifact.wait()

    return model_artifact_name


def save_script_artifact(
    user: WandbUser,
    *,
    script_path: Path,
    artifact_name: str = SCRIPT_ARTIFACT_NAME,
) -> str:
    """Save this Python script as a standalone W&B code artifact."""
    with wandb.init(
        entity=user.entity,
        project=user.project,
        job_type="save_script",
    ) as run:
        script_artifact = wandb.Artifact(
            name=artifact_name,
            type="code",
            description="Python script used for the Zoo registry workflow.",
            metadata={
                "filename": script_path.name,
            },
        )
        script_artifact.add_file(str(script_path), name=script_path.name)
        logged_artifact = run.log_artifact(script_artifact)
        logged_artifact.wait()

    return artifact_name


def publish_model_registry(
    user: WandbUser,
    *,
    registry: str,
    collection: str,
    model_artifact_name: str,
    version: int = 0,
) -> None:
    """Link the trained model artifact into a W&B Registry collection."""
    artifact_name = f"{user.entity}/{user.project}/{model_artifact_name}:v{version}"
    target_path = f"wandb-registry-{registry}/{collection}"

    LOGGER.info("Artifact name: %s", artifact_name)
    LOGGER.info("Target path: %s", target_path)

    with wandb.init(entity=user.entity, project=user.project) as run:
        model_artifact = run.use_artifact(
            artifact_or_name=artifact_name,
            type="model",
        )
        run.link_artifact(artifact=model_artifact, target_path=target_path)


def positive_int(value: str) -> int:
    """Parse a positive integer CLI argument."""
    parsed = int(value)
    if parsed < 1:
        raise argparse.ArgumentTypeError("must be 1 or greater")
    return parsed


def non_negative_int(value: str) -> int:
    """Parse a non-negative integer CLI argument."""
    parsed = int(value)
    if parsed < 0:
        raise argparse.ArgumentTypeError("must be 0 or greater")
    return parsed


def positive_float(value: str) -> float:
    """Parse a positive float CLI argument."""
    parsed = float(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be greater than 0")
    return parsed


def output_dir_path(value: str) -> Path:
    """Parse an output directory CLI argument."""
    path = Path(value).expanduser()
    if path.exists() and not path.is_dir():
        raise argparse.ArgumentTypeError(
            f"must be a directory, got file: {path}"
        )
    return path


def resolve_output_dir(path: Path) -> Path:
    """Resolve and create the output directory for generated files."""
    output_dir = path.expanduser().resolve()
    if output_dir.exists() and not output_dir.is_dir():
        raise NotADirectoryError(
            f"--output-dir must be a directory, got file: {output_dir}"
        )
    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Publish Zoo dataset tensors and a trained model to W&B Registry.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--entity",
        default=DEFAULT_ENTITY,
        help="W&B entity that owns the publishing project.",
    )
    parser.add_argument(
        "--project",
        default=DEFAULT_PROJECT,
        help="W&B project used for the publishing runs.",
    )
    parser.add_argument(
        "--registry",
        default=DEFAULT_REGISTRY,
        help="W&B Registry name to link the dataset artifacts into.",
    )
    parser.add_argument(
        "--output-dir",
        type=output_dir_path,
        default=SCRIPT_DIR,
        help="Directory where the tensor files are written before publishing.",
    )
    parser.add_argument(
        "--skip-publish",
        action="store_true",
        help="Create local tensor files without publishing to W&B or training.",
    )
    parser.add_argument(
        "--skip-dataset-publish",
        action="store_true",
        help=(
            "Do not publish dataset artifacts before training. Use this when "
            "the split dataset artifact is already available in the registry."
        ),
    )
    parser.add_argument(
        "--dataset-version",
        type=non_negative_int,
        default=0,
        help="Version of the split dataset registry artifact to train on.",
    )
    parser.add_argument(
        "--model-collection",
        default=MODEL_COLLECTION,
        help="Registry collection to link the trained model artifact into.",
    )
    parser.add_argument(
        "--learning-rate",
        type=positive_float,
        default=0.1,
        help="SGD learning rate for model training.",
    )
    parser.add_argument(
        "--epochs",
        type=positive_int,
        default=1000,
        help="Number of training epochs.",
    )
    parser.add_argument(
        "--model-filename",
        default=MODEL_FILENAME,
        help="Filename used when saving the trained PyTorch state dict.",
    )
    return parser.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> None:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    args = parse_args(argv)
    output_dir = resolve_output_dir(args.output_dir)

    user = WandbUser(entity=args.entity, project=args.project)
    full_dataset_entry, split_dataset_entry = build_registry_entries(
        output_dir=output_dir,
        registry=args.registry,
    )

    features, labels = fetch_data()
    dataset, label_tensor = process_data(features, labels, output_dir)
    split_config = split_data(dataset, label_tensor, output_dir)

    if args.skip_publish:
        LOGGER.info("Created dataset tensors in %s", output_dir)
        return

    if not args.skip_dataset_publish:
        publish_dataset_registry(full_dataset_entry, user)
        publish_dataset_registry(split_dataset_entry, user, config=split_config)

    hyperparameter_config = build_hyperparameter_config(
        learning_rate=args.learning_rate,
        epochs=args.epochs,
    )
    model_artifact_name = train_model_from_registry(
        user,
        registry=args.registry,
        split_collection=split_dataset_entry.collection,
        dataset_version=args.dataset_version,
        output_dir=output_dir,
        model_filename=args.model_filename,
        hyperparameter_config=hyperparameter_config,
    )
    save_script_artifact(user, script_path=SCRIPT_PATH)
    publish_model_registry(
        user,
        registry=args.registry,
        collection=args.model_collection,
        model_artifact_name=model_artifact_name,
    )


if __name__ == "__main__":
    main()
```

Next, run the training script using `uv`:

```bash theme={null}
uv train.py
```

## Create your first notebook

1. Navigate to your project's workspace.
2. Select **Notebooks** from the project sidebar.
3. Click **Create notebook**.

See the [Create and manage notebooks](/models/notebooks/create-notebook) for more information.

## Install dependencies

Within your notebook install the required dependencies:

1. Select **Manage packages** (<Icon icon="cube" />) from the notebook sidebar.
2. Enter `torch`, `ucimlrepo`, and `scikit-learn`.
3. Select **Add**.

For more information, see the [Manage packages and environments](/models/notebooks/manage-packages-environments) guide or refer to the [marimo documentation](https://docs.marimo.io/guides/package_management/).

## Run the notebook example