Update app.py

app.py

@@ -5,12 +5,15 @@ import streamlit as st
 import torch
 from transformers import pipeline, set_seed
 from transformers import AutoTokenizer, AutoModelForCausalLM
-
+import logging
+logger = logging.getLogger()
+logger.addHandler(logging.StreamHandler())
 
 HF_AUTH_TOKEN = os.environ.get("HF_AUTH_TOKEN", None)
 DEVICE = os.environ.get("DEVICE", "cpu")  # cuda:0
 if DEVICE != "cpu" and not torch.cuda.is_available():
     DEVICE = "cpu"
+logger.info(f"DEVICE {DEVICE}")
 DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
 MODEL_NAME = os.environ.get("MODEL_NAME", "bertin-project/bertin-gpt-j-6B")
 MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 1024))
@@ -133,15 +136,13 @@ class TextGeneration:
         )[0]["generated_text"]
 
 
-#@st.cache(allow_output_mutation=True)
-@st.cache(allow_output_mutation=True, hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
+#@st.cache(hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
+@st.cache(allow_output_mutation=True)
 def load_text_generator():
     text_generator = TextGeneration()
     text_generator.load()
     return text_generator
 
-generator = load_text_generator()
-
 
 def main():
     st.set_page_config(
@@ -151,7 +152,7 @@ def main():
         initial_sidebar_state="expanded"
     )
     style()
-
+    generator = load_text_generator()
     st.sidebar.markdown(SIDEBAR_INFO, unsafe_allow_html=True)
 
     max_length = st.sidebar.slider(