Update app.py

app.py

@@ -7,12 +7,14 @@ import shap
 import lime.lime_tabular
 import optuna
 import wandb
+import json
 from smolagents import HfApiModel, CodeAgent
 from huggingface_hub import login
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.model_selection import train_test_split, cross_val_score
 from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
 from sklearn.preprocessing import LabelEncoder
+from PIL import Image
 
 # Authenticate with Hugging Face
 hf_token = os.getenv("HF_TOKEN")
@@ -79,6 +81,7 @@ def run_agent(_):
        - At least 3 visualizations showing important trends.
     4. Derive at least 3 actionable real-world insights.
        5. Save all visualizations to ./figures/ directory.
+       6. "Ensure that all visualizations have figure size at least 8x6 and saved at 150+ dpi."
        Return a JSON object with keys:
          - 'insights': clean bullet-point insights.
          - 'figures': list of file paths of generated visualizations.
@@ -89,17 +92,29 @@ def run_agent(_):
         additional_args={"source_file": temp_file.name}
     )
 
-    # Process result as a dictionary if possible
+    if isinstance(result, str):
+    try:
+        result = json.loads(result)
+    except json.JSONDecodeError:
+        insights = "Failed to parse result from agent."
+        return insights, []
+
     if isinstance(result, dict):
         insights = result.get("insights", "No insights generated.")
         image_paths = result.get("figures", [])
     else:
-        insights = "Error: The result is not in the expected format."
+        insights = "Error: Unexpected result format from agent."
         image_paths = []
 
-    return insights, image_paths
-
+    images = []
+    for path in image_paths:
+        try:
+            images.append(Image.open(path))
+        except Exception as e:
+            print(f"Error loading image {path}: {e}")
+    return insights, images
 
+    
 def train_model(_):
     wandb.login(key=os.environ.get("WANDB_API_KEY"))
     run_counter = 1
@@ -247,12 +262,6 @@ with gr.Blocks() as demo:
         shap_img = gr.Image(label="SHAP Summary Plot")
         lime_img = gr.Image(label="LIME Explanation")
 
-    with gr.Row():
-        agent_btn = gr.Button("Run AI Agent (5 Insights + 5 Visualizations)")
-        insights_output = gr.Textbox(label="Insights from SmolAgent", lines=15)
-        #visual_output = gr.Gallery(label="Generated Visualizations").style(grid=3, height="auto")
-        visual_output = gr.Gallery(label="Generated Visualizations", columns=[3], height=400)
-
 
     #agent_btn.click(fn=run_agent, inputs=df_output, outputs=insights_output)
     agent_btn.click(fn=run_agent, inputs=df_output, outputs=[insights_output, visual_output])