How do I download a pickle file of my trained model using the st.download_button
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)?
Hi @snehankekre ,
I am using scikit learn.
@Jeffry_Christopher, here’s an example based on the scikit-learn docs:
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
This topic was automatically closed 365 days after the last reply. New replies are no longer allowed.