skops_secure_sklearn_model_save_load_quickstart.py

1import numpy as np
2from sklearn.linear_model import LogisticRegression
3from sklearn.datasets import load_iris
4import skops.io as sio
5
6# Load data and train a model
7X, y = load_iris(return_X_y=True)
8clf = LogisticRegression(random_state=0, max_iter=1000).fit(X, y)
9
10# Save the model using skops
11obj_save = "iris_model.skops"
12sio.dump(clf, obj_save)
13
14# Define trusted types for secure loading
15# In a real scenario, you should only trust the types you expect
16trusted_types = sio.get_untrusted_types(file=obj_save)
17
18# Load the model back
19loaded_clf = sio.load(obj_save, trusted=trusted_types)
20
21# Make a prediction to verify
22predictions = loaded_clf.predict(X[:5])
23print(f"Predictions: {predictions}")