Update app.py

app.py

@@ -9,9 +9,9 @@ from huggingface_hub import hf_hub_download, list_repo_files
 import logging
 import traceback
 import sklearn
+import pkg_resources
 
 # 版本检查
-import pkg_resources
 required_versions = {
     'numpy': '1.23.5',
     'scipy': '1.10.1',
@@ -23,17 +23,17 @@ def check_versions():
     for package, required_version in required_versions.items():
         try:
             installed_version = pkg_resources.get_distribution(package).version
-            st.write(f"{package} version: {installed_version} (required: {required_version})")
             if installed_version != required_version:
-                st.warning(f"Warning: {package} version mismatch. Required: {required_version}, Installed: {installed_version}")
+                logger.warning(f"Warning: {package} version mismatch. Required: {required_version}, Installed: {installed_version}")
+                return False
         except pkg_resources.DistributionNotFound:
-            st.error(f"Error: {package} not found!")
+            logger.error(f"Error: {package} not found!")
             return False
     return True
 
 # Set up logging
 logging.basicConfig(
-    level=logging.DEBUG,
+    level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
@@ -45,35 +45,18 @@ st.set_page_config(
     layout="centered"
 )
 
-# 在加载模型之前检查版本
-if not check_versions():
-    st.error("Package version requirements not met. Please check the logs.")
-    st.stop()
-
 @st.cache_resource
 def load_model():
     """Load model from Hugging Face Hub"""
     try:
-        # 检查scikit-learn版本
-        logger.info(f"Using scikit-learn version: {sklearn.__version__}")
-        st.write(f"Using scikit-learn version: {sklearn.__version__}")
-        
-        # 首先列出仓库中的所有文件
-        logger.info("Listing repository files...")
-        try:
-            files = list_repo_files("404Brain-Not-Found-yeah/healing-music-classifier")
-            logger.info(f"Repository files: {files}")
-            st.write("Available files in repository:", files)
-        except Exception as e:
-            logger.error(f"Error listing repository files: {str(e)}\n{traceback.format_exc()}")
-            st.error(f"Error listing repository files: {str(e)}")
+        # 检查版本
+        if not check_versions():
+            logger.error("Package version requirements not met")
             return None, None
 
         # 创建临时目录
         os.makedirs("temp_models", exist_ok=True)
-        logger.info("Created temp_models directory")
 
-        logger.info("Downloading model from Hugging Face Hub...")
         # 下载模型文件
         try:
             model_path = hf_hub_download(
@@ -81,48 +64,21 @@ def load_model():
                 filename="models/model.joblib",
                 local_dir="temp_models"
             )
-            logger.info(f"Model downloaded to: {model_path}")
-            st.write(f"Model downloaded to: {model_path}")
-        except Exception as e:
-            logger.error(f"Error downloading model: {str(e)}\n{traceback.format_exc()}")
-            st.error(f"Error downloading model: {str(e)}")
-            return None, None
-
-        # 下载scaler文件
-        try:
             scaler_path = hf_hub_download(
                 repo_id="404Brain-Not-Found-yeah/healing-music-classifier",
                 filename="models/scaler.joblib",
                 local_dir="temp_models"
             )
-            logger.info(f"Scaler downloaded to: {scaler_path}")
-            st.write(f"Scaler downloaded to: {scaler_path}")
         except Exception as e:
-            logger.error(f"Error downloading scaler: {str(e)}\n{traceback.format_exc()}")
-            st.error(f"Error downloading scaler: {str(e)}")
+            logger.error(f"Error downloading model files: {str(e)}")
             return None, None
         
         # 加载模型文件
         try:
-            logger.info("Loading model and scaler...")
-            # 检查文件是否存在
-            if not os.path.exists(model_path):
-                logger.error(f"Model file not found at: {model_path}")
-                st.error(f"Model file not found at: {model_path}")
-                return None, None
-            if not os.path.exists(scaler_path):
-                logger.error(f"Scaler file not found at: {scaler_path}")
-                st.error(f"Scaler file not found at: {scaler_path}")
+            if not os.path.exists(model_path) or not os.path.exists(scaler_path):
+                logger.error("Model files not found")
                 return None, None
 
-            # 检查文件大小
-            model_size = os.path.getsize(model_path)
-            scaler_size = os.path.getsize(scaler_path)
-            logger.info(f"Model file size: {model_size} bytes")
-            logger.info(f"Scaler file size: {scaler_size} bytes")
-            st.write(f"Model file size: {model_size} bytes")
-            st.write(f"Scaler file size: {scaler_size} bytes")
-
             # 尝���使用不同的pickle协议加载
             try:
                 model = joblib.load(model_path)
@@ -136,17 +92,13 @@ def load_model():
                 with open(scaler_path, 'rb') as f:
                     scaler = pickle.load(f, encoding='latin1')
 
-            logger.info("Model and scaler loaded successfully")
-            st.success("Model and scaler loaded successfully!")
             return model, scaler
         except Exception as e:
-            logger.error(f"Error loading model/scaler files: {str(e)}\n{traceback.format_exc()}")
-            st.error(f"Error loading model/scaler files: {str(e)}")
+            logger.error(f"Error loading model/scaler files: {str(e)}")
             return None, None
             
     except Exception as e:
-        logger.error(f"Unexpected error in load_model: {str(e)}\n{traceback.format_exc()}")
-        st.error(f"Unexpected error in load_model: {str(e)}")
+        logger.error(f"Unexpected error in load_model: {str(e)}")
         return None, None
 
 def main():
@@ -167,7 +119,6 @@ def main():
         try:
             # Create temporary file
             with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
-                # Write uploaded file content
                 tmp_file.write(uploaded_file.getvalue())
                 tmp_file_path = tmp_file.name
 
@@ -178,7 +129,7 @@ def main():
             # Load model
             model, scaler = load_model()
             if model is None or scaler is None:
-                st.error("Model loading failed. Please check the logs for details.")
+                st.error("Model loading failed. Please try again later.")
                 return
 
             progress_bar.progress(50)
@@ -197,8 +148,8 @@ def main():
                 healing_probability = model.predict_proba(scaled_features)[0][1]
                 progress_bar.progress(90)
             except Exception as e:
-                logger.error(f"Error during prediction: {str(e)}\n{traceback.format_exc()}")
-                st.error(f"Error during prediction: {str(e)}")
+                logger.error(f"Error during prediction: {str(e)}")
+                st.error("Error during prediction. Please try again.")
                 return
 
             # Display results
@@ -220,7 +171,7 @@ def main():
                 st.warning("This music has limited healing potential. 🎵")
 
         except Exception as e:
-            st.error(f"An unexpected error occurred: {str(e)}")
+            st.error("An unexpected error occurred. Please try again.")
             logger.exception("Unexpected error")
 
         finally: