refactor evaluation

angle_emb/__init__.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
-from .angle import *
+from .angle import *  # NOQA
+from .evaluation import *  # NOQA
 
 
 __version__ = '0.4.8'

angle_emb/angle.py

@@ -10,15 +10,10 @@
 from typing import Any, Dict, Optional, List, Union, Tuple, Callable
 from dataclasses import dataclass
 
-import scipy
-import scipy.stats
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import bitsandbytes as bnb
-from tqdm import tqdm
-from boltons.iterutils import chunked_iter
 from datasets import Dataset
 from transformers import (
     AutoModelForCausalLM, AutoModel, AutoTokenizer,
@@ -35,6 +30,7 @@
 from peft.tuners.lora import LoraLayer
 
 from .utils import logger
+from .evaluation import CorrelationEvaluator
 
 
 DEFAULT_LLM_PATTERNS = [r'.*llama.*', r'.*qwen.*', r'.*baichuan.*', r'.*mistral.*']
@@ -237,44 +233,6 @@ def contrastive_with_negative_loss(
     return nn.CrossEntropyLoss()(scores, labels)
 
 
-def compute_corrcoef(x: np.ndarray, y: np.ndarray) -> float:
-    """
-    Compute correlation coefficients
-
-    :param x: np.ndarry, x array
-    :param y: np.ndarry, y array
-
-    :return: float
-    """
-    return scipy.stats.spearmanr(x, y).correlation
-
-
-def l2_normalize(arr: np.ndarray) -> np.ndarray:
-    """
-    Normalize array using L2
-
-    :param arr: np.ndarray, input array
-
-    :return: np.ndarray
-    """
-    norms = (arr**2).sum(axis=1, keepdims=True)**0.5
-    return arr / np.clip(norms, 1e-8, np.inf)
-
-
-def optimal_threshold(y_true: np.ndarray, y_pred: np.ndarray) -> Tuple[float, float]:
-    """
-    Compute optimal threshold
-
-    :param y_true: np.ndarray, y_true
-    :param y_pred: np.ndarray, y_true
-
-    :return: Tuple[float, float]
-    """
-    loss = lambda t: -np.mean((y_true > 0.5) == (y_pred > np.tanh(t)))  # NOQA
-    result = scipy.optimize.minimize(loss, 1, method='Powell')
-    return np.tanh(result.x), -result.fun
-
-
 def check_llm(model_name_or_path: str, llm_regex_patterns: List[str] = None) -> bool:
     if llm_regex_patterns is not None:
         llm_regex_patterns += DEFAULT_LLM_PATTERNS
@@ -1499,35 +1457,21 @@ def fit(self,
             trainer.push_to_hub()
         self.backbone.save_pretrained(output_dir)
 
-    def evaluate(self, data: Dataset, batch_size: int = 32, threshold: Optional[float] = None, device: Any = None):
-        self.backbone.eval()
-        data_collator = AngleDataCollator(
-            self.tokenizer,
-            return_tensors="pt",
-            max_length=self.max_length,
-            filter_duplicate=False,
-        )
-        y_trues, y_preds = [], []
-        # for X, y in data.make_iter(random=False):
-        for features in tqdm(chunked_iter(data, batch_size), desc='Evaluate'):
-            X = data_collator(features)
-            y = X.pop('labels', None)
-            y_trues.extend(y[::2, 0].detach().cpu().numpy())
-            with torch.no_grad():
-                X.to(device or self.device)
-                x_vecs = self.pooler(X,
-                                     pooling_strategy=self.pooling_strategy).detach().float().cpu().numpy()
-            x_vecs = l2_normalize(x_vecs)
-            pred = (x_vecs[::2] * x_vecs[1::2]).sum(1)
-            y_preds.extend(pred)
-
-        y_trues, y_preds = np.array(y_trues), np.array(y_preds)
-        corrcoef = compute_corrcoef(y_trues, y_preds)
-        if threshold is None:
-            _, accuracy = optimal_threshold(y_trues, y_preds)
-        else:
-            accuracy = np.mean((y_trues > 0.5) == (y_preds > threshold))
-        return corrcoef, accuracy
+    def evaluate(self, data: Dataset, batch_size: int = 32, metric: str = 'spearman_cosine') -> float:
+        """ evaluate
+
+        :param data: Dataset, DatasetFormats.A is required
+        :param batch_size: int. Default 32.
+        :param metric: str. Default 'spearman_cosine'.
+
+        :return: float.
+        """
+        return CorrelationEvaluator(
+            text1=data['text1'],
+            text2=data['text2'],
+            labels=data['label'],
+            batch_size=batch_size,
+        )(self)[metric]
 
     def encode(self,
                inputs: Union[List[str], Tuple[str], List[Dict], str],
@@ -1656,7 +1600,7 @@ def __init__(self,
         self.hub_private_repo = hub_private_repo
 
     def on_epoch_end(self, args, state, control, **kwargs):
-        corrcoef, accuracy = self.evaluate_fn(self.valid_ds)
+        corrcoef = self.evaluate_fn(self.valid_ds)
         if corrcoef > self.best_corrcoef:
             self.best_corrcoef = corrcoef
             print('new best corrcoef!')
@@ -1669,4 +1613,4 @@ def on_epoch_end(self, args, state, control, **kwargs):
                         private=self.hub_private_repo,
                         exist_ok=True,
                         commit_message='new best checkpoint')
-        print(f'corrcoef: {corrcoef}, accuracy: {accuracy}, best corrcoef: {self.best_corrcoef}')
+        logger.info(f'corrcoef: {corrcoef}, best corrcoef: {self.best_corrcoef}')

angle_emb/evaluation.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+
+from typing import List
+
+import numpy as np
+from boltons.iterutils import chunked_iter
+from tqdm import tqdm
+from sklearn.metrics.pairwise import (
+    paired_cosine_distances,
+    paired_euclidean_distances,
+    paired_manhattan_distances
+)
+from scipy.stats import pearsonr, spearmanr
+
+from .angle import AnglE
+
+
+class CorrelationEvaluator(object):
+    def __init__(
+        self,
+        text1: List[str],
+        text2: List[str],
+        labels: List[float],
+        batch_size: int = 32
+    ):
+        assert len(text1) == len(text2) == len(labels), "text1, text2, and labels must have the same length"
+        self.text1 = text1
+        self.text2 = text2
+        self.labels = labels
+        self.batch_size = batch_size
+
+    def __call__(self, model: AnglE, **kwargs) -> dict:
+        """ Evaluate the model on the given dataset.
+
+        :param model: AnglE, the model to evaluate.
+        :param kwargs: Additional keyword arguments to pass to the `encode` method of the model.
+
+        :return: dict, The evaluation results.
+        """
+        embeddings1 = []
+        embeddings2 = []
+        for chunk in tqdm(chunked_iter(range(len(self.text1)), self.batch_size)):
+            batch_text1 = [self.text1[i] for i in chunk]
+            batch_text2 = [self.text2[i] for i in chunk]
+
+            batch_embeddings1 = model.encode(batch_text1, **kwargs)
+            batch_embeddings2 = model.encode(batch_text2, **kwargs)
+            embeddings1.append(batch_embeddings1)
+            embeddings2.append(batch_embeddings2)
+
+        embeddings1 = np.concatenate(embeddings1, axis=0)
+        embeddings2 = np.concatenate(embeddings2, axis=0)
+
+        cosine_labels = 1 - (paired_cosine_distances(embeddings1, embeddings2))
+        manhattan_distances = -paired_manhattan_distances(embeddings1, embeddings2)
+        euclidean_distances = -paired_euclidean_distances(embeddings1, embeddings2)
+        dot_products = [np.dot(emb1, emb2) for emb1, emb2 in zip(embeddings1, embeddings2)]
+
+        pearson_cosine, _ = pearsonr(self.labels, cosine_labels)
+        spearman_cosine, _ = spearmanr(self.labels, cosine_labels)
+
+        pearson_manhattan, _ = pearsonr(self.labels, manhattan_distances)
+        spearman_manhattan, _ = spearmanr(self.labels, manhattan_distances)
+
+        pearson_euclidean, _ = pearsonr(self.labels, euclidean_distances)
+        spearman_euclidean, _ = spearmanr(self.labels, euclidean_distances)
+
+        pearson_dot, _ = pearsonr(self.labels, dot_products)
+        spearman_dot, _ = spearmanr(self.labels, dot_products)
+
+        metrics = {
+            "pearson_cosine": pearson_cosine,
+            "spearman_cosine": spearman_cosine,
+            "pearson_manhattan": pearson_manhattan,
+            "spearman_manhattan": spearman_manhattan,
+            "pearson_euclidean": pearson_euclidean,
+            "spearman_euclidean": spearman_euclidean,
+            "pearson_dot": pearson_dot,
+            "spearman_dot": spearman_dot,
+        }
+        return metrics
+
+    def list_all_metrics(self) -> List[str]:
+        """ Get a list of all the metrics that can be computed by this evaluator.
+
+        :return: List[str], A list of all the metrics that can be computed by this evaluator.
+        """
+        return [
+            "pearson_cosine",
+            "spearman_cosine",
+            "pearson_manhattan",
+            "spearman_manhattan",
+            "pearson_euclidean",
+            "spearman_euclidean",
+            "pearson_dot",
+            "spearman_dot",
+        ]

requirements.txt

@@ -6,4 +6,5 @@ prettytable
 transformers>=4.32.1
 scipy
 einops
-wandb
+wandb
+scikit-learn

tests/test_eval.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+
+
+def test_eval():
+    from datasets import load_dataset
+    from angle_emb import AnglE, CorrelationEvaluator
+
+    angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls')
+    eval_dataset = load_dataset('sentence-transformers/stsb', split="test")
+
+    spearman = CorrelationEvaluator(
+        text1=eval_dataset["sentence1"],
+        text2=eval_dataset["sentence2"],
+        labels=eval_dataset["score"],
+    )(angle)['spearman_cosine']
+    assert spearman > 0.9
+
+    spearman = angle.evaluate(eval_dataset)
+    assert spearman > 0.9