tau-eval 0.1.0__py3-none-any.whl

tau_eval/__init__.py

@@ -0,0 +1,18 @@
+"""Top-level package for tau-eval."""
+
+from . import (
+    logger,
+    models,
+    utils,
+)
+from .experiment import Experiment, ExperimentConfig
+
+
+__all__ = [
+    "Experiment",
+    "ExperimentConfig",
+    "models",
+    "utils",
+    "logger",
+]
+__version__ = "0.1.0"

tau_eval/experiment.py

@@ -0,0 +1,231 @@
+import json
+import os
+from dataclasses import dataclass, field
+
+import pandas as pd
+from rich.table import Table
+from tasknet import Task
+
+from tau_eval.tasks.customtask import CustomTask
+
+from .logger import logger
+from .models import Anonymizer
+from .utils import run_models_on_custom_task, run_models_on_task
+
+
+def rich_display_dataframe(df, title="Dataframe") -> None:
+    """Display dataframe as table using rich library.
+    Args:
+        df (pd.DataFrame): dataframe to display
+        title (str, optional): title of the table. Defaults to "Dataframe".
+    Raises:
+        NotRenderableError: if dataframe cannot be rendered
+    Returns:
+        rich.table.Table: rich table
+    """
+    from rich import print
+
+    # ensure dataframe contains only string values
+    df = df.astype(str)
+
+    table = Table(title=title)
+    for col in df.columns:
+        table.add_column(col)
+    for row in df.values:
+        table.add_row(*row)
+    print(table)
+
+@dataclass
+class ExperimentConfig:
+    """Evaluation experiment config"""
+
+    exp_name: str = os.path.basename(__file__)[: -len(".py")]
+    classifier_name: str = "answerdotai/ModernBERT-base"
+    train_task_models: bool = False
+    train_with_generations: bool = False
+    device: str | None = "cuda"
+    classifier_args: dict = field(default_factory=dict)
+
+
+class Experiment:
+    def __init__(
+        self,
+        models: list[Anonymizer],
+        metrics: list[str],
+        tasks: list[Task],
+        config: ExperimentConfig = ExperimentConfig(),
+    ):
+        self.models = models
+        self.metrics = metrics
+        self.tasks = tasks
+        self.args = config
+
+    def run(self, output_dir="", device="cuda"):
+        self.output_dir = output_dir
+        logger.info("Running experiment...")
+        out = {}
+
+        for i, task in enumerate(self.tasks):
+            logger.info(f"Running task: {i}")
+            if isinstance(task, CustomTask):
+                results = run_models_on_custom_task(self.models, task, self.metrics)
+            else:
+                results = run_models_on_task(
+                    self.models,
+                    task,
+                    self.metrics,
+                    self.args.classifier_name,
+                    self.args.train_task_models,
+                    self.args.train_with_generations,
+                    device,
+                )
+            if hasattr(task, "name"):
+                out[f"{task.name}_{i}"] = results
+            else:
+                out[f"{task.__class__.__name__}_{i}"] = results
+
+        self.results = out
+        with open(self.output_dir, "w") as f:
+            json.dump(out, f)
+            logger.info("Results saved")
+
+    def summary(self, output_dir=None, to_rich=False):
+        if output_dir is not None:
+            results = json.load(open(output_dir))
+        else:
+            results = json.load(open(self.output_dir))
+
+        task_dataframes = {}
+
+        for task_name, task in results.items():
+            # Prepare columns for DataFrame
+            if isinstance(task, CustomTask):
+                columns = ["Model Name"] + task["metrics"]
+                rows = []
+
+                # Add original metrics
+                original_metrics = task.get("original_metrics", {})
+                if original_metrics:
+                    original_row = ["Original"]
+                    for m in task["metrics"]:
+                        if m == "cola":
+                            original_row.append(round(original_metrics.get("cola", 0), 4))
+                        else:
+                            original_row.append("-")
+                    rows.append(original_row)
+                    # Add rewritten metrics
+                for model_name, model in task.items():
+                    if model_name in ["original_metrics", "metrics"]:
+                        continue
+
+                    row = [model_name]
+
+                    for m in task["metrics"]:
+                        row.append(round(model.get(m, 0), 4))
+
+                    rows.append(row)
+
+            else:
+                columns = ["Model Name", "Accuracy", "F1"] + task["metrics"]
+                rows = []
+
+                # Add original metrics
+                original_metrics = task.get("original_metrics", {})
+                if original_metrics:
+                    original_row = ["Original"]
+                    original_row.append(round(original_metrics.get("test_accuracy", 0), 4))
+                    original_row.append(round(original_metrics.get("test_f1", 0), 4))
+                    for m in task["metrics"]:
+                        if m == "cola":
+                            original_row.append(round(original_metrics.get("cola", 0), 4))
+                        else:
+                            original_row.append("-")
+                    rows.append(original_row)
+
+                # Add rewritten metrics
+                for model_name, model in task.items():
+                    if model_name in ["original_metrics", "metrics"]:
+                        continue
+
+                    row = [model_name]
+                    row.append(model.get("test_accuracy", 0))
+                    row.append(model.get("test_f1", 0))
+
+                    for m in task["metrics"]:
+                        row.append(round(model.get(m, 0), 4))
+
+                    rows.append(row)
+
+            # Create DataFrame for the task
+            task_dataframe = pd.DataFrame(rows, columns=columns)
+            task_dataframes[task_name] = task_dataframe
+        if to_rich:
+            for t, value in task_dataframes.items():
+                rich_display_dataframe(value)
+
+        return task_dataframes
+
+    def plot_tradeoff(
+        self,
+        task_name1,
+        target1,
+        task_name2,
+        target2,
+        x_label="Metric 1",
+        y_label="Metric 2",
+        title="Tradeoff Between Metrics",
+        figsize=(4, 4),
+    ):
+        from itertools import cycle
+
+        import matplotlib.pyplot as plt
+
+        task1 = self.results[task_name1]
+        task2 = self.results[task_name2]
+        model_names = []
+        x_vals = []
+        y_vals = []
+
+        for model_name, model in task1.items():
+            model_names.append(model_name)
+            if model_name in ["original_metrics", "metrics"]:
+                continue
+            else:
+                x_vals.append(model.get(target1))
+
+        for model_name, model in task2.items():
+            model_names.append(model_name)
+            if model_name in ["original_metrics", "metrics"]:
+                continue
+            else:
+                y_vals.append(model.get(target2))
+
+        markers = cycle(
+            ["o", "s", "D", "^", "v", "<", ">", "p", "*", "h", "H"]
+        )  # Cycle through different marker styles
+        plt.figure(figsize=figsize)
+
+        for i, (x_val, y_val) in enumerate(zip(x_vals, y_vals)):
+            marker = next(markers)
+            plt.scatter(
+                x_val,
+                y_val,
+                label=model_names[i] if model_names else f"Model {i + 1}",
+                marker=marker,
+                s=100,
+                alpha=0.7,
+                edgecolor="k",
+            )
+            # Annotate points with model names if provided
+            if model_names:
+                plt.text(x_val, y_val, model_names[i], fontsize=9, ha="right", va="bottom")
+
+        plt.axhline(0, color="gray", linewidth=0.5, linestyle="--")
+        plt.axvline(0, color="gray", linewidth=0.5, linestyle="--")
+        plt.grid(alpha=0.3)
+
+        plt.xlabel(x_label)
+        plt.ylabel(y_label)
+        plt.title(title)
+        plt.legend(loc="best", fontsize=9)
+        plt.tight_layout()

