Update app.py

app.py

@@ -36,7 +36,7 @@ COMPANIES = [
 ]
 
 class StockPredictor:
-    def __init__(self, data, model_type):
+    def __init__(self, data, model_type='Prophet'):
         self.data = data
         self.model_type = model_type
         self.model = None
@@ -349,7 +349,7 @@ def main():
         predict_stock_prices()
 
 def test_model():
-    st.header("Test Prophet Model")
+    st.header("Test Stock Prediction Model")
 
     col1, col2 = st.columns(2)
 
@@ -357,6 +357,9 @@ def test_model():
         company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
         test_split = st.slider("Test Data Split", 0.1, 0.5, 0.2, 0.05)
 
+    with col2:
+        model_type = st.selectbox("Select Model Type", ['Prophet', 'LSTM', 'SARIMA', 'XGBoost', 'RandomForest'])
+
     if st.button("Train and Test Model"):
         with st.spinner("Fetching data and training model..."):
             company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
@@ -375,7 +378,7 @@ def test_model():
                 train_data = data.iloc[:split_index]
                 test_data = data.iloc[split_index:]
 
-                predictor = StockPredictor(train_data)  # Updated: removed model_type argument
+                predictor = StockPredictor(train_data, model_type)  # Updated: added model_type argument
                 predictor.preprocess_data()
                 if predictor.train_model():
                     test_pred = predictor.predict(days=len(test_data))
@@ -392,7 +395,7 @@ def test_model():
                         plot = create_test_plot(predictor.data, test_data, test_pred, company_name)
                         st.plotly_chart(plot, use_container_width=True)
                 else:
-                    st.error("Failed to train the Prophet model. Please try a different dataset.")
+                    st.error(f"Failed to train the {model_type} model. Please try a different dataset or model type.")
 
 def predict_stock_prices():
     st.header("Predict Stock Prices")
@@ -403,6 +406,9 @@ def predict_stock_prices():
         company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
         days_to_predict = st.slider("Days to Predict", 1, 365, 30)
 
+    with col2:
+        model_type = st.selectbox("Select Model Type", ['Prophet', 'LSTM', 'SARIMA', 'XGBoost', 'RandomForest'])
+
     if st.button("Predict Stock Prices"):
         with st.spinner("Fetching data and making predictions..."):
             company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
@@ -417,7 +423,7 @@ def predict_stock_prices():
 
                 st.markdown(get_table_download_link(data), unsafe_allow_html=True)
 
-                predictor = StockPredictor(data)  # Updated: removed model_type argument
+                predictor = StockPredictor(data, model_type)  # Updated: added model_type argument
                 predictor.preprocess_data()
                 if predictor.train_model():
                     predictions = predictor.predict(days=days_to_predict)
@@ -439,7 +445,7 @@ def predict_stock_prices():
                         for item in news:
                             st.markdown(f"[{item['title']}]({item['link']}) ({item['pubDate']})")
                 else:
-                    st.error("Failed to train the Prophet model. Please try a different dataset.")
+                    st.error(f"Failed to train the {model_type} model. Please try a different dataset or model type.")
 
 def explore_data():
     st.header("Explore Stock Data")