Snowflake 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:

m.unset_tag("live_version")

Renaming a model¶

Use the rename method to rename or move a model. Specify a fully qualified name as the new name to move the model to a different database or schema.

m.rename("MY_MODEL_TOO")

Deleting model versions¶

You can delete a model version by using the model’s delete_version method:

m.delete_version("rc1")

Default version¶

A version of a model can 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"

Model version aliases¶

You can assign an alias to a model version by using the SQL ALTER MODEL command. You can use an alias wherever a version name is required, such as when getting a reference to a model version, in Python or in SQL. A given alias can be assigned to only one model version at a time.

In addition to aliases you create, the following system aliases are available in all models:

• DEFAULT refers to the default version of the model.

• FIRST refers to the oldest version of the model by creation time.

• LAST refers to the newest version of the model by creation time.

Alias names you create must not be the same as any existing version name or alias in the model, including system aliases.

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:

m = reg.get_model("MyModel")

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

After 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 you can set these values as metrics:

# 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")

Retrieving model explanations¶

The model registry can explain a model’s results, telling you which input features contribute most to predictions, by calculating Shapley values (https://towardsdatascience.com/the-shapley-value-for-ml-models-f1100bff78d1). This preview feature is available by default in all model views created in Snowflake 8.31 and later through the underlying model’s explain method. You can call explain from SQL or via a model view’s run method in Python.

For details on this feature, see Model Explainability.

Exporting a model version¶

Use mv.export to export a model’s files to a local directory; the directory is created if it does not exist:

mv.export("~/mymodel/")

By default, the exported files include the code, the environment to load the model, and model weights. To also export the files needed to run the model in a warehouse, specify export_mode = ExportMode.FULL:

mv.export("~/mymodel/", export_mode=ExportMode.FULL)

Loading a model version¶

Use mv.load to load the original Python model object that was originally added to the registry. You can then use the model for inference just as though you had defined it in your Python code:

clf = mv.load()

To ensure proper functionality of a model loaded from the registry, the target Python environment (that is, the versions of the Python interpreter and of all libraries) should be identical to the environment from which the model was logged. Specify force=True in the load call to force the model to be loaded even if the environment is different.

conda_env = session.file.get("snow://model/<modelName>/versions/<versionName>/runtimes/python_runtime/env/conda.yml", ".")
open("~/conda.yml", "w").write(conda_env)

conda env create --name newenv --file=~/conda.yml
conda activate newenv

The optional options argument is a dictionary of options for loading the model. Currently, the argument supports only the use_gpu option.

The following example illustrates the use of the options argument:

clf = mv.load(options={"use_gpu": True})

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 can be used in the options dictionary when you call 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 can be used in the options dictionary when you call 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)
clf = ensemble.RandomForestClassifier(random_state=42)
clf.fit(iris_X, iris_y)
model_ref = registry.log_model(
    clf,
    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 can be used in the options dictionary when you call 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 like this 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 to 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 (https://www.tensorflow.org/guide/function). If it extends Model, it must be compiled; see compile documentation (https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile).

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 like this 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 to 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 can be used in the options dictionary when you call 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 (link removed) 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 only if the pipeline contains one task from the following list:

• conversational

• fill-mask

• question-answering

• summarization

• table-question-answering

• text2text-generation

• text-classification (also called sentiment-analysis)

• text-generation

• token-classification (also called ner)

• translation

• translation_xx_to_yy

• zero-shot-classification

The sample_input_data argument is completely ignored for Hugging Face models. Specify the signatures argument when logging a Hugging Face model that is not in the above list so that the registry knows the signatures of the target methods.

To see the inferred signature, use the show_functions method. The following dictionary, 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}]