Update evaluator.py (#16)

- Update evaluator.py (260dd1f6f4d58ed1d98270526903191fc129a3a4)


Co-authored-by: Manay Porwal <manayporwal07@users.noreply.huggingface.co>

evaluator.py

@@ -204,258 +204,257 @@
 #
 ###############################################################################################################################
 
+# evaluator.py
 """
-Evaluation logic for Agentic Evaluation Framework.
+Upgraded Evaluation logic for the Agentic Evaluation Framework.
+Provides scoring functions, visualization generation, and summary outputs.
 """
 
-import os
+import math
+import uuid
+from typing import List, Dict, Tuple
 import numpy as np
 import pandas as pd
-import torch
 import matplotlib.pyplot as plt
-
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSequenceClassification,
-    AutoModelForCausalLM,
-    pipeline,
-)
-from sentence_transformers import SentenceTransformer
-import evaluate
+import seaborn as sns
 
 # -----------------------------
-# Global Config
+# Lazy model loading
 # -----------------------------
-NLI_MODEL = "microsoft/deberta-v2-xlarge-mnli"
-EMBED_MODEL = "all-MiniLM-L6-v2"
-LLM_JUDGE_MODEL = "microsoft/DialoGPT-small"
-FLUENCY_MODEL = "textattack/roberta-base-CoLA"
-
-device = 0 if torch.cuda.is_available() else -1
+NLI_MODEL = "textattack/roberta-base-MNLI"
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-# Caches
-_nli_model, _nli_tokenizer = None, None
+_nli_tokenizer = None
+_nli_model = None
 _embed_model = None
-_judge_model, _judge_tokenizer = None, None
-_fluency_checker = None
-
-# Metrics
-bertscore = evaluate.load("bertscore")
-bleu = evaluate.load("bleu")
-rouge = evaluate.load("rouge")
-
+_id2label = None
 
-# -----------------------------
-# Lazy Model Loaders
-# -----------------------------
-def get_nli_model():
-    global _nli_model, _nli_tokenizer
+def ensure_models_loaded():
+    global _nli_tokenizer, _nli_model, _embed_model, _id2label
+    if _embed_model is None:
+        from sentence_transformers import SentenceTransformer, util
+        _embed_model = SentenceTransformer(EMBED_MODEL)
+        globals()["util"] = util
     if _nli_model is None:
+        from transformers import AutoTokenizer, AutoModelForSequenceClassification
         _nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL)
-        _nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL).to(
-            torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        )
+        _nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL)
+        _nli_model.to("cpu")
         _nli_model.eval()
-    return _nli_model, _nli_tokenizer
-
+        _id2label = {int(k): v.upper() for k, v in _nli_model.config.id2label.items()}
 
 def get_embed_model():
-    global _embed_model
-    if _embed_model is None:
-        _embed_model = SentenceTransformer(EMBED_MODEL, device="cuda" if torch.cuda.is_available() else "cpu")
+    ensure_models_loaded()
     return _embed_model
 
-
-def get_judge_model():
-    global _judge_model, _judge_tokenizer
-    if _judge_model is None:
-        _judge_tokenizer = AutoTokenizer.from_pretrained(LLM_JUDGE_MODEL)
-        _judge_model = AutoModelForCausalLM.from_pretrained(LLM_JUDGE_MODEL).to(
-            torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        )
-    return _judge_model, _judge_tokenizer
-
-
-def get_fluency_checker():
-    global _fluency_checker
-    if _fluency_checker is None:
-        _fluency_checker = pipeline(
-            "text-classification", model=FLUENCY_MODEL, device=device
-        )
-    return _fluency_checker
-
+def get_nli_tokenizer_and_model():
+    ensure_models_loaded()
+    return _nli_tokenizer, _nli_model, _id2label
 
 # -----------------------------
-# Evaluation Functions
+# Metric functions
 # -----------------------------
-def check_instruction_following(prompt, response):
-    try:
-        nli_model, nli_tokenizer = get_nli_model()
-        inputs = nli_tokenizer(prompt, response, return_tensors="pt", truncation=True, padding=True).to(
-            nli_model.device
-        )
-        with torch.no_grad():
-            logits = nli_model(**inputs).logits
-        probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
-        entailment_score = probs[2]  # entailment index
-        return float(entailment_score)
-    except Exception:
-        return 0.0
 
-
-def check_hallucination(reference, response):
-    try:
-        nli_model, nli_tokenizer = get_nli_model()
-        inputs = nli_tokenizer(reference, response, return_tensors="pt", truncation=True, padding=True).to(
-            nli_model.device
-        )
-        with torch.no_grad():
-            logits = nli_model(**inputs).logits
-        probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
-        contradiction_score = probs[0]  # contradiction index
-        return 1.0 - float(contradiction_score)
-    except Exception:
+def check_instruction_following(prompt: str, response: str) -> float:
+    if not prompt or not response:
         return 0.0
