Update app.py

app.py

@@ -119,11 +119,15 @@ def prediction(df):
   #predicition = []
   #for i in list(y_pred):
   if y_pred ==0:
-      st.write('Rouge Component is Good')
+      return 'Rouge Component is Good'
   else:
-      st.write('Rouge Component is not good')
+      return 'Rouge Component is not good'
   #X_test['Actual_time_to_repair'] = y_test
   #X_test['Predicted_time_to_repair'] = predicition
   # X_test.to_csv(r'/content/drive/MyDrive/Colab Notebooks/HAL/rogue_test_data.csv')
   #print(X_test.head())
-prediction(user_data)
+
+y_pred = prediction(user_data)
+
+if st.button("Predict"):
+    st.subheader(y_pred)