make both functions respect the "visualizations" toggle. (#1)

- make both functions respect the "visualizations" toggle. (fceae67877a85f806911c24cd5d7482b61f6c0b3)


Co-authored-by: Karim "Sins" Yehia <mewhenmonkeyavatar@users.noreply.huggingface.co>

app.py

@@ -492,41 +492,40 @@ def process_comprehensive_analysis(model_input: str, onnx_file, llms_to_query: l
             full_response += f"{interpretation}\n\n---\n\n"
 
     # Visualization Generation
-    viz_llm = llms_to_query[0]
-    print(f"Using {viz_llm} for visualization generation...")
-
     viz_outputs = {}
     viz_types = ["shap", "lime", "attention", "architecture", "parameter"]
-    
-    for viz_type in viz_types:
-        try:
-            prompt = get_visualization_prompt(viz_type)
-            if not prompt:
-                viz_outputs[viz_type] = None
-                continue
+    for i in viz_types:
+        viz_outputs[i] = None
+
+    if "visualizations" in analysis_types:
+        viz_llm = llms_to_query[0]
+        print(f"Using {viz_llm} for visualization generation...")
+        
+        for viz_type in viz_types:
+            try:
+                prompt = get_visualization_prompt(viz_type)
+                if not prompt:
+                    continue
+                    
+                print(f"Generating {viz_type} visualization...")
+                generated_code_response = get_llm_response(prompt, model_description, viz_llm)
                 
-            print(f"Generating {viz_type} visualization...")
-            generated_code_response = get_llm_response(prompt, model_description, viz_llm)
-            
-            if "Error" in generated_code_response:
-                print(f"LLM Error for {viz_type}: {generated_code_response}")
-                viz_outputs[viz_type] = None
-                continue
+                if "Error" in generated_code_response:
+                    print(f"LLM Error for {viz_type}: {generated_code_response}")
+                    continue
+                    
+                cleaned_code = extract_python_code(generated_code_response)
                 
-            cleaned_code = extract_python_code(generated_code_response)
-            
-            if not cleaned_code:
-                print(f"Could not extract valid code for {viz_type}.")
-                viz_outputs[viz_type] = None
-                continue
+                if not cleaned_code:
+                    print(f"Could not extract valid code for {viz_type}.")
+                    continue
+                    
+                # Execute code, save plot, and get the file path
+                image_path = safely_execute_visualization_code(cleaned_code, viz_type, run_dir)
+                viz_outputs[viz_type] = image_path
                 
-            # Execute code, save plot, and get the file path
-            image_path = safely_execute_visualization_code(cleaned_code, viz_type, run_dir)
-            viz_outputs[viz_type] = image_path
-            
-        except Exception as e:
-            print(f"Unexpected error in main loop for {viz_type}: {e}")
-            viz_outputs[viz_type] = None
+            except Exception as e:
+                print(f"Unexpected error in main loop for {viz_type}: {e}")
 
     return (
         full_response,
@@ -579,41 +578,40 @@ def process_text_comprehensive_analysis(model_input: str, llms_to_query, analysi
             full_response += f"{interpretation}\n\n---\n\n"
 
     # Visualization Generation
-    viz_llm = llms_to_query[0]
-    print(f"Using {viz_llm} for visualization generation...")
-
     viz_outputs = {}
     viz_types = ["shap", "lime", "attention", "architecture", "parameter"]
+    for i in viz_types:
+        viz_outputs[i] = None
 
-    for viz_type in viz_types:
-        try:
-            prompt = get_visualization_prompt(viz_type)
-            if not prompt:
-                viz_outputs[viz_type] = None
-                continue
+    if "visualizations" in analysis_types:
+        viz_llm = llms_to_query[0]
+        print(f"Using {viz_llm} for visualization generation...")
+        
+        for viz_type in viz_types:
+            try:
+                prompt = get_visualization_prompt(viz_type)
+                if not prompt:
+                    continue
 
-            print(f"Generating {viz_type} visualization...")
-            generated_code_response = get_llm_response(prompt, model_description, viz_llm)
+                print(f"Generating {viz_type} visualization...")
+                generated_code_response = get_llm_response(prompt, model_description, viz_llm)
 
-            if "Error" in generated_code_response:
-                print(f"LLM Error for {viz_type}: {generated_code_response}")
-                viz_outputs[viz_type] = None
-                continue
+                if "Error" in generated_code_response:
+                    print(f"LLM Error for {viz_type}: {generated_code_response}")
+                    continue
 
-            cleaned_code = extract_python_code(generated_code_response)
+                cleaned_code = extract_python_code(generated_code_response)
 
-            if not cleaned_code:
-                print(f"Could not extract valid code for {viz_type}.")
-                viz_outputs[viz_type] = None
-                continue
+                if not cleaned_code:
+                    print(f"Could not extract valid code for {viz_type}.")
+                    continue
 
-            # Execute code, save plot, and get the file path
-            image_path = safely_execute_visualization_code(cleaned_code, viz_type, run_dir)
-            viz_outputs[viz_type] = image_path
+                # Execute code, save plot, and get the file path
+                image_path = safely_execute_visualization_code(cleaned_code, viz_type, run_dir)
+                viz_outputs[viz_type] = image_path
 
-        except Exception as e:
-            print(f"Unexpected error in main loop for {viz_type}: {e}")
-            viz_outputs[viz_type] = None
+            except Exception as e:
+                print(f"Unexpected error in main loop for {viz_type}: {e}")
 
     return (
         full_response,