Update setup.py

setup.py

@@ -1,6 +1,7 @@
 import os
 from pathlib import Path
 from urllib.request import urlretrieve
+import streamlit as st
 
 HOME_DIR = Path.home()
 APP_CACHE = HOME_DIR / "olmo_demo"
@@ -9,18 +10,6 @@ OLMO_MODEL_FILE = os.environ.get("OLMO_MODEL_FILE", "OLMo-7B-Instruct-Q4_K_M.ggu
 OLMO_MODEL = APP_CACHE / OLMO_MODEL_FILE
 
 def download_olmo_model(model_file: str | None = None, force=False) -> Path:
-    """Download the OLMO model from the Hugging Face model hub.
-    Parameters
-    ----------
-    model_file : str | None, optional
-        The name of the model file to download, by default None
-    force : bool, optional
-        Whether to force the download even if the file already exists, by default False
-    Returns
-    -------
-    pathlib.Path
-        The path to the downloaded model file
-    """
     if not OLMO_MODEL.exists() or force:
         if model_file is None:
             model_file = OLMO_MODEL_FILE
@@ -28,14 +17,19 @@ def download_olmo_model(model_file: str | None = None, force=False) -> Path:
         else:
             olmo_model = APP_CACHE / model_file
         olmo_model_url = f"https://huggingface.co/ssec-uw/OLMo-7B-Instruct-GGUF/resolve/main/{model_file}"
-        print(f"Downloading model from {olmo_model_url} to {olmo_model}")
-        urlretrieve(olmo_model_url, olmo_model)
-        print("Download complete!")
+        
+        st.info(f"Downloading model from {olmo_model_url} to {olmo_model}")
+        try:
+            APP_CACHE.mkdir(parents=True, exist_ok=True)
+            urlretrieve(olmo_model_url, olmo_model)
+            st.success("Download complete!")
+        except Exception as e:
+            st.error(f"Failed to download model: {str(e)}")
+            raise
         return olmo_model
     else:
-        print(f"Model already exists at {OLMO_MODEL}")
+        st.info(f"Model already exists at {OLMO_MODEL}")
         return OLMO_MODEL
 
 if __name__ == "__main__":
-    APP_CACHE.mkdir(parents=True, exist_ok=True)
-    download_olmo_model()  
+    download_olmo_model()