Snowpark ML Ops: Model Registry¶

Retrieving and Updating Comments¶

Use the model’s comment attribute to retrieve and update the model’s comment.

print(m.comment)
m.comment = "A better description than the one I provided originally"

print(m.description)
m.description = "A better description than the one I provided originally"

Retrieving and Updating Tags¶

Tags are metadata used to record a model’s purpose, algorithm, training data set, lifecycle stage, or other information you choose. You can set tags when the model is registered or at any time afterward. You can also update the values of existing tags or remove tags entirely.

To get all of a model’s tags as a Python dictionary, use show_tags.

print(m.show_tags())

To add a new tag or change the value of an existing tag, use set_tag.

m.set_tag("live_version", "v1")

To retrieve the value of a tag, use get_tag:

m.get_tag("live_version")

To remove a tag, use unset_tag.

model.unset_tag("live_version")

Deleting Model Versions¶

You can delete a model version using the model’s delete_version method (in package version 1.2.3 or later).

m.delete_version("rc1")

Default Version¶

A version of a model may be designated as the default model. Retrieve or set the model’s default attribute to obtain the current default version (as a ModelVersion object) or to change it (using a string).

default_version = m.default
m.default = "v2"

Getting a Reference to a Model Version¶

To get a reference to a specific version of a model as a ModelVersion instance, use the model’s version method. Use the model’s default attribute to get the default version of the model.

mv = m.version("v1")
mv = m.default

Once you have a reference to a specific version of a model (such as the variable mv in this example), you can retrieve or update its comments or metrics and call the model’s methods (or functions) as shown in the following sections.

Retrieving and Updating Comments¶

As with models, model versions can have comments, which can be accessed and set via the model version’s comment or description attribute.

print(mv.comment)
print(mv.description)

mv.comment = "A model version comment"
mv.description = "Same as setting the comment"

Retrieving and Updating Metrics¶

Metrics are key-value pairs used to track prediction accuracy and other model version characteristics. You can set metrics when creating a model version or set them using the set_metric method. A metric value can be any Python object that can be serialized to JSON, including numbers, strings, lists, and dictionaries. Unlike tags, metric names and possible values do not need to be defined in advance.

A test accuracy metric might be generated using sklearn’s accuracy_score:

from sklearn import metrics

test_accuracy = metrics.accuracy_score(test_labels, prediction)

The confusion matrix can be generated similarly using sklearn:

test_confusion_matrix = metrics.confusion_matrix(test_labels, prediction)

Then we can set these values as metrics as follows.

# scalar metric
mv.set_metric("test_accuracy", test_accuracy)

# hierarchical (dictionary) metric
mv.set_metric("evaluation_info", {"dataset_used": "my_dataset", "accuracy": test_accuracy, "f1_score": f1_score})

# multivalent (matrix) metric
mv.set_metric("confusion_matrix", test_confusion_matrix)

To retrieve a model version’s metrics as a Python dictionary, use show_metrics.

metrics = mv.show_metrics()

To delete a metric, call delete_metric.

mv.delete_metric("test_accuracy")

Calling Model Methods in SQL¶

You can call or invoke methods of a model in SQL using the model_name!method_name(...) syntax. The methods available on a model are determined by the underlying Python model class. For example, many types of models use a method named predict for inference.

To invoke a method of the latest version of a model, use the syntax shown here, passing arguments to the method, if any, between the parentheses, and passing the name of the table containing the inference data in the FROM clause.

SELECT <model_name>!<method_name>(...) FROM <table_name>;

To invoke a method of a specific version of a model, first create an alias to the specific version of the model, then invoke the desired method through the alias.

WITH <model_version_alias> AS MODEL <model_name> VERSION <version>
    SELECT <model_version_alias>!<method_name>(...) FROM <table_name>;

Snowpark ML¶

The registry supports models created using Snowpark ML modeling APIs (models derived from snowpark.ml.modeling.framework.base.BaseEstimator). The following additional options may be used in the options dictionary when calling log_model.

You do not need to specify sample_input_data or signatures when logging a Snowpark ML model; these are automatically inferred during fitting.

import pandas as pd
import numpy as np
from sklearn import datasets
from snowflake.ml.modeling.xgboost import XGBClassifier

