sklearn_classifier_save_load_with_skops_secure_format.py

1import numpy as np
2from sklearn.datasets import load_iris
3from sklearn.linear_model import LogisticRegression
4from sklearn.model_selection import train_test_split
5import skops.io as sio
6
7# Load data and train a model
8X, y = load_iris(return_X_y=True)
9X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
10model = LogisticRegression(random_state=42, max_iter=1000).fit(X_train, y_train)
11
12# Save the model using skops
13obj_save = "iris_model.skops"
14sio.dump(model, obj_save)
15
16# To load the model, we first need to define which types are trusted. 
17# We can use get_untrusted_types to see what's in the file.
18unknown_types = sio.get_untrusted_types(file=obj_save)
19print(f"Unknown types: {unknown_types}")
20
21# Now we can load the model by trusting the required types
22# In this case, we trust the types required by sklearn
23loaded_model = sio.load(obj_save, trusted=unknown_types)
24
25# Use the loaded model for prediction
26predictions = loaded_model.predict(X_test)
27print(f"Predictions: {predictions[:5]}")