claude refactoring the evaluation visualization

src/know_lang_bot/evaluation/chatbot_evaluation_visualize.py

@@ -3,21 +3,33 @@ import json
 import pandas as pd
 from rich.console import Console
 from rich.table import Table
-from typing import List
+from typing import List, Dict
+from enum import Enum
 from know_lang_bot.evaluation.chatbot_evaluation import EvalSummary
 
+class RetrievalMethod(str, Enum):
+    EMBEDDING = "embedding"
+    EMBEDDING_RERANKING = "embedding_reranking"
+    EMBEDDING_WITH_CODE = "embedding_reranking_with_code"
+
 class ResultAnalyzer:
-    def __init__(self, base_dir: Path):
+    def __init__(self, base_dir: Path, baseline_method: RetrievalMethod = RetrievalMethod.EMBEDDING):
         self.console = Console()
-        self.embedding_dir = base_dir / "embedding"
-        self.reranking_dir = base_dir / "embedding_reranking"
-        
+        self.base_dir = base_dir
+        self.baseline_method = baseline_method
+        # Map each method to its directory
+        self.method_dirs = {
+            RetrievalMethod.EMBEDDING: self.base_dir / RetrievalMethod.EMBEDDING.value,
+            RetrievalMethod.EMBEDDING_RERANKING: self.base_dir / RetrievalMethod.EMBEDDING_RERANKING.value,
+            RetrievalMethod.EMBEDDING_WITH_CODE: self.base_dir / RetrievalMethod.EMBEDDING_WITH_CODE.value
+        }
+
     def load_results(self, file_path: Path) -> List[EvalSummary]:
         """Load evaluation results from JSON file"""
         with open(file_path) as f:
             obj_list = json.load(f)
             return [EvalSummary.model_validate(obj) for obj in obj_list]
-        
+
     def create_dataframe(self, results: List[EvalSummary]) -> pd.DataFrame:
         """Convert results to pandas DataFrame with flattened metrics"""
         rows = []
@@ -29,116 +41,154 @@ class ResultAnalyzer:
                 "chunk_relevance": result.eval_response.chunk_relevance,
                 "answer_correctness": result.eval_response.answer_correctness,
                 "code_reference": result.eval_response.code_reference,
-                "weighted_total": result.eval_response.weighted_total
+                "weighted_total": result.eval_response.weighted_total,
+                "environment": getattr(result.case, 'environment', 'default')  # Added environment
             }
             rows.append(row)
         
         return pd.DataFrame(rows)
 
-    def analyze_results(self):
-        """Analyze and display results comparison"""
-        # Load all results
-        all_results = {
-            "embedding": [],
-            "reranking": []
-        }
-        
-        for file in self.embedding_dir.glob("*.json"):
-            all_results["embedding"].extend(self.load_results(file))
-        
-        for file in self.reranking_dir.glob("*.json"):
-            all_results["reranking"].extend(self.load_results(file))
-
-        # Convert to DataFrames
-        embedding_df = self.create_dataframe(all_results["embedding"])
-        reranking_df = self.create_dataframe(all_results["reranking"])
-
-        # Calculate statistics by evaluator model
-        def get_model_stats(df: pd.DataFrame) -> pd.DataFrame:
-            return df.groupby("evaluator_model").agg({
-                "chunk_relevance": ["mean", "std"],
-                "answer_correctness": ["mean", "std"],
-                "code_reference": ["mean", "std"],
-                "weighted_total": ["mean", "std"]
-            }).round(2)
-
-        embedding_stats = get_model_stats(embedding_df)
-        reranking_stats = get_model_stats(reranking_df)
-
-        # Display comparison tables
-        self.display_comparison_table(embedding_stats, reranking_stats)
-        self.display_improvement_metrics(embedding_df, reranking_df)
-        
-        # Save detailed results to CSV
-        self.save_detailed_results(embedding_df, reranking_df)
+    def load_all_results(self) -> Dict[RetrievalMethod, pd.DataFrame]:
+        """Load results for all available methods"""
+        results = {}
+        for method in RetrievalMethod:
+            method_dir = self.method_dirs.get(method)
+            if method_dir and method_dir.exists():
+                all_results = []
+                for file in method_dir.glob("*.json"):
+                    all_results.extend(self.load_results(file))
+                if all_results:  # Only include methods with results
+                    results[method] = self.create_dataframe(all_results)
+        return results
 
