Update app.py

app.py

@@ -15,9 +15,20 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from setup import download_olmo_model
-model_path = download_olmo_model()
-# import subprocess
+from setup import download_olmo_model, OLMO_MODEL
+
+# Ensure model is downloaded before proceeding
+@st.cache_resource
+def ensure_model_downloaded():
+    try:
+        model_path = download_olmo_model()
+        st.success(f"Model successfully loaded from {model_path}")
+        return model_path
+    except Exception as e:
+        st.error(f"Failed to download or locate the model: {str(e)}")
+        st.stop()
+
+model_path = ensure_model_downloaded()# import subprocess
 
 # # Define the path to your bash script
 # script_path = "./start.sh"
@@ -90,13 +101,17 @@ def get_chain(temperature):
     retriever = load_retriever(docstore_path,chroma_path,embeddings,child_splitter,parent_splitter)
     
     # Replace the local OLMOLLM with the Hugging Face model
-    llm = LlamaCpp(
-        model_path=str(model_path),
-        temperature=temperature,
-        max_tokens=3000,
-        verbose=False,
-        echo=False
-    )
+    try:
+        llm = LlamaCpp(
+            model_path=str(model_path),
+            temperature=temperature,
+            max_tokens=3000,
+            verbose=False,
+            echo=False
+        )
+    except Exception as e:
+        st.error(f"Failed to initialize LlamaCpp: {str(e)}")
+        st.stop()