fix: batching in Metric

belar/metrics/__init__.py

@@ -1,3 +1,4 @@
 from belar.metrics.base import Evaluation, Metric
+from belar.metrics.similarity import *
 from belar.metrics.simple import *
 from belar.metrics.similarity import SBERTScore

belar/metrics/base.py

@@ -23,24 +23,21 @@ def is_batchable(self) -> bool:
     def score(self, ground_truth, generated_text) -> float | list[float]:
         ...
 
-    def __call__(self, row):
-        score = self.score(row["ground_truth"], row["generated_text"])
-        row[f"{self.name}_score"] = score
-
-        return row
-
 
 @dataclass
 class Evaluation:
     metrics: list[Metric]
+    batched: bool = False
 
-    def eval(
-        self, ground_truth: Dataset, generated_text: t.Sequence, batched: bool = False
-    ):
+    def eval(self, ground_truth: list[list[str]], generated_text: list[list[str]]):
         ds = ground_truth.add_column("generated_text", generated_text)
-        scores_list = []
+        ds = ds.map(self._get_score, batched=self.batched)
+
+        return ds
+
+    def _get_score(self, row):
         for metric in self.metrics:
-            scores = ds.map(metric, batched=batched)[f"{metric.name}_score"]
-            scores_list.append(scores)
+            score = metric.score(row["ground_truth"], row["generated_text"])
+            row[f"{metric.name}_score"] = score
 
-        return scores_list
+        return row

belar/metrics/similarity.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-from ast import List
 
 import typing as t
 from dataclasses import dataclass
@@ -13,41 +12,51 @@
 
 @dataclass
 class SBERTScore(Metric):
-
     similarity_metric: t.Literal[SBERT_METRIC] = "cosine"
     model_path: str = "all-MiniLM-L6-v2"
     batch_size: int = 1000
 
     def __post_init__(self):
-
         self.model = SentenceTransformer(self.model_path)
 
-    def name(self,):
-        return f"SBERT-{self.similarity_metric}-score"
+    @property
+    def name(
+        self,
+    ):
+        return f"SBERT_{self.similarity_metric}"
 
     def is_batchable(self):
         return True
-
-    def score(self, ground_truth: t.Union[str, t.List[str]], generated_text: t.Union[str, t.List[str]]):
 
+    def score(
+        self,
+        ground_truth: str | list[str],
+        generated_text: str | list[str],
+    ):
         if isinstance(ground_truth, str):
             ground_truth = [ground_truth]
         if isinstance(generated_text, str):
             generated_text = [generated_text]
-
-        gndtruth_emb = self.model.encode(ground_truth, batch_size=self.batch_size,
-                                         convert_to_numpy=True)
-        gentext_emb = self.model.encode(generated_text, batch_size=self.batch_size,
-                                        convert_to_numpy=True)
-
+
+        gndtruth_emb = self.model.encode(
+            ground_truth, batch_size=self.batch_size, convert_to_numpy=True
+        )
+        gentext_emb = self.model.encode(
+            generated_text, batch_size=self.batch_size, convert_to_numpy=True
+        )
+
         if self.similarity_metric == "cosine":
-            score = np.dot(gndtruth_emb, gentext_emb.T) / (norm(gndtruth_emb) * norm(gentext_emb))
+            score = np.dot(gndtruth_emb, gentext_emb.T) / (
+                norm(gndtruth_emb) * norm(gentext_emb)
+            )
 
         elif self.similarity_metric == "euclidean":
             score = norm(gndtruth_emb - gentext_emb, ord=2)
-        
+
         else:
             raise ValueError(f"Unkown metrics {self.similarity_metric}")
 
         return score
-
+
+
+__all__ = ["SBERTScore"]

belar/metrics/simple.py

@@ -20,9 +20,11 @@ def __post_init__(self):
             [self.type], use_stemmer=self.use_stemmer
         )
 
+    @property
     def name(self):
         return self.type
 
+    @property
     def is_batchable(self):
         return False