-    def display_comparison_table(self, embedding_stats: pd.DataFrame, reranking_stats: pd.DataFrame):
-        """Display rich table comparing embedding and reranking results"""
-        table = Table(title="Embedding vs Reranking Comparison")
+    def calculate_improvement(self, new_val: float, baseline_val: float) -> str:
+        """Calculate and format improvement percentage"""
+        if baseline_val == 0:
+            return "N/A"
+        improvement = ((new_val - baseline_val) / baseline_val * 100).round(1)
+        return f"{improvement:+.1f}%" if improvement else "0%"
+
+    def get_stats_by_group(self, df: pd.DataFrame, group_by: str) -> pd.DataFrame:
+        """Calculate statistics grouped by specified column"""
+        return df.groupby(group_by).agg({
+            "chunk_relevance": ["mean", "std"],
+            "answer_correctness": ["mean", "std"],
+            "code_reference": ["mean", "std"],
+            "weighted_total": ["mean", "std"]
+        }).round(2)
+
+    def display_comparison_table(self, results: Dict[RetrievalMethod, pd.DataFrame]):
+        """Display rich table comparing all methods"""
+        table = Table(title="Method Comparison by Evaluator Model")
         
+        # Add columns
         table.add_column("Metric", style="cyan")
         table.add_column("Model", style="magenta")
-        table.add_column("Embedding", style="blue")
-        table.add_column("Reranking", style="green")
-        table.add_column("Improvement", style="yellow")
+        for method in results.keys():
+            table.add_column(method.value.replace("_", " ").title(), style="blue")
+            if method != self.baseline_method:
+                table.add_column(f"{method.value} Improvement", style="yellow")
+
+        # Calculate stats for each method
+        stats_by_method = {
+            method: self.get_stats_by_group(df, "evaluator_model")
+            for method, df in results.items()
+        }
 
         metrics = ["chunk_relevance", "answer_correctness", "code_reference", "weighted_total"]
-        
+        baseline_stats = stats_by_method[self.baseline_method]
+
         for metric in metrics:
-            for model in embedding_stats.index:
-                emb_mean = embedding_stats.loc[model, (metric, "mean")]
-                emb_std = embedding_stats.loc[model, (metric, "std")]
-                rer_mean = reranking_stats.loc[model, (metric, "mean")]
-                rer_std = reranking_stats.loc[model, (metric, "std")]
-                
-                improvement = ((rer_mean - emb_mean) / emb_mean * 100).round(1)
-                
-                table.add_row(
+            for model in baseline_stats.index:
+                row_data = [
                     metric.replace("_", " ").title(),
-                    model.split(":")[-1],
-                    f"{emb_mean:.2f} ±{emb_std:.2f}",
-                    f"{rer_mean:.2f} ±{rer_std:.2f}",
-                    f"{improvement:+.1f}%" if improvement else "0%"
-                )
+                    model.split(":")[-1]
+                ]
+                
+                # Add data for each method
+                for method in results.keys():
+                    stats = stats_by_method[method]
+                    mean = stats.loc[model, (metric, "mean")]
+                    std = stats.loc[model, (metric, "std")]
+                    row_data.append(f"{mean:.2f} ±{std:.2f}")
+                    
+                    # Add improvement column if not baseline
+                    if method != self.baseline_method:
+                        baseline_mean = baseline_stats.loc[model, (metric, "mean")]
+                        row_data.append(self.calculate_improvement(mean, baseline_mean))
+
+                table.add_row(*row_data)
 
         self.console.print(table)
 