-
-
-def check_assumption(prompt, response):
-    try:
-        judge_model, judge_tokenizer = get_judge_model()
-        input_text = f"Does this response make assumptions not in the prompt?\nPrompt: {prompt}\nResponse: {response}\nAnswer yes or no:"
-        inputs = judge_tokenizer.encode(input_text, return_tensors="pt").to(judge_model.device)
-        outputs = judge_model.generate(inputs, max_length=50)
-        judgment = judge_tokenizer.decode(outputs[0], skip_special_tokens=True).lower()
-        if "yes" in judgment:
-            return 0.0
-        elif "no" in judgment:
-            return 1.0
-        return 0.5
-    except Exception:
-        return 0.5
-
-
-def check_coherence(response):
-    try:
-        emb = get_embed_model().encode(response, convert_to_tensor=True, normalize_embeddings=True)
-        coherence = float(torch.mean(emb).cpu().item())
-        return coherence
-    except Exception:
+    embed_model = get_embed_model()
+    p_emb = embed_model.encode(prompt, convert_to_tensor=True)
+    r_emb = embed_model.encode(response, convert_to_tensor=True)
+    sim = float(util.cos_sim(p_emb, r_emb).item())
+    return round(max(0.0, min(1.0, sim)), 3)
+
+def check_hallucination(prompt: str, response: str) -> float:
+    if not prompt or not response:
         return 0.0
-
-
-def check_accuracy(reference, response):
-    try:
-        bert_results = bertscore.compute(predictions=[response], references=[reference], lang="en")
-        bert_f1 = bert_results["f1"][0]
-    except Exception:
-        bert_f1 = 0.0
-
-    try:
-        bleu_results = bleu.compute(predictions=[response], references=[[reference]])
-        bleu_score = bleu_results["bleu"]
-    except Exception:
-        bleu_score = 0.0
-
-    try:
-        rouge_results = rouge.compute(predictions=[response], references=[reference])
-        rouge_l = rouge_results["rougeL"]
-    except Exception:
-        rouge_l = 0.0
-
-    return float((bert_f1 + bleu_score + rouge_l) / 3)
-
-
-def check_relevance(prompt, response):
-    try:
-        model = get_embed_model()
-        emb1 = model.encode(prompt, convert_to_tensor=True)
-        emb2 = model.encode(response, convert_to_tensor=True)
-        cos_sim = torch.nn.functional.cosine_similarity(emb1, emb2, dim=0)
-        return float(cos_sim.item())
-    except Exception:
+    tokenizer, model, id2label = get_nli_tokenizer_and_model()
+    inputs = tokenizer.encode_plus(prompt, response, return_tensors="pt", truncation=True)
+    outputs = model(**inputs)
+    probs = outputs.logits.softmax(dim=1).detach().cpu().numpy()[0]
+    labels = [id2label[i] for i in range(len(probs))]
+    if "ENTAILMENT" in labels:
+        entailment_prob = float(probs[labels.index("ENTAILMENT")])
+    else:
+        entailment_prob = float(probs.max())
+    return round(entailment_prob, 3)
+
+def check_accuracy(reference: str, response: str) -> float:
+    if not reference or not response:
         return 0.0
-
-
-def check_fluency(response):
-    try:
-        fluency_checker = get_fluency_checker()
-        result = fluency_checker(response)[0]
-        return float(result["score"]) if result["label"] == "LABEL_1" else 1.0 - float(result["score"])
-    except Exception:
-        return 0.5
-
+    embed_model = get_embed_model()
+    ref_emb = embed_model.encode(reference, convert_to_tensor=True)
+    r_emb = embed_model.encode(response, convert_to_tensor=True)
+    sim = float(util.cos_sim(ref_emb, r_emb).item())
+    return round(max(0.0, min(1.0, sim)), 3)
+
+def check_coherence(response: str) -> float:
+    if not response or not isinstance(response, str):
+        return 0.0
+    sentences = [s.strip() for s in response.split(".") if s.strip()]
+    if not sentences:
+        return 0.0
+    lengths = [len(s.split()) for s in sentences]
+    avg_len = np.mean(lengths)
+    std = np.std(lengths)
+    score = max(0.0, min(1.0, 1.0 - (std / (avg_len + 1e-6))))
+    return round(score, 3)
+
+def check_fluency(response: str) -> float:
+    if not response or not isinstance(response, str):
+        return 0.0
+    letters = sum(ch.isalpha() for ch in response)
+    total = len(response)
+    return round(max(0.0, min(1.0, letters / max(1, total))), 3)
 
 # -----------------------------
-# Row Evaluation
+# Visualization helpers
 # -----------------------------
