Update app.py

app.py

@@ -5,67 +5,149 @@ import importlib.util
 import streamlit as st
 import sys
 import os
+import shutil
 from pathlib import Path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
-import shutil
 
-#___________ cache issue 03_08
-if os.path.exists("private_space_cache"):
-    shutil.rmtree("private_space_cache")  # Forcefully remove corrupted cache folder
-incomplete_files = Path(".").rglob("*.incomplete")
-for file in incomplete_files:
-    try:
-        file.unlink()
-    except Exception as e:
-        print(f"Could not delete {file}: {e}")
-#____________________________________
+# #___________ cache issue 03_08
+# if os.path.exists("private_space_cache"):
+#     shutil.rmtree("private_space_cache")  # Forcefully remove corrupted cache folder
+# incomplete_files = Path(".").rglob("*.incomplete")
+# for file in incomplete_files:
+#     try:
+#         file.unlink()
+#     except Exception as e:
+#         print(f"Could not delete {file}: {e}")
+# #____________________________________
 
-HF_TOKEN_LLAMA = os.environ.get("HF_TOKEN_LLAMA")
-login(token=HF_TOKEN_LLAMA)
+# HF_TOKEN_LLAMA = os.environ.get("HF_TOKEN_LLAMA")
+# login(token=HF_TOKEN_LLAMA)
+
+# HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN 
+# login(token=HF_TOKEN)
+
+
+# USER_NAME = os.getenv("USER_NAME", "").strip().strip('"')
+# PRIVATE_SPACE_NAME = os.getenv("PRIVATE_SPACE_NAME", "").strip().strip('"')
+
+# #Construct the repo ID
+# REPO_ID = f"{USER_NAME}/{PRIVATE_SPACE_NAME}"
+
+# REPO_TYPE = "space"
+
+# # sys.path.append(repo_dir)
+
+# # Download the entire space, including the fine-tuned model folder
+# repo_dir = snapshot_download(
+#     repo_id=REPO_ID,
+#     repo_type=REPO_TYPE,
+#     token=HF_TOKEN,
+#     cache_dir="private_space_cache",
+#     force_download=True  # Forces redownload
+# )
+
+
+# # Change the working directory to the downloaded snapshot directory
+# # This step is very imporptant
+# os.chdir(repo_dir)
+
+# # # Add repo directory to sys.path so Python can find modules inside it
+# sys.path.append(repo_dir)
 
-HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN 
-login(token=HF_TOKEN)
 
+# # Download specific files (if snapshot_download wasn't used)
+# app_path = hf_hub_download(
+#     repo_id=REPO_ID,
+#     filename="app.py",
+#     repo_type=REPO_TYPE
+# )
 
+# # Load and execute `app.py`
+# spec_app = importlib.util.spec_from_file_location("*", app_path)
+# app_module = importlib.util.module_from_spec(spec_app)
+# spec_app.loader.exec_module(app_module)
+# # Now you can use functions from utils_module
+# result = app_module.main()
+
+
+
+# Helper function to clean up cache
+def clear_cache(cache_dir="private_space_cache"):
+    """Remove cache directory and all .incomplete files."""
+    if os.path.exists(cache_dir):
+        shutil.rmtree(cache_dir, ignore_errors=True)
+    for file in Path(".").rglob("*.incomplete"):
+        try:
+            file.unlink()
+        except Exception as e:
+            st.warning(f"Could not delete {file}: {e}")
+
+# Streamlit app setup
+st.title("Math Question Labeling Tool (Public Interface)")
+
+# Get environment variables
+HF_TOKEN_LLAMA = os.environ.get("HF_TOKEN_LLAMA")
+HF_TOKEN = os.environ.get("HF_TOKEN")
 USER_NAME = os.getenv("USER_NAME", "").strip().strip('"')
 PRIVATE_SPACE_NAME = os.getenv("PRIVATE_SPACE_NAME", "").strip().strip('"')
 
-#Construct the repo ID
-REPO_ID = f"{USER_NAME}/{PRIVATE_SPACE_NAME}"
+# Log in to Hugging Face
+login(token=HF_TOKEN_LLAMA)
+login(token=HF_TOKEN)
 
+# Construct the repo ID
+REPO_ID = f"{USER_NAME}/{PRIVATE_SPACE_NAME}"
 REPO_TYPE = "space"
 
-# sys.path.append(repo_dir)
-
-# Download the entire space, including the fine-tuned model folder
-repo_dir = snapshot_download(
-    repo_id=REPO_ID,
-    repo_type=REPO_TYPE,
-    token=HF_TOKEN,
-    cache_dir="private_space_cache",
-    force_download=True  # Forces redownload
-)
+# Always clear cache before downloading
+with st.spinner("Clearing old cache..."):
+    clear_cache()
 
+# Download the private space
+with st.spinner("Downloading private space..."):
+    try:
+        repo_dir = snapshot_download(
+            repo_id=REPO_ID,
+            repo_type=REPO_TYPE,
+            token=HF_TOKEN,
+            cache_dir="private_space_cache",
+            force_download=True,  # Always redownload to avoid stale cache
+            local_dir_use_symlinks=False  # Avoid symlinks that might break
+        )
+    except Exception as e:
+        st.error(f"Failed to download private space: {e}")
+        st.stop()
 
-# Change the working directory to the downloaded snapshot directory
-# This step is very imporptant
+# Change working directory and update sys.path
 os.chdir(repo_dir)
-
-# # Add repo directory to sys.path so Python can find modules inside it
 sys.path.append(repo_dir)
 
+# Download and load app.py from the private space
+try:
+    app_path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename="app.py",
+        repo_type=REPO_TYPE,
+        token=HF_TOKEN,
+        force_download=True  # Ensure fresh download
+    )
+except Exception as e:
+    st.error(f"Failed to download app.py: {e}")
+    st.stop()
 
-# Download specific files (if snapshot_download wasn't used)
-app_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename="app.py",
-    repo_type=REPO_TYPE
-)
-
-# Load and execute `app.py`
-spec_app = importlib.util.spec_from_file_location("*", app_path)
+# Load and execute app.py
+spec_app = importlib.util.spec_from_file_location("app_module", app_path)
 app_module = importlib.util.module_from_spec(spec_app)
 spec_app.loader.exec_module(app_module)
-# Now you can use functions from utils_module
-result = app_module.main()
+
+# Run the private app's main function
+try:
+    result = app_module.main()
+except Exception as e:
+    st.error(f"Error running private app: {e}")
+    st.stop()
+
+# Optional: Clean up after execution
+with st.spinner("Cleaning up..."):
+    clear_cache()