Can't cache fairseq model

Summary

I’m trying to cache a RoBERTa model which i’m using for a masked language model task, as I don’t want it to be loaded in everytime something changes on the page. I’m trying to use @st.cache for this end, but it fails with the following error trace:

AttributeError: 'function' object has no attribute 'memo'

Streamlit encountered an error while caching the body of load_model(). This is likely due to a bug in /groups/wall2-ilabt-iminds-be/nlp-students/users/sdeblanc/mt2021_streamlit/src/fairseq/fairseq/utils.py near line 455:


if module_path not in import_user_module.memo:
    import_user_module.memo.add(module_path)
Please modify the code above to address this.

If you think this is actually a Streamlit bug, you may file a bug report here.

Traceback:
File "/groups/wall2-ilabt-iminds-be/nlp-students/users/sdeblanc/mt2021_streamlit/WebApp.py", line 25, in <module>
    def load_model(path, modelname):

Steps to reproduce

Environment is reproducable with the following requirements.txt file

pandas
nltk
scikit-learn
pillow
streamlit
torch 
torchvision
-e git+https://github.com/pytorch/fairseq.git#egg=fairseq

Then I tried to load in the model as follows:

# Load in the model
@st.cache()
def load_model(path, modelname):
    model = RobertaModel.from_pretrained(path, checkpoint_file=modelname)
    model.eval()
    return model

model = load_model('../models/robbert/','RobBERT-base.pt')

Expected behavior:

The model should not be reloaded everytime something to the page changes

Actual behavior:

An error is displayed upon loading in the page.

Is this a regression?

No

Debug info

• Streamlit version: 0.70.0
• Python version: 3.8.6
• Using Conda
• OS version: Ubuntu 20
• Browser version: Google Chrome

[sdblanc]

If it makes it any easier for you, I published my basic skeleton on this repository : https://github.com/sdblanc/mt2021_streamlit . When you run it, you’ll see that the model shows the prediction for the sentence ‘Het meisje <mask> daar loopt.’ without any problems. If you then enter the same text within the text field and click ‘Predict’ , it will take quite a long time to make the same prediction, and I think this is because the model is loaded again (as this waiting time is not present in a similar notebook). Using caching, I won’t to mitigate this waiting time.

[jadore801120]

I think it is solved with the experimental_singleton decorator now.

@st.experimental_singleton
def load_model(path, modelname):
    model = RobertaModel.from_pretrained(path, checkpoint_file=modelname)
    model.eval()
    return model