| import logging |
| from typing import List, Dict, Any |
| import matplotlib.pyplot as plt |
| from pathlib import Path |
| from evaluation.metrics import EvaluationMetrics |
|
|
| logger = logging.getLogger(__name__) |
|
|
| class RAGEvaluator: |
| """Evaluates RAG pipeline improvements and optimizes Top-K retrieval.""" |
| |
| def __init__(self, comparator): |
| self.comparator = comparator |
| self.metrics = EvaluationMetrics() |
| logger.info(f"RAGEvaluator using BERTScore/CodeBERTScore backbone: {getattr(comparator, 'bertscore_model', 'models/codebert-base')}") |
|
|
| def measure_rag_improvement(self, rag_pipeline, references, queries, lora_preds, top_k=3): |
| logger.info(f"Measuring RAG improvement over baseline LoRA (top_k={top_k})...") |
| rag_preds = [] |
| for q in queries: |
| out, _ = rag_pipeline.generate_with_rag(q, top_k=top_k) |
| rag_preds.append(out) |
| |
| bertscore_model = getattr(self.comparator, "bertscore_model", "models/codebert-base") |
|
|
| rag_metrics = { |
| "BLEU": self.metrics.compute_bleu(references, rag_preds), |
| "BERTScore": self.metrics.compute_bertscore(references, rag_preds, device=self.comparator.device, model_type=bertscore_model) |
| |
| } |
| |
| baseline_metrics = { |
| "BLEU": self.metrics.compute_bleu(references, lora_preds), |
| "BERTScore": self.metrics.compute_bertscore(references, lora_preds, device=self.comparator.device, model_type=bertscore_model) |
| |
| } |
| |
| gain = {k: round(rag_metrics[k] - baseline_metrics[k], 4) for k in rag_metrics} |
| |
| return { |
| "metrics": {"RAG_Pipeline": rag_metrics, "Baseline_LoRA": baseline_metrics}, |
| "gain_over_baseline": gain, |
| "baseline_used_for_gain": "Baseline_LoRA" |
| } |
|
|
| def sweep_top_k(self, rag_pipeline, references, queries, baseline_preds, k_values=[1, 2, 3, 5, 8], verbose=True): |
| sweep_results = {} |
| |
| |
| baseline_bleu = self.metrics.compute_bleu(references, baseline_preds) |
| baseline_rouge = self.metrics.compute_rouge(references, baseline_preds).get("ROUGE-1", 0.0) |
| |
| for k in k_values: |
| logger.info(f"Sweeping top_k={k}...") |
| preds = [] |
| for q in queries: |
| out, _ = rag_pipeline.generate_with_rag(q, top_k=k) |
| preds.append(out) |
| |
| |
| bleu = self.metrics.compute_bleu(references, preds) |
| rouge1 = self.metrics.compute_rouge(references, preds).get("ROUGE-1", 0.0) |
| |
| metrics = {"BLEU": bleu, "ROUGE-1": rouge1} |
| gain = { |
| "BLEU": round(bleu - baseline_bleu, 4), |
| "ROUGE-1": round(rouge1 - baseline_rouge, 4) |
| } |
| |
| sweep_results[k] = {"metrics": metrics, "gain_over_baseline": gain, "predictions": preds} |
|
|
| if verbose: |
| print("\n" + self.format_sweep_table(sweep_results)) |
|
|
| return sweep_results |
|
|
| @staticmethod |
| def format_sweep_table(sweep_results: Dict[int, Dict[str, Any]]) -> str: |
| ks = list(sweep_results.keys()) |
| if not ks: |
| return "(no sweep results to display)" |
|
|
| columns = [ |
| ("Top-K", lambda k, r: str(k)), |
| ("BLEU", lambda k, r: f"{r['metrics'].get('BLEU', 0.0):.4f}"), |
| ("ROUGE-1", lambda k, r: f"{r['metrics'].get('ROUGE-1', 0.0):.4f}"), |
| ("BLEU Gain", lambda k, r: f"{r['gain_over_baseline'].get('BLEU', 0.0):+.4f}"), |
| ("ROUGE-1 Gain", lambda k, r: f"{r['gain_over_baseline'].get('ROUGE-1', 0.0):+.4f}"), |
| ] |
|
|
| rows = [[render(k, sweep_results[k]) for _, render in columns] for k in ks] |
| headers = [name for name, _ in columns] |
|
|
| widths = [ |
| max(len(headers[i]), *(len(row[i]) for row in rows)) |
| for i in range(len(columns)) |
| ] |
|
|
| def fmt_row(cells): |
| return "| " + " | ".join(c.center(w) for c, w in zip(cells, widths)) + " |" |
|
|
| separator = "+-" + "-+-".join("-" * w for w in widths) + "-+" |
|
|
| lines = ["Top-K Retrieval Sweep Results", separator, fmt_row(headers), separator] |
| lines += [fmt_row(row) for row in rows] |
| lines.append(separator) |
| return "\n".join(lines) |
|
|
| def best_k(self, sweep_results, metric="BLEU"): |
| best_k_val = -999 |
| best_k = None |
| for k, res in sweep_results.items(): |
| val = res["gain_over_baseline"].get(metric, -999) |
| if val > best_k_val: |
| best_k_val = val |
| best_k = k |
| return best_k if best_k is not None else list(sweep_results.keys())[0] |
|
|
| def plot_gain_vs_k(self, sweep_results, save_path="output/plots/topk_gain_sweep.png"): |
| Path(save_path).parent.mkdir(parents=True, exist_ok=True) |
| ks = list(sweep_results.keys()) |
| bleu_gains = [sweep_results[k]["gain_over_baseline"].get("BLEU", 0) for k in ks] |
| |
| plt.figure(figsize=(8, 5)) |
| plt.plot(ks, bleu_gains, marker='o', linestyle='-', color='b') |
| plt.title("RAG Gain over Baseline by Top-K") |
| plt.xlabel("Top-K Retrieved Contexts") |
| plt.ylabel("BLEU Score Gain") |
| plt.grid(True) |
| plt.savefig(save_path) |
| plt.close() |
| logger.info(f"✅ Saved Top-K sweep plot to {save_path}") |
|
|