tau_eval/logger.py

@@ -0,0 +1,28 @@
+import logging
+
+from rich.logging import RichHandler
+
+
+logger = logging.getLogger(__name__)
+
+# the handler determines where the logs go: stdout/file
+shell_handler = RichHandler()
+file_handler = logging.FileHandler("debug.log")
+
+logger.setLevel(logging.DEBUG)
+shell_handler.setLevel(logging.DEBUG)
+file_handler.setLevel(logging.DEBUG)
+
+# the formatter determines what our logs will look like
+fmt_shell = "%(message)s"
+fmt_file = "%(levelname)s %(asctime)s [%(filename)s:%(funcName)s:%(lineno)d] %(message)s"
+
+shell_formatter = logging.Formatter(fmt_shell)
+file_formatter = logging.Formatter(fmt_file)
+
+# here we hook everything together
+shell_handler.setFormatter(shell_formatter)
+file_handler.setFormatter(file_formatter)
+
+logger.addHandler(shell_handler)
+logger.addHandler(file_handler)

tau_eval/metrics/bertscore.py

@@ -0,0 +1,35 @@
+import evaluate
+
+
+# Loading evaluation metric
+bert = evaluate.load("bertscore")
+
+
+def compute_bertscore(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+    model_id: str = "distilbert-base-uncased",
+) -> dict[str, list[float]]:
+    """
+    Computes BERTScore for a list of input and output text pairs.
+
+    Args:
+        input_texts: A string or a list of input text strings.
+        output_texts: A string or a list of output text strings.
+        model_id: Bert specification, HuggingFace model to use.
+
+    Returns:
+        A dictionary containing BERTScore scores for each input-output pair.
+        The dictionary will contain keys "precision", "recall", and "f1".
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    return bert.compute(
+        predictions=output_texts,
+        references=input_texts,
+        model_type=model_id,
+    )

tau_eval/metrics/cola.py

@@ -0,0 +1,89 @@
+import numpy as np
+import torch
+from nltk.tokenize import sent_tokenize
+
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+
+def load_cola(
+    model_name: str = "textattack/roberta-base-CoLA", device: str = "cuda"
+) -> tuple[PreTrainedModel, PreTrainedTokenizer]:
+    """
+    Loads the CoLA (Corpus of Linguistic Acceptability) model and tokenizer.
+
+    Args:
+        model_name: HuggingFace model to load
+        device: The device to load the model onto ("cuda" or "cpu").
+
+    Returns:
+        A tuple containing the loaded model and tokenizer.
+    """
+    cola_tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(model_name, map_location=device)
+    cola_model: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+    return (cola_tokenizer, cola_model)
+
+
+def cola_score(
+    text: str,
+    cola_tokenizer: PreTrainedTokenizer,
+    cola_model: PreTrainedModel,
+    device: str = "cuda",
+) -> float:
+    """
+    Calculates the CoLA score for a single piece of text.
+
+    Args:
+        text: The text to score.
+        cola_tokenizer: The tokenizer for the CoLA model.
+        cola_model: The CoLA model.
+        device: The device to run the model on ("cuda" or "cpu").
+
+    Returns:
+        The CoLA score (a float between 0 and 1).
+    """
+    tokenize_input = cola_tokenizer.tokenize(text)
+    tensor_input = torch.tensor([cola_tokenizer.convert_tokens_to_ids(tokenize_input)]).to(device)
+    output = cola_model(tensor_input)
+    return output.logits.softmax(-1)[0][1].item()
+
+
+def compute_cola(
+    output_texts: str | list[str],
+    cola_tokenizer: PreTrainedTokenizer,
+    cola_model: PreTrainedModel,
+    device: str = "cuda",
+) -> dict[str, list[float]]:
+    """
+    Computes CoLA scores for a list of input texts.
+
+    Args:
+        output_texts: A list of text strings.
+        cola_tokenizer: The tokenizer for the CoLA model.
+        cola_model: The CoLA model.
+        device: The device to run the model on ("cuda" or "cpu").
+
+    Returns:
+        A dictionary containing CoLA scores for each input text.
+    """
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]
+
+    cola_ls: list[float] = []
+    for i, text in enumerate(output_texts):
+        if text == "":
+            continue
+        # if text is too big, break it up
+        sentences = sent_tokenize(text)
+        cola_score_list: list[float] = []
+
+        for sent in sentences:
+            cola_score_list.append(cola_score(sent, cola_tokenizer, cola_model, device))
+
+        cola_ls.append(np.mean(cola_score_list))
+
+    return {"cola": cola_ls}

tau_eval/metrics/luar.py

@@ -0,0 +1,54 @@
+"""Sentence transformers version of LUAR"""
+import torch
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
+
+
+def load_luar(
+    model_name: str = "gabrielloiseau/LUAR-MUD-sentence-transformers",
+    device: str = "cuda",
+) -> SentenceTransformer:
+    """
+    Loads the LUAR (Language Understanding and Reasoning) sentence transformer model.
+
+    Args:
+        model_name: SentenceTransformers model to load.
+        device: The device to load the model onto ("cuda" or "cpu").
+
+    Returns:
+        The loaded SentenceTransformer model.
+    """
+    return SentenceTransformer(model_name, device=device)
+
+
+def compute_luar(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+    sim_model: SentenceTransformer,
+) -> dict[str, list[float]]:
+    """
+    Computes LUAR scores based on cosine similarity between the embeddings of original and rewritten texts.
+
+    Args:
+        original: A string or list of original texts.
+        rewrites: A string or list of rewritten texts.
+        sim_model: The loaded SentenceTransformer model.
+
+    Returns:
+        A dictionary containing the LUAR scores for each pair of input texts.
+        The dictionary has the key "luar" with a list of float values.
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    outputs = []
+    embedding_orig: torch.Tensor = sim_model.encode(input_texts, convert_to_tensor=True, show_progress_bar=False)
+    embedding_rew: torch.Tensor = sim_model.encode(output_texts, convert_to_tensor=True, show_progress_bar=False)
+
+    for orig, new in zip(embedding_orig, embedding_rew):
+        outputs.append(cos_sim(orig, new).item())
+
+    return {"luar": outputs}