-def evaluate_row(row):
-    scores = {
-        "instruction_following": check_instruction_following(row["prompt"], row["response"]),
-        "hallucination": check_hallucination(row["reference"], row["response"]),
-        "assumption": check_assumption(row["prompt"], row["response"]),
-        "coherence": check_coherence(row["response"]),
-        "accuracy": check_accuracy(row["reference"], row["response"]),
-        "relevance": check_relevance(row["prompt"], row["response"]),
-        "fluency": check_fluency(row["response"]),
-    }
-    scores["final_score"] = np.mean(list(scores.values()))
-    return pd.Series(scores)
 
+def spider_net_multi(labels: List[str], rows: List[Dict], title: str, fill_alpha: float = 0.12):
+    N = len(labels)
+    angles = [n / float(N) * 2 * math.pi for n in range(N)]
+    angles += angles[:1]
+    fig = plt.figure(figsize=(6.5, 6.5))
+    ax = plt.subplot(111, polar=True)
+    ax.set_xticks(angles[:-1])
+    ax.set_xticklabels(labels, fontsize=9)
+    ax.set_ylim(0, 1)
+    ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
+    for r in rows:
+        values = r["values"]
+        values_closed = values + values[:1]
+        ax.plot(angles, values_closed, linewidth=1.5, label=r["name"])
+        ax.fill(angles, values_closed, alpha=fill_alpha)
+    ax.set_title(title, y=1.08, fontsize=12)
+    ax.legend(loc="upper right", bbox_to_anchor=(1.25, 1.1))
+    return fig
+
+def heatmap_plot(df: pd.DataFrame, metric_cols: List[str], title: str = "Metric Correlations"):
+    fig, ax = plt.subplots(figsize=(7, 5))
+    sns.heatmap(df[metric_cols].corr(), annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
+    ax.set_title(title)
+    return fig
+
+def task_agent_heatmap(leaderboard: pd.DataFrame, metric: str):
+    pivot = leaderboard.pivot(index="task", columns="agent", values=metric)
+    fig, ax = plt.subplots(figsize=(7, 5))
+    sns.heatmap(pivot, annot=True, fmt=".2f", cmap="YlGnBu", ax=ax)
+    ax.set_title(f"Task-Agent Performance ({metric})")
+    return fig
+
+def leaderboard_barplot(leaderboard: pd.DataFrame, metric_cols: List[str]):
+    melted = leaderboard.melt(id_vars=["agent"], value_vars=metric_cols, var_name="metric", value_name="score")
+    fig, ax = plt.subplots(figsize=(8, 5))
+    sns.barplot(x="metric", y="score", hue="agent", data=melted, ax=ax)
+    ax.set_title("Leaderboard Bar Chart")
+    ax.set_ylim(0, 1)
+    return fig
+
+def distribution_plot(metrics_df: pd.DataFrame, metric: str):
+    fig, ax = plt.subplots(figsize=(7, 5))
+    sns.boxplot(x="agent", y=metric, data=metrics_df, ax=ax)
+    sns.stripplot(x="agent", y=metric, data=metrics_df, ax=ax, color="black", alpha=0.4, jitter=True)
+    ax.set_title(f"Distribution of {metric} Scores per Agent")
+    ax.set_ylim(0, 1)
+    return fig
+
+def scatter_two_metrics(metrics_df: pd.DataFrame, metric_x: str, metric_y: str):
+    fig, ax = plt.subplots(figsize=(6, 6))
+    sns.scatterplot(x=metric_x, y=metric_y, hue="agent", data=metrics_df, ax=ax, alpha=0.7)
+    ax.set_title(f"{metric_x} vs {metric_y}")
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    return fig
 
 # -----------------------------
-# Visualization Helpers
+# Main evaluation entrypoint
 # -----------------------------
-def plot_radar_chart(metrics_df, out_path="/tmp/radar.png"):
-    import seaborn as sns
-
-    mean_scores = metrics_df.mean(numeric_only=True).drop("final_score", errors="ignore")
-    categories = list(mean_scores.index)
-    values = mean_scores.values.tolist()
 
-    values += values[:1]
-    categories += categories[:1]
+def evaluate_dataframe(df: pd.DataFrame) -> Tuple[pd.DataFrame, List[Tuple[str,str]], pd.DataFrame]:
+    df = df.copy()
 
-    angles = np.linspace(0, 2 * np.pi, len(categories), endpoint=False).tolist()
-    angles += angles[:1]
-
-    plt.figure(figsize=(6, 6))
-    ax = plt.subplot(111, polar=True)
-    ax.plot(angles, values, "o-", linewidth=2)
-    ax.fill(angles, values, alpha=0.25)
-    ax.set_thetagrids(np.degrees(angles[:-1]), categories)
-    plt.savefig(out_path)
-    plt.close()
-    return out_path, "Radar Chart (Mean Scores)"
-
-
-def plot_leaderboard(metrics_df, out_path="/tmp/leaderboard.png"):
-    agent_means = metrics_df.groupby("agent")["final_score"].mean().sort_values(ascending=False)
-    plt.figure(figsize=(10, 5))
-    agent_means.plot(kind="bar", colormap="Set3", ax=plt.gca())
-    plt.title("Leaderboard: Avg Final Score per Agent")
-    plt.ylabel("Score")
-    plt.tight_layout()
-    plt.savefig(out_path)
-    plt.close()
-    return out_path, "Leaderboard"
+    # compute scores per row
+    scores = []
+    for _, row in df.iterrows():
+        s = {}
+        s["instruction_following"] = check_instruction_following(str(row.get("prompt", "")), str(row.get("response", "")))
+        s["hallucination"] = check_hallucination(str(row.get("prompt", "")), str(row.get("response", "")))
+        s["accuracy"] = check_accuracy(str(row.get("reference", "")), str(row.get("response", "")))
+        s["coherence"] = check_coherence(str(row.get("response", "")))
+        s["fluency"] = check_fluency(str(row.get("response", "")))
+        scores.append(s)
 
+    metrics_df = pd.concat([df.reset_index(drop=True), pd.DataFrame(scores)], axis=1)
 
-# -----------------------------
-# Main Evaluation Entry
-# -----------------------------
-def evaluate_dataframe(df: pd.DataFrame):
-    metrics_df = df.apply(evaluate_row, axis=1, result_type="expand")
-    metrics_df = pd.concat([df, metrics_df], axis=1)
-
+    # leaderboard: average per agent & task
+    metric_cols = ["instruction_following", "hallucination", "accuracy", "coherence", "fluency"]
     leaderboard = (
-        metrics_df.groupby("agent")["final_score"]
+        metrics_df.groupby(["agent", "task"])[metric_cols]
         .mean()
         .reset_index()
-        .sort_values("final_score", ascending=False)
     )
 
+    # -------------------
+    # Visualization images
+    # -------------------
     images = []
-    images.append(plot_radar_chart(metrics_df))
-    images.append(plot_leaderboard(metrics_df))
 
-    return metrics_df, images, leaderboard
+    try:
+        rows = []
+        for agent in leaderboard["agent"].unique():
+            vals = leaderboard[leaderboard["agent"] == agent][metric_cols].mean().tolist()
+            rows.append({"name": agent, "values": vals})
+        fig1 = spider_net_multi(metric_cols, rows, "Agent Performance Radar")
+        path1 = f"/tmp/radar_{uuid.uuid4().hex}.png"
+        fig1.savefig(path1, bbox_inches="tight")
+        plt.close(fig1)
+        images.append((path1, "Radar Plot"))
+    except Exception as e:
+        print("Radar plot failed:", e)
+
+    try:
+        fig2 = heatmap_plot(metrics_df, metric_cols, title="Metric Correlation Heatmap")
+        path2 = f"/tmp/heatmap_{uuid.uuid4().hex}.png"
+        fig2.savefig(path2, bbox_inches="tight")
+        plt.close(fig2)
+        images.append((path2, "Correlation Heatmap"))
+    except Exception as e:
+        print("Heatmap failed:", e)
+
+    try:
+        fig3 = task_agent_heatmap(leaderboard, "accuracy")
+        path3 = f"/tmp/task_agent_{uuid.uuid4().hex}.png"
+        fig3.savefig(path3, bbox_inches="tight")
+        plt.close(fig3)
+        images.append((path3, "Task-Agent Heatmap (Accuracy)"))
+    except Exception as e:
+        print("Task-agent heatmap failed:", e)
+
+    try:
+        fig4 = leaderboard_barplot(leaderboard, metric_cols)
+        path4 = f"/tmp/barplot_{uuid.uuid4().hex}.png"
+        fig4.savefig(path4, bbox_inches="tight")
+        plt.close(fig4)
+        images.append((path4, "Leaderboard Bar Chart"))
+    except Exception as e:
+        print("Barplot failed:", e)
+
+    try:
+        fig5 = distribution_plot(metrics_df, "accuracy")
+        path5 = f"/tmp/distribution_{uuid.uuid4().hex}.png"
+        fig5.savefig(path5, bbox_inches="tight")
+        plt.close(fig5)
+        images.append((path5, "Accuracy Distribution"))
+    except Exception as e:
+        print("Distribution plot failed:", e)
+
+    try:
+        fig6 = scatter_two_metrics(metrics_df, "instruction_following", "accuracy")
+        path6 = f"/tmp/scatter_{uuid.uuid4().hex}.png"
+        fig6.savefig(path6, bbox_inches="tight")
+        plt.close(fig6)
+        images.append((path6, "Instruction Following vs Accuracy"))
+    except Exception as e:
+        print("Scatter plot failed:", e)
+
+    return metrics_df, images, leaderboard