Update app.py

app.py

@@ -9,6 +9,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 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"
 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))
@@ -132,11 +134,13 @@ class TextGeneration:
 
 
 #@st.cache(allow_output_mutation=True)
-@st.cache(hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
+@st.cache(allow_output_mutation=True, hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
 def load_text_generator():
-    generator = TextGeneration()
-    generator.load()
-    return generator
+    text_generator = TextGeneration()
+    text_generator.load()
+    return text_generator
+
+generator = load_text_generator()
 
 
 def main():
@@ -147,8 +151,6 @@ def main():
         initial_sidebar_state="expanded"
     )
     style()
-    with st.spinner('Cargando el modelo. Por favor, espere...'):
-        generator = load_text_generator()
 
     st.sidebar.markdown(SIDEBAR_INFO, unsafe_allow_html=True)