tau_eval/metrics/meteor.py

@@ -0,0 +1,46 @@
+import evaluate
+
+
+# Loading evaluation metric
+meteor = evaluate.load("meteor")
+
+
+def compute_meteor(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+    alpha: float = 0.9,
+    beta: float = 3,
+    gamma: float = 0.5,
+) -> dict[str, list[float]]:
+    """
+    Computes METEOR scores for a list of input and output text pairs.
+
+    Args:
+        input_texts: A list of input text strings.
+        output_texts: A list of output text strings.
+        alpha: Parameter for controlling relative weights of precision and recall.
+        beta: Parameter for controlling shape of penalty function.
+        gamma: Relative weight of fragmentation penalty.
+
+    Returns:
+        A dictionary containing METEOR scores for each input-output pair.
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]  # type: ignore
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]  # type: ignore
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    scores = []
+    for input_text, output_text in zip(input_texts, output_texts):
+        scores.append(
+            meteor.compute(
+                predictions=[output_text],
+                references=[input_text],
+                alpha=alpha,
+                beta=beta,
+                gamma=gamma,
+            )["meteor"]
+        )
+
+    return {"meteor": scores}

tau_eval/metrics/nli.py

@@ -0,0 +1,78 @@
+import torch
+
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+
+def load_nli(
+    model_name: str = "alisawuffles/roberta-large-wanli", device: str = "cuda"
+) -> tuple[PreTrainedTokenizer, PreTrainedModel]:
+    """
+    Loads the NLI (Natural Language Inference) model and tokenizer.
+
+    Args:
+        model_name: HuggingFace model to load
+        device: The device to load the model onto ("cuda" or "cpu").
+
+    Returns:
+        A tuple containing the loaded tokenizer and model.
+    """
+    # Download necessary models for evaluation
+    nli_tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(model_name, map_location=device)
+    nli_model: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+
+    return (nli_tokenizer, nli_model)
+
+
+def compute_nli(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+    nli_tokenizer: PreTrainedTokenizer,
+    nli_model: PreTrainedModel,
+    batch_size: int = 16,
+    device: str = "cuda",
+    max_length: int = 128,
+) -> dict[str, list[float]]:
+    """
+    Computes the probability of entailment between two texts using the NLI model.
+
+    Args:
+        input_text: The premise text.
+        output_text: The hypothesis text.
+        nli_tokenizer: The tokenizer for the NLI model.
+        nli_model: The NLI model.
+
+    Returns:
+        A dictionary containing the probability of entailment.
+        The dictionary has the key "entailment" with a float value.
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]  # type: ignore
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]  # type: ignore
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    entailment_prob = []
+
+    for start_index in range(0, len(output_texts), batch_size):
+        end_index = min(start_index + batch_size, len(output_texts))
+        input_batch = input_texts[start_index:end_index]
+        output_batch = output_texts[start_index:end_index]
+
+        x = nli_tokenizer(
+            input_batch,
+            output_batch,
+            return_tensors="pt",
+            max_length=max_length,
+            padding=True,
+            truncation=True,
+        ).to(device)
+        with torch.no_grad():
+            logits = nli_model.to(device)(**x).logits
+        entailment_prob += logits.softmax(dim=1)[:, 1].tolist()
+
+    return {"entailment": entailment_prob}