iris = datasets.load_iris()
df = pd.DataFrame(data=np.c_[iris["data"], iris["target"]], columns=iris["feature_names"] + ["target"])
df.columns = [s.replace(" (CM)", "").replace(" ", "") for s in df.columns.str.upper()]

input_cols = ["SEPALLENGTH", "SEPALWIDTH", "PETALLENGTH", "PETALWIDTH"]
label_cols = "TARGET"
output_cols = "PREDICTED_TARGET"

clf_xgb = XGBClassifier(
        input_cols=input_cols, output_cols=output_cols, label_cols=label_cols, drop_input_cols=True
)
clf_xgb.fit(df)
model_ref = registry.log_model(
    clf_xgb,
    model_name="XGBClassifier",
    version_name="v1",
)
model_ref.run(df.drop(columns=label_cols).head(10), function_name='predict_proba')

scikit-learn¶

The registry supports models created using scikit-learn (models derived from sklearn.base.BaseEstimator or sklearn.pipeline.Pipeline). The following additional options may be used in the options dictionary when calling log_model.

You must specify either the sample_input_data or signatures parameter when logging a scikit-learn model so that the registry knows the signatures of the target methods.

from sklearn import datasets, ensemble

iris_X, iris_y = datasets.load_iris(return_X_y=True, as_frame=True)
model = ensemble.RandomForestClassifier(random_state=42)
model.fit(iris_X, iris_y)
model_ref = registry.log_model(
    model,
    model_name="RandomForestClassifier",
    version_name="v1",
    sample_input_data=iris_X,
    options={
        "method_options": {
            "predict": {"case_sensitive": True},
            "predict_proba": {"case_sensitive": True},
            "predict_log_proba": {"case_sensitive": True},
        }
    },
)
model_ref.run(iris_X[-10:], function_name='"predict_proba"')

XGBoost¶

The registry supports models created using XGBoost (models derived from xgboost.XGBModel or xgboost.Booster). The following additional options may be used in the options dictionary when calling log_model.

You must specify either the sample_input_data or signatures parameter when logging an XGBoost model so that the registry knows the signatures of the target methods.

import xgboost
from sklearn import datasets, model_selection

cal_X, cal_y = datasets.load_breast_cancer(as_frame=True, return_X_y=True)
cal_X_train, cal_X_test, cal_y_train, cal_y_test = model_selection.train_test_split(cal_X, cal_y)
params = dict(n_estimators=100, reg_lambda=1, gamma=0, max_depth=3, objective="binary:logistic")
regressor = xgboost.train(params, xgboost.DMatrix(data=cal_X_train, label=cal_y_train))
model_ref = registry.log_model(
    regressor,
    model_name="xgBooster",
    version_name="v1",
    sample_input_data=cal_X_test,
    options={
        "target_methods": ["predict"],
        "method_options": {
            "predict": {"case_sensitive": True},
        },
    },
)
model_ref.run(cal_X_test[-10:])

PyTorch¶

The registry supports PyTorch models (classes derived from torch.nn.Module or torch.jit.ModuleScript) if the model’s forward method accepts one or more torch.Tensor instances as input and returns a torch.Tensor or a tuple of them. The registry converts between pandas DataFrames and tensors when calling the model and returning the results. Tensors correspond to columns in the dataframe.

For example, suppose your model accepts two tensors like this:

import torch

class TorchModel(torch.nn.Module):
    def __init__(self, n_input: int, n_hidden: int, n_out: int, dtype: torch.dtype = torch.float32) -> None:
        super().__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Linear(n_input, n_hidden, dtype=dtype),
            torch.nn.ReLU(),
            torch.nn.Linear(n_hidden, n_out, dtype=dtype),
            torch.nn.Sigmoid(),
        )

    def forward(self, tensor_1: torch.Tensor, tensor_2: torch.Tensor) -> torch.Tensor:
        return self.model(tensor_1) + self.model(tensor_2)

If you want to pass torch.Tensor([[1,2],[3,4]]) as tensor_1 and torch.Tensor([[5,6], [7,8]]) as tensor_2, create a DataFrame as follows to pass to the model.

import pandas as pd
tensors = pd.DataFrame([[[1,2],[5,6]],[[3,4],[7,8]]])