-    def display_improvement_metrics(self, embedding_df: pd.DataFrame, reranking_df: pd.DataFrame):
-        """Display additional improvement metrics"""
-        # Calculate improvements by difficulty
-        difficulties = sorted(embedding_df["difficulty"].unique())
+    def display_environment_comparison(self, results: Dict[RetrievalMethod, pd.DataFrame]):
+        """Display comparison across different evaluation environments"""
+        table = Table(title="Method Comparison by Environment")
+        
+        table.add_column("Environment", style="cyan")
+        for method in results.keys():
+            table.add_column(method.value.replace("_", " ").title(), style="blue")
+            if method != self.baseline_method:
+                table.add_column(f"{method.value} Improvement", style="yellow")
+
+        # Get environments from all results
+        environments = sorted(set().union(*[
+            set(df["environment"].unique()) 
+            for df in results.values()
+        ]))
+
+        baseline_df = results[self.baseline_method]
         
-        table = Table(title="Improvements by Difficulty")
-        table.add_column("Difficulty", style="cyan")
-        table.add_column("Embedding", style="blue")
-        table.add_column("Reranking", style="green")
-        table.add_column("Improvement", style="yellow")
-
-        for diff in difficulties:
-            emb_score = embedding_df[embedding_df["difficulty"] == diff]["weighted_total"].mean()
-            rer_score = reranking_df[reranking_df["difficulty"] == diff]["weighted_total"].mean()
-            improvement = ((rer_score - emb_score) / emb_score * 100).round(1)
+        for env in environments:
+            row_data = [env]
+            
+            for method in results.keys():
+                df = results[method]
+                env_score = df[df["environment"] == env]["weighted_total"].mean()
+                row_data.append(f"{env_score:.2f}")
+                
+                if method != self.baseline_method:
+                    baseline_score = baseline_df[
+                        baseline_df["environment"] == env
+                    ]["weighted_total"].mean()
+                    row_data.append(self.calculate_improvement(env_score, baseline_score))
             
-            table.add_row(
-                str(diff),
-                f"{emb_score:.2f}",
-                f"{rer_score:.2f}",
-                f"{improvement:+.1f}%"
-            )
+            table.add_row(*row_data)
 
         self.console.print(table)
 
-    def save_detailed_results(self, embedding_df: pd.DataFrame, reranking_df: pd.DataFrame):
+    def analyze_results(self):
+        """Analyze and display results comparison"""
+        results = self.load_all_results()
+        if not results:
+            self.console.print("[red]No results found!")
+            return
+
+        # Display comparisons
+        self.display_comparison_table(results)
+        self.console.print("\n")
+        self.display_environment_comparison(results)
+        
+        # Save detailed results
+        self.save_detailed_results(results)
+
+    def save_detailed_results(self, results: Dict[RetrievalMethod, pd.DataFrame]):
         """Save detailed results to CSV"""
-        # Add method column
-        embedding_df["method"] = "embedding"
-        reranking_df["method"] = "reranking"
+        # Combine all results with method column
+        dfs = []
+        for method, df in results.items():
+            df = df.copy()
+            df["method"] = method.value
+            dfs.append(df)
         
-        # Combine and save
-        combined_df = pd.concat([embedding_df, reranking_df])
-        combined_df.to_csv("evaluation_comparison.csv", index=False)
-        self.console.print(f"\nDetailed results saved to evaluation_comparison.csv")
+        combined_df = pd.concat(dfs)
+        output_path = self.base_dir / "evaluation_comparison.csv"
+        combined_df.to_csv(output_path, index=False)
+        self.console.print(f"\nDetailed results saved to {output_path}")
 
 if __name__ == "__main__":
-    analyzer = ResultAnalyzer(Path("evaluations"))
+    analyzer = ResultAnalyzer(
+        Path("evaluations"), 
+        baseline_method=RetrievalMethod.EMBEDDING
+    )
     analyzer.analyze_results()