How to download a trained model

I’m wondering how I can download a model that was trained within the app. I know that you can download a csv file using the following code:

    csv = output_df.to_csv(index=False)
    b64 = base64.b64encode(csv.encode()).decode()  # some strings <-> bytes conversions necessary here
    href = f'<a href="data:file/csv;base64,{b64}">Download Test Set Predictions CSV File</a> (right-click and save as &lt;some_name&gt;.csv)'
    st.markdown(href, unsafe_allow_html=True)

How could I apply something similar using pickle?

Figured it out. Here’s the code if anyone else is wondering:

def download_model(model):
    output_model = pickle.dumps(model)
    b64 = base64.b64encode(output_model).decode()
    href = f'<a href="data:file/output_model;base64,{b64}">Download Trained Model .pkl File</a> (right-click and save as &lt;some_name&gt;.pkl)'
    st.markdown(href, unsafe_allow_html=True)
2 Likes

You can also put download="filename.pickle" in the a tag, that way it’ll tell the browser it’s a download link and give the download a filename by default.

could you give a code snippet of how to implement that? Thanks!

Here you go:

import streamlit as st
import pickle
import base64

x = {"my": "data"}

def download_model(model):
    output_model = pickle.dumps(model)
    b64 = base64.b64encode(output_model).decode()
    href = f'<a href="data:file/output_model;base64,{b64}" download="myfile.pkl">Download Trained Model .pkl File</a>'
    st.markdown(href, unsafe_allow_html=True)


download_model(x)

No need to right click and download, and it’ll name the file.

4 Likes

Thanks for sharing this! I did something similar for downloading an arbitrary file, and I found that it worked for a small file but I got a “Hmm that address doesn’t look right” for a larger file (~3 MB). That’s to be expected because of URL length limits, right? I’m guessing for that size file I’ll need to use nginx or something else to serve those, unless you know a way around that with just streamlit.

Hey @benlindsay,

You can try alternatives from this issue, I like the more ugly way of running a separate Python server pointing on a particular folder where you will save your model then reference it in your Streamlit app through the href trick.

1 Like

Good point, I’ve not hit that as an issue before so didn’t consider it!

I think you’d either need to store it elsewhere (one easy option is upload to S3, then create a signed URL for the download) or sneakily save it to wherever streamlit stores its static files.

1 Like