Update app.py

app.py

@@ -121,7 +121,20 @@ def main():
     def restore_prediction_state():
         if "prediction" in st.session_state:
             del st.session_state.prediction
-    option = st.selectbox("Select a model to use:", models_available, on_change=restore_prediction_state)
+    
+    option_ai = st.selectbox("Generated/Revised", ["Generated", "Revised"], on_change=restore_prediction_state)
+    if option_ai == "Generated":
+        option = st.selectbox("Select a model to use:", {"LR Generated":"models/baseline_model_lr_generated.joblib", 
+                                                         "NB Generated": "models/baseline_model_nb_generated.joblib",
+                                                         "D-BERT Generated": "GradApplicationDocuments/DistilBERTGradGen",
+                                                         "BERT Generated": "GradApplicationDocuments/BERTGradGen"}, 
+                              on_change=restore_prediction_state)
+    elif option_ai == "Revised":
+        option = st.selectbox("Select a model to use:", {"LR Revised":"models/baseline_model_lr_revised.joblib", 
+                                                         "NB Revised": "models/baseline_model_nb_revised.joblib",
+                                                         "D-BERT Revised": "GradApplicationDocuments/DistilBERTGradRevised",
+                                                         "BERT Revised": "GradApplicationDocuments/BERTGradRevised"}, 
+                              on_change=restore_prediction_state)
     
 
     # Load the selected trained model