Update app.py

app.py

@@ -5,6 +5,7 @@ import pandas as pd
 from matminer.featurizers.conversions import StrToComposition
 from matminer.featurizers.base import MultipleFeaturizer
 from matminer.featurizers import composition as cf
+import joblib
 
 # Define feature calculators
 feature_calculators = MultipleFeaturizer([
@@ -49,24 +50,85 @@ def get_features_single(formula):
     mlmd = mlmdd_single(formula)
     ext_mag = generate_single(formula)
     result = pd.concat([ext_mag, mlmd], axis=1)
-    return result.iloc[:, :10]  # Select the first 10 features
+    return result
 
 def predict_features(formula):
-    """
-    Gradio prediction function to return top 10 features for the input formula.
-    """
     try:
+        # Generate features for the input formula
         df = get_features_single(formula)
-        return df.head(1).to_dict(orient="records")[0]
+        X_user = df.iloc[:, 2:].fillna(0)
+
+        # Load saved model
+        model_path = "saved_model/lgbm_model.pkl"
+        loaded_model = joblib.load(model_path)
+
+        # Load saved LabelEncoder
+        label_encoder_path = "saved_model/label_encoder.pkl"
+        label_encoder = joblib.load(label_encoder_path)
+
+        # Load Layer Group Mapping
+        mapping_path = "saved_model/layer_group_mapping.pkl"
+        layer_group_mapping = joblib.load(mapping_path)
+
+        # Predict probabilities
+        prediction_probs = loaded_model.predict_proba(X_user)[0]
+
+        # Get top 5 predictions
+        top_5_indices = np.argsort(prediction_probs)[-5:][::-1]
+        top_5_numbers = label_encoder.inverse_transform(top_5_indices)
+        top_5_names = [layer_group_mapping.get(num, "Unknown Layer Group") for num in top_5_numbers]
+        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)
+        ]
+
+        # 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")
+
+        return {
+            "Formula": formula,
+            "Predicted Layer Group Number": decoded_prediction,
+            "Predicted Layer Group Name": layer_group_name,
+            "Top 5 Predictions": top_5_results
+        }
     except Exception as e:
         return {"Error": str(e)}
 
-# Define Gradio interface
-iface = gr.Interface(
-    fn=predict_features,
-    inputs=gr.Textbox(label="Enter Chemical Formula", placeholder="E.g., BrIPtZrS2"),
-    outputs=gr.JSON(label="Top 10 Features")
-)
+
+# Define a more visually appealing Gradio interface
+def gradio_ui():
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # 🔬 2D Material Layer Group Predictor
+        Enter a chemical formula of 2D Material below to get the predicted layer group information along with the top 5 probable groups.
+        """)
+
+        with gr.Row():
+            formula_input = gr.Textbox(label="Enter Chemical Formula", placeholder="E.g., BrIPtZrS2")
+            predict_button = gr.Button("Predict")
+
+        with gr.Row():
+            final_prediction = gr.Textbox(label="Final Prediction (Layer Group Name)", interactive=False)
+
+        with gr.Row():
+            top_5_table = gr.Dataframe(headers=["Layer Group Number", "Layer Group Name", "Probability"], interactive=False)
+
+        def update_output(formula):
+            result = predict_features(formula)
+            return result["Predicted Layer Group Name"], result["Top 5 Predictions"]
+
+        predict_button.click(fn=update_output, inputs=[formula_input], outputs=[final_prediction, top_5_table])
+
+    return demo
 
 if __name__ == "__main__":
-    iface.launch()
+    gradio_ui().launch()