Then the tensors DataFrame looks like this.

        0       1
0  [1, 2]  [5, 6]
1  [3, 4]  [7, 8]

Similarly, if your model returns two tensors, such as (torch.Tensor([[1,2],[3,4]]), torch.Tensor([[5,6], [7,8]])), the result is a DataFrame like the one above.

When providing sample input data for a PyTorch model, you must provide either a list of tensors (which will be converted ta a pandas DataFrame) or a DataFrame. A list may contain a single tensor, but a tensor on its own is not accepted.

import torch
import numpy as np

class TorchModel(torch.nn.Module):
        def __init__(self, n_input: int, n_hidden: int, n_out: int, dtype: torch.dtype = torch.float32) -> None:
                super().__init__()
                self.model = torch.nn.Sequential(
                        torch.nn.Linear(n_input, n_hidden, dtype=dtype),
                        torch.nn.ReLU(),
                        torch.nn.Linear(n_hidden, n_out, dtype=dtype),
                        torch.nn.Sigmoid(),
                )

        def forward(self, tensor: torch.Tensor) -> torch.Tensor:
                return self.model(tensor)

n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01
dtype = torch.float32
x = np.random.rand(batch_size, n_input)
data_x = torch.from_numpy(x).to(dtype=dtype)
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).to(dtype=dtype)

model = TorchModel(n_input, n_hidden, n_out, dtype=dtype)
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
for _epoch in range(100):
        pred_y = model.forward(data_x)
        loss = loss_function(pred_y, data_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


model_ref = registry.log_model(
    model,
    model_name="torchModel",
    version_name="v1",
    sample_input_data=[data_x],
)
model_ref.run([data_x])

TensorFlow¶

Models that extend tensorflow.Module or tensorflow.keras.Model are supported when they accept and return tensors and are compilable or compiled.

• The __call__ method for a tensorflow.Module or the call method for a tensorflow.keras.Model accepts one or more tensorflow.Tensor or tensorflow.Variable as input and returns a tensorflow.Tensor or tensorflow.Variable or a tuple of one of these types.

• If your model extends Module, it must be compilable, meaning the __call__ method is decorated with @tensorflow.function; see tf.function documentation. If it extends Model, it must be compiled; see compile documentation.

The registry converts between pandas DataFrames and tensors when calling the model and returning the results. Tensors correspond to columns in the dataframe.

For example, suppose your model accepts two tensors like this:

import tensorflow as tf

class KerasModel(tf.keras.Model):
    def  __init__(self, n_hidden: int, n_out: int) -> None:
        super().__init__()
        self.fc_1 = tf.keras.layers.Dense(n_hidden, activation="relu")
        self.fc_2 = tf.keras.layers.Dense(n_out, activation="sigmoid")

    def call(self, tensor_1: tf.Tensor, tensor_2: tf.Tensor) -> tf.Tensor:
        input = tensor_1 + tensor_2
        x = self.fc_1(input)
        x = self.fc_2(x)
        return x

If you want to pass tf.Tensor([[1,2],[3,4]]) as tensor_1 and tf.Tensor([[5,6], [7,8]]) as tensor_2, create a DataFrame as follows to pass to the model.

import pandas as pd
tensors = pd.DataFrame([[[1,2],[5,6]],[[3,4],[7,8]]])

Then the tensors DataFrame looks like this.

        0       1
0  [1, 2]  [5, 6]
1  [3, 4]  [7, 8]

Similarly, if your model returns two tensors, such as (tf.Tensor([[1,2],[3,4]]), tf.Tensor([[5,6], [7,8]])), the result is a DataFrame like the one above.

When providing sample input data for a TensorFlow model, you must provide either a list of tensors (which will be converted ta a pandas DataFrame) or a DataFrame. A list may contain a single tensor, but a tensor on its own is not accepted.

import tensorflow as tf
import numpy as np

class KerasModel(tf.keras.Model):
        def __init__(self, n_hidden: int, n_out: int) -> None:
                super().__init__()
                self.fc_1 = tf.keras.layers.Dense(n_hidden, activation="relu")
                self.fc_2 = tf.keras.layers.Dense(n_out, activation="sigmoid")

        def call(self, tensor: tf.Tensor) -> tf.Tensor:
                input = tensor
                x = self.fc_1(input)
                x = self.fc_2(x)
                return x

n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01
dtype = tf.float32
x = np.random.rand(batch_size, n_input)
data_x = tf.convert_to_tensor(x, dtype=dtype)
raw_data_y = tf.random.uniform((batch_size, 1))
raw_data_y = tf.where(raw_data_y > 0.5, tf.ones_like(raw_data_y), tf.zeros_like(raw_data_y))
data_y = tf.cast(raw_data_y, dtype=dtype)

model = KerasModel(n_hidden, n_out)
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate), loss=tf.keras.losses.MeanSquaredError())
model.fit(data_x, data_y, batch_size=batch_size, epochs=100)

