Update app.py

app.py

@@ -1193,22 +1193,44 @@ def plot_compression_tradeoff(summaries_by_ratio: Dict[float, Dict[str, Any]],
     logger.info(f"Compression trade-off plots saved: {plot_path}")
     return plot_path
 
-def plot_ablation_results(ablation_results: Dict[str, Dict[str, Any]], baseline_summary: Dict[str, Any]) -> str:
+def plot_ablation_results(ablation_results: Dict[str, Dict[str, Any]], baseline_summary: Dict[str, Any] = None) -> str:
     """Generate publication-grade ablation study plots."""
+    # Find baseline summary if not provided
+    if baseline_summary is None:
+        if "baseline" in ablation_results and "summary" in ablation_results["baseline"]:
+            baseline_summary = ablation_results["baseline"]["summary"]
+        else:
+            # Find any NONE compression type as baseline
+            for config_name, result in ablation_results.items():
+                if "summary" in result:
+                    baseline_summary = result["summary"]
+                    logger.warning(f"Using {config_name} as baseline for ablation plots")
+                    break
+    
+    if baseline_summary is None:
+        logger.error("No valid baseline found for ablation plots")
+        return None
+    
     fig, axes = plt.subplots(2, 3, figsize=(18, 10))
     
+    # Filter out configs with errors
+    valid_configs = {k: v for k, v in ablation_results.items() if "summary" in v}
+    if not valid_configs:
+        logger.error("No valid ablation results to plot")
+        return None
+    
     # Prepare data
-    configs = list(ablation_results.keys())
-    compression_ratios = [ablation_results[c]['summary']['compression_ratio'] for c in configs]
-    gen_ppls = [ablation_results[c]['summary']['generation_perplexity'] for c in configs]
-    decode_times = [ablation_results[c]['summary']['decode_time_ms'] for c in configs]
-    kv_memories = [ablation_results[c]['summary']['kv_cache_memory_mb'] for c in configs]
-    throughputs = [ablation_results[c]['summary'].get('end_to_end_throughput', 0) for c in configs]
+    configs = list(valid_configs.keys())
+    compression_ratios = [valid_configs[c]['summary']['compression_ratio'] for c in configs]
+    gen_ppls = [valid_configs[c]['summary']['generation_perplexity'] for c in configs]
+    decode_times = [valid_configs[c]['summary']['decode_time_ms'] for c in configs]
+    kv_memories = [valid_configs[c]['summary']['kv_cache_memory_mb'] for c in configs]
+    throughputs = [valid_configs[c]['summary'].get('end_to_end_throughput', 0) for c in configs]
     
     baseline_gen_ppl = baseline_summary['generation_perplexity']
     baseline_decode_time = baseline_summary['decode_time_ms']
     baseline_kv_memory = baseline_summary['kv_cache_memory_mb']
-    baseline_throughput = baseline_summary.get('end_to_end_throughput', 0)
+    baseline_throughput = baseline_summary.get('end_to_end_throughput', 1)
     
     # 1. Compression Ratio by Component
     ax1 = axes[0, 0]
@@ -3448,9 +3470,26 @@ def create_research_interface():
                 tradeoff_path = None
         
         # Generate ablation plots if ablation study was done
-        if enable_ablation and ablation_results and "baseline" in ablation_results:
+        if enable_ablation and ablation_results:
             try:
-                ablation_plots_path = plot_ablation_results(ablation_results, ablation_results["baseline"]["summary"])
+                # Find a valid baseline for comparison
+                baseline_for_plots = None
+                if "baseline" in ablation_results and "summary" in ablation_results["baseline"]:
+                    baseline_for_plots = ablation_results["baseline"]["summary"]
+                elif "NONE" in all_summaries:
+                    baseline_for_plots = all_summaries["NONE"]
+                else:
+                    # Use the first available summary
+                    for result in ablation_results.values():
+                        if "summary" in result:
+                            baseline_for_plots = result["summary"]
+                            break
+                
+                if baseline_for_plots:
+                    ablation_plots_path = plot_ablation_results(ablation_results, baseline_for_plots)
+                else:
+                    logger.warning("No baseline available for ablation plots")
+                    ablation_plots_path = None
             except Exception as e:
                 logger.error(f"Failed to generate ablation plots: {e}")
                 ablation_plots_path = None