I would like to use the great Huggingface transformers in a streamlit app. Unfortunately, the caching of the loaded models does not work as expected. It gives cache misses even though I previously had loaded the model for that input/key. Sometimes the
load function that loads the current model is even re-executed when I only change the text in the text field. Note that I do use the
allow_output_mutation=True argument. I added a sample script that works after installing
transformers. The models are community models that can be accessed by anyone.
And then there is this other annoying thing. When I rerun the script after some code change, I get this error:
ValueError: Custom>TFBertMainLayer has already been registered to <class 'transformers.modeling_tf_bert.TFBertMainLayer'>
Thanks for the help!
import streamlit as st from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch @st.cache(allow_output_mutation=True) def load_bert(classif): if classif=="hatespeech": name = "deepset/bert-base-german-cased-hatespeech-GermEval18Coarse" elif classif=="sentiment": name = "nlptown/bert-base-multilingual-uncased-sentiment" else: return None, None t = AutoTokenizer.from_pretrained(name) m = AutoModelForSequenceClassification.from_pretrained(name) m.eval() return m, t classifier = st.sidebar.selectbox('Classifier', ["hatespeech", "sentiment"]) model, tokenizer = load_bert(classifier) text = st.text_area("Please enter text to classify", "Ich bin ein Berliner.") if text: input_ids = torch.tensor(tokenizer.encode(text, add_special_tokens=True)).unsqueeze(0) # Batch size 1 output = model(input_ids) st.text(output)