model_ref = registry.log_model(
    model,
    model_name="tfModel",
    version_name="v1",
    sample_input_data=[data_x],
)
model_ref.run([data_x])

MLFlow¶

MLFlow models that provide a PyFunc flavor are supported. If your MLFlow model has a signature, the signature argument is inferred from the model. Otherwise, you must provide either signature or sample_input_data.

The following additional options may be used in the options dictionary when calling log_model.

import mlflow
from sklearn import datasets, model_selection, ensemble

db = datasets.load_diabetes(as_frame=True)
X_train, X_test, y_train, y_test = model_selection.train_test_split(db.data, db.target)
with mlflow.start_run() as run:
    rf = ensemble.RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
    rf.fit(X_train, y_train)

    # Use the model to make predictions on the test dataset.
    predictions = rf.predict(X_test)
    signature = mlflow.models.signature.infer_signature(X_test, predictions)
    mlflow.sklearn.log_model(
        rf,
        "model",
        signature=signature,
    )
    run_id = run.info.run_id


model_ref = registry.log_model(
    mlflow.pyfunc.load_model(f"runs:/{run_id}/model"),
    model_name="mlflowModel",
    version_name="v1",
    conda_dependencies=["mlflow<=2.4.0", "scikit-learn", "scipy"],
    options={"ignore_mlflow_dependencies": True}
)
model_ref.run(X_test)

Hugging Face Pipeline¶

The registry supports Hugging Face model classes defined as transformers that derive from transformers.Pipeline. The following code is an example of logging a compatible model.

lm_hf_model = transformers.pipeline(
    task="text-generation",
    model="bigscience/bloom-560m",
    token="...",  # Put your HuggingFace token here.
    return_full_text=False,
    max_new_tokens=100,
)

lmv = reg.log_model(lm_hf_model, model_name='bloom', version_name='v560m')

The registry infers the signatures argument as long as the pipeline contains only tasks from the following list.

• conversational

• fill-mask

• question-answering

• summarization

• table-question-answering

• text2text-generation

• text-classification (alias sentiment-analysis)

• text-generation

• token-classification (alias ner)

• translation

• translation_xx_to_yy

• zero-shot-classification

To see the inferred signature, use the show_functions method. The following, for example, is the result of lmv.show_functions() where lmv is the model logged above.

{'name': '__CALL__',
  'target_method': '__call__',
  'signature': ModelSignature(
                      inputs=[
                          FeatureSpec(dtype=DataType.STRING, name='inputs')
                      ],
                      outputs=[
                          FeatureSpec(dtype=DataType.STRING, name='outputs')
                      ]
                  )}]

With this information, you can call the model as follows.

import pandas as pd
remote_prediction = lmv.run(pd.DataFrame(["Hello, how are you?"], columns=["inputs"]))

# Prepare model
import transformers
import pandas as pd

finbert_model = transformers.pipeline(
    task="text-classification",
    model="ProsusAI/finbert",
    top_k=2,
)

# Log the model
mv = registry.log_model(
    finbert_model,
    model_name="finbert",
    version_name="v1",
)

# Use the model
mv.run(pd.DataFrame(
        [
            ["I have a problem with my Snowflake that needs to be resolved asap!!", ""],
            ["I would like to have udon for today's dinner.", ""],
        ]
    )
)

0  [{"label": "negative", "score": 0.8106237053871155}, {"label": "neutral", "score": 0.16587384045124054}]
1  [{"label": "neutral", "score": 0.9263970851898193}, {"label": "positive", "score": 0.05286872014403343}]