tau_eval/metrics/perplexity.py

@@ -0,0 +1,24 @@
+import evaluate
+
+
+# Loading evaluation metric
+perplexity = evaluate.load("perplexity", module_type="metric")
+
+
+def compute_perplexity(output_texts: str | list[str], model_id: str = "gpt2") -> dict[str, list[float]]:
+    """
+    Computes perplexity scores for a list of output texts.
+
+    Args:
+        output_texts: A string or list of output text strings.
+        model_id: HuggingFace model to use
+
+    Returns:
+        A dictionary containing perplexity scores for each input text.
+    """
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]
+
+    result = perplexity.compute(predictions=output_texts, model_id=model_id)
+
+    return result

tau_eval/metrics/rouge.py

@@ -0,0 +1,28 @@
+import evaluate
+
+
+# Loading evaluation metric
+rouge = evaluate.load("rouge")
+
+
+def compute_rouge(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+) -> dict[str, list[float]]:
+    """
+    Computes ROUGE scores for a list of input and output text pairs.
+
+    Args:
+        input_texts: A list of input text strings.
+        output_texts: A list of output text strings.
+
+    Returns:
+        A dictionary containing ROUGE scores for each input-output pair.
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]  # type: ignore
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]  # type: ignore
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    return rouge.compute(predictions=output_texts, references=input_texts, use_aggregator=False)

tau_eval/metrics/sbert.py

@@ -0,0 +1,52 @@
+import torch
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
+
+
+def load_sbert(
+    model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str = "cuda"
+) -> SentenceTransformer:
+    """
+    Loads the sentence similarity model.
+
+    Args:
+        model_name: SentenceTransformers model to load.
+        device: The device to load the model onto ("cuda" or "cpu").
+
+    Returns:
+        The loaded SentenceTransformer model.
+    """
+    return SentenceTransformer(model_name, device=device)
+
+
+def compute_sbert(
+    input_texts: str | list[str],
+    output_texts: str | list[str],
+    sim_model: SentenceTransformer,
+) -> dict[str, list[float]]:
+    """
+    Computes the cosine similarity between the embeddings of original and rewritten texts.
+
+    Args:
+        original: A string or a list of original texts.
+        rewrites: A string or a list of rewritten texts.
+        sim_model: The loaded SentenceTransformer model.
+
+    Returns:
+        A dictionary containing the similarity scores for each input text pair.
+        The dictionary has the key "similarity" with a list of float values.
+    """
+    if not isinstance(input_texts, list):
+        input_texts = [input_texts]
+    if not isinstance(output_texts, list):
+        output_texts = [output_texts]
+    assert len(input_texts) == len(output_texts), "inputs are different lengths"
+
+    outputs = []
+    embedding_orig: torch.Tensor = sim_model.encode(input_texts, convert_to_tensor=True, show_progress_bar=False)
+    embedding_rew: torch.Tensor = sim_model.encode(output_texts, convert_to_tensor=True, show_progress_bar=False)
+
+    for orig, new in zip(embedding_orig, embedding_rew):
+        outputs.append(cos_sim(orig, new).item())
+
+    return {"sbert": outputs}

tau_eval/models/__init__.py

@@ -0,0 +1,4 @@
+from .anonymizer import Anonymizer
+
+
+__all__ = ["Anonymizer"]