Download pickle file of trained model using st.download_button

[Jeffry_Christopher]

How do I download a pickle file of my trained model using the st.download_button

[snehankekre]

Hi @Jeffry_Christopher, welcome to the community!!

Most ML libraries have an inbuilt method to save models. What library are you using to train your model (e.g. scikit-learn, xgboost, etc)?

[Jeffry_Christopher]

Hi @snehankekre ,

I am using scikit learn.

[snehankekre]

@Jeffry_Christopher, here’s an example based on the scikit-learn docs:

9. Model persistence

Solution

Say your model is clf. You can then pass pickle.dumps(clf) to the data argument of st.download_button, like so:

st.download_button(
    "Download Model",
    data=pickle.dumps(clf),
    file_name="model.pkl",
)

Example

Here’s a working example:

import streamlit as st
from sklearn import svm
from sklearn import datasets
import pickle

clf = svm.SVC()
X, y = datasets.load_iris(return_X_y=True)
clf.fit(X, y)


st.download_button(
    "Download Model",
    data=pickle.dumps(clf),
    file_name="model.pkl",
)

uploaded_file = st.file_uploader("Upload Model")

if uploaded_file is not None:
    clf2 = pickle.loads(uploaded_file.read())
    st.write("Model loaded")
    st.write(clf2)
    st.write("Predicting...")
    st.write(clf2.predict(X[0:1]))
    st.write(y[0])
    st.write("Done!")

Output

Happy Streamlit-ing!
Snehan