fix

app.py

@@ -30,8 +30,11 @@ ATTRIBUTE_MODELS = {
 }
 
 LANGUAGE_MODELS = {
-    "Russian": ('sberbank-ai/rugpt3small_based_on_gpt2',),
-    "English": ("distilgpt2", )
+    "Russian": (
+        'sberbank-ai/rugpt3small_based_on_gpt2',
+        "sberbank-ai/rugpt3large_based_on_gpt2"
+    ),
+    "English": ("distilgpt2", "gpt2", "EleutherAI/gpt-neo-1.3B")
 }
 
 ATTRIBUTE_MODEL_LABEL = {
@@ -81,10 +84,10 @@ def main():
         label2id = cls_model_config.label2id
         print(list(label2id.keys()))
         label_key = st.selectbox(ATTRIBUTE_LABEL[language], [list(label2id.keys())[-1]])
-        target_label_id = 0
+        target_label_id = 1
     prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language])
     alpha = st.slider("Alpha", min_value=-10, max_value=10, step=1, value=0)
-    entropy_threshold = st.slider("Entropy threshold", min_value=0., max_value=5., step=.1, value=5.)
+    entropy_threshold = st.slider("Entropy threshold", min_value=0., max_value=5., step=.1, value=2.)
     auth_token = os.environ.get('TOKEN') or True
     with st.spinner('Running inference...'):
         text = inference(
@@ -101,13 +104,13 @@ def main():
 
 
 
-@st.cache(hash_funcs={str: lambda lm_model_name: hash(lm_model_name)}, allow_output_mutation=True)
+@st.cache
 def load_generator(lm_model_name: str) -> Generator:
     with st.spinner('Loading language model...'):
         generator = Generator(lm_model_name=lm_model_name, device=device)
         return generator
 
-
+@st.cache(hash_funcs={tokenizers.Tokenizer: lambda lm_tokenizer: hash(lm_tokenizer.to_str)}, allow_output_mutation=True)
 def load_sampler(cls_model_name, lm_tokenizer):
     with st.spinner('Loading classifier model...'):
         sampler = CAIFSampler(classifier_name=cls_model_name, lm_tokenizer=lm_tokenizer, device=device)