Update app.py

app.py

@@ -54,52 +54,55 @@ def get_features_single(formula):
 
 def predict_features(formula):
     try:
-        print(f"Received formula: {formula}")  # 打印接收到的化学式
+        # 验证化学式的有效性
+        try:
+            comp = Composition(formula)  # 如果无效会抛出异常
+            _ = comp.get_reduced_formula_and_factor()  # 进一步验证
+        except Exception:
+            return {"Error": "Invalid chemical formula. Please check your input and try again."}
+
+        # Generate features for the input formula
         df = get_features_single(formula)
-        print(f"Generated features dataframe:\n{df.head()}")  # 打印生成的特征数据
+        if df.empty:
+            return {"Error": "Unable to generate features. Please check your input and try again."}
 
         X_user = df.iloc[:, 2:].fillna(0)
-        print(f"Input features for model:\n{X_user.head()}")  # 打印输入特征
 
-        # 加载模型
+        # Load saved model
         model_path = "saved_model/lgbm_model.pkl"
         loaded_model = joblib.load(model_path)
-        print("Model loaded successfully.")  # 确认模型加载成功
 
-        # 加载编码器和映射
+        # Load saved LabelEncoder
         label_encoder_path = "saved_model/label_encoder.pkl"
         label_encoder = joblib.load(label_encoder_path)
-        print("Label encoder loaded successfully.")
 
+        # Load Layer Group Mapping
         mapping_path = "saved_model/layer_group_mapping.pkl"
         layer_group_mapping = joblib.load(mapping_path)
-        print("Layer group mapping loaded successfully.")
 
-        # 模型预测
+        # Predict probabilities
         prediction_probs = loaded_model.predict_proba(X_user)[0]
-        print(f"Prediction probabilities: {prediction_probs}")
 
-        # 获取 Top 5
+        # Get top 5 predictions
         top_5_indices = np.argsort(prediction_probs)[-5:][::-1]
-        print(f"Top 5 indices: {top_5_indices}")
-
         top_5_numbers = label_encoder.inverse_transform(top_5_indices)
-        print(f"Top 5 layer group numbers: {top_5_numbers}")
-
         top_5_names = [layer_group_mapping.get(num, "Unknown Layer Group") for num in top_5_numbers]
-        print(f"Top 5 layer group names: {top_5_names}")
-
         top_5_probs = prediction_probs[top_5_indices]
+
+        # Prepare top 5 results as list of lists
         top_5_results = [
             [num, name, f"{prob:.2%}"]
             for num, name, prob in zip(top_5_numbers, top_5_names, top_5_probs)
         ]
-        print(f"Top 5 results: {top_5_results}")
 
+        # Predict using the loaded model
         prediction = loaded_model.predict(X_user)
+
+        # Decode Layer Group Number
         decoded_prediction = label_encoder.inverse_transform(prediction)[0]
+
+        # Map to Layer Group Name
         layer_group_name = layer_group_mapping.get(decoded_prediction, "Unknown Layer Group")
-        print(f"Final prediction: {layer_group_name}")
 
         return {
             "Formula": formula,
@@ -108,9 +111,7 @@ def predict_features(formula):
             "Top 5 Predictions": top_5_results
         }
     except Exception as e:
-        print(f"Error in predict_features: {str(e)}")
-        return {"Error": str(e)}
-
+        return {"Error": f"An unexpected error occurred: {str(e)}"}
 
 
 # Define a more visually appealing Gradio interface
@@ -133,7 +134,10 @@ def gradio_ui():
 
         def update_output(formula):
             result = predict_features(formula)
-            return result["Predicted Layer Group Name"], result["Top 5 Predictions"]
+            if "Error" in result:
+                # 返回错误消息到 Gradio 界面
+                return f"Error: {result['Error']}", []
+            return result.get("Predicted Layer Group Name", "Unknown"), result.get("Top 5 Predictions", [])
 
         predict_button.click(fn=update_output, inputs=[formula_input], outputs=[final_prediction, top_5_table])