Update app.py

app.py

@@ -17,7 +17,9 @@ from catboost import CatBoostRegressor
 import warnings
 warnings.filterwarnings('ignore')
 
-def load_embeddings(embeddings_file_path):
+# Load datasets (Assuming the datasets are in the same directory)
+def load_embeddings():
+    embeddings_file_path = 'county_embeddings.csv'  # Adjust the file name if necessary
     county_embeddings = pd.read_csv(embeddings_file_path).set_index('place')
     numeric_cols = county_embeddings.select_dtypes(include=['number']).columns
     county_embeddings_numeric = county_embeddings[numeric_cols]
@@ -28,7 +30,8 @@ def load_embeddings(embeddings_file_path):
     county_embeddings_pca = pca.transform(county_embeddings_imputed)
     return county_embeddings, county_embeddings_pca, pca, imputer
 
-def load_unemployment_data(unemployment_file_path):
+def load_unemployment_data():
+    unemployment_file_path = 'county_unemployment.csv'  # Adjust the file name if necessary
     unemployment_data = pd.read_csv(unemployment_file_path).set_index('place')
     unemployment_long = unemployment_data.reset_index().melt(id_vars='place', var_name='date', value_name='unemployment_rate')
     return unemployment_long
@@ -65,15 +68,18 @@ def preprocess_data(county_embeddings, county_embeddings_pca, unemployment_long,
 
     return X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train
 
-def train_and_evaluate_models(X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train):
+def train_and_evaluate_models(X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train, selected_models):
     # Define models
-    models = {
+    all_models = {
         "Random Forest": RandomForestRegressor(n_estimators=100, random_state=42),
         "XGBoost": XGBRegressor(n_estimators=100, random_state=42, tree_method='gpu_hist'),
         "Ridge Regression": Ridge(alpha=1.0),
         "CatBoost": CatBoostRegressor(iterations=100, random_seed=42, task_type="GPU")
     }
 
+    # Filter selected models
+    models = {name: model for name, model in all_models.items() if name in selected_models}
+
     results = {}
     feature_importances = {}
 
@@ -123,10 +129,10 @@ def plot_metrics(results):
 
     return rmse_plot, r2_plot
 
-def main(embeddings_file_path, unemployment_file_path):
+def main(selected_models):
     # Load data
-    county_embeddings, county_embeddings_pca, pca, imputer = load_embeddings(embeddings_file_path)
-    unemployment_long = load_unemployment_data(unemployment_file_path)
+    county_embeddings, county_embeddings_pca, pca, imputer = load_embeddings()
+    unemployment_long = load_unemployment_data()
 
     # Preprocess data
     X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train = preprocess_data(
@@ -135,7 +141,7 @@ def main(embeddings_file_path, unemployment_file_path):
 
     # Train and evaluate models
     results, feature_importances, feature_names = train_and_evaluate_models(
-        X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train
+        X_train_pca, X_test_pca, y_train, y_test, numeric_cols_train, selected_models
     )
 
     # Plot metrics
@@ -151,41 +157,44 @@ def main(embeddings_file_path, unemployment_file_path):
 
 def gradio_app():
     with gr.Blocks() as demo:
-        gr.Markdown("# County-Level Unemployment Rate Forecasting")
-        gr.Markdown("Upload county embeddings and unemployment data to train models and visualize results.")
+        gr.Markdown("<h1 style='text-align: center'>County-Level Unemployment Rate Forecasting</h1>")
+        gr.Markdown("This app forecasts county-level unemployment rates using various machine learning models with GPU acceleration.")
 
         with gr.Row():
-            embeddings_file = gr.File(label="Upload County Embeddings CSV")
-            unemployment_file = gr.File(label="Upload Unemployment Data CSV")
+            with gr.Column(scale=1):
+                gr.Markdown("### Select Models to Train")
+                model_choices = ["Random Forest", "XGBoost", "Ridge Regression", "CatBoost"]
+                selected_models = gr.CheckboxGroup(choices=model_choices, value=model_choices, label="Models")
 
-        run_button = gr.Button("Run Analysis")
+                run_button = gr.Button("Run Analysis")
 
-        output_results = gr.JSON(label="Model Performance Metrics")
-        output_rmse_plot = gr.Plot(label="RMSE Comparison")
-        output_r2_plot = gr.Plot(label="R-squared Comparison")
-        output_feature_importance = gr.Plot(label="Feature Importances")
+            with gr.Column(scale=2):
+                output_results = gr.JSON(label="Model Performance Metrics")
+                output_rmse_plot = gr.Plot(label="RMSE Comparison")
+                output_r2_plot = gr.Plot(label="R-squared Comparison")
 
-        def run_analysis(embeddings_file, unemployment_file):
-            if embeddings_file is None or unemployment_file is None:
-                return gr.update(value="Please upload both embeddings and unemployment data files."), None, None, None
+        gr.Markdown("### Feature Importances")
+        output_feature_importance = gr.TabbedInterface([], [])
 
-            # Read files
-            embeddings_file_path = embeddings_file.name
-            unemployment_file_path = unemployment_file.name
+        def run_analysis(selected_models):
+            if not selected_models:
+                return gr.update(value="Please select at least one model to train."), None, None, gr.update(tabs=[], contents=[])
 
             # Run main analysis
-            results, rmse_plot, r2_plot, feature_importance_plots = main(embeddings_file_path, unemployment_file_path)
+            results, rmse_plot, r2_plot, feature_importance_plots = main(selected_models)
 
-            # For simplicity, display feature importance of Random Forest (if available)
-            fi_plot = None
-            if 'Random Forest' in feature_importance_plots:
-                fi_plot = feature_importance_plots['Random Forest']
+            # Prepare feature importance plots
+            fi_tabs = []
+            fi_plots = []
+            for model_name, fig in feature_importance_plots.items():
+                fi_tabs.append(model_name)
+                fi_plots.append(fig)
 
-            return results, rmse_plot, r2_plot, fi_plot
+            return results, rmse_plot, r2_plot, gr.update(tabs=fi_tabs, contents=fi_plots)
 
         run_button.click(
             run_analysis,
-            inputs=[embeddings_file, unemployment_file],
+            inputs=selected_models,
             outputs=[output_results, output_rmse_plot, output_r2_plot, output_feature_importance]
         )