Update app.py

app.py

@@ -449,10 +449,18 @@ with gr.Blocks(css=".gradio-container {max-width: 1200px;}") as demo:
     # Second Tab: Co-occurrence on demand
     with gr.Tab("Co-occurrence"):
         gr.Markdown("## Co-Occurrence Plot\nSelect two features + one label to see a distribution.")
-        with gr.Row():
-            feat1_dd = gr.Dropdown(sorted(df.columns), label="Feature 1")
-            feat2_dd = gr.Dropdown(sorted(df.columns), label="Feature 2")
-            label_dd = gr.Dropdown(sorted(df.columns), label="Label Column")
+
+        # 1) We gather the 'input features' from input_mapping keys:
+        list_of_features = sorted(input_mapping.keys())
+
+        # 2) We gather the 'label columns' from predictor.prediction_map keys:
+        list_of_labels = sorted(predictor.prediction_map.keys())
+
+        # 3) Build three dropdowns:
+        feat1_dd = gr.Dropdown(choices=list_of_features, label="Feature 1")
+        feat2_dd = gr.Dropdown(choices=list_of_features, label="Feature 2")
+        label_dd = gr.Dropdown(choices=list_of_labels,  label="Label Column")
+
         generate_btn = gr.Button("Generate Plot")
         co_occ_output = gr.Plot()