Update Space (evaluate main: e4a27243)

frugalscore.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 """FrugalScore metric."""
 
+from dataclasses import dataclass
+from typing import Optional
+
 import datasets
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
@@ -54,13 +57,28 @@ Examples:
 """
 
 
+@dataclass
+class FRUGALSCOREConfig(evaluate.info.Config):
+
+    name: str = "default"
+
+    batch_size: int = 32
+    max_length: int = 128
+    device: Optional[str] = None
+
+
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class FRUGALSCORE(evaluate.Metric):
-    def _info(self):
+
+    CONFIG_CLASS = FRUGALSCOREConfig
+    ALLOWED_CONFIG_NAMES = ["default"]
+
+    def _info(self, config):
         return evaluate.MetricInfo(
             description=_DESCRIPTION,
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
+            config=config,
             features=datasets.Features(
                 {
                     "predictions": datasets.Value("string"),
@@ -78,26 +96,20 @@ class FRUGALSCORE(evaluate.Metric):
         self.model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
         self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 
-    def _compute(
-        self,
-        predictions,
-        references,
-        batch_size=32,
-        max_length=128,
-        device=None,
-    ):
+    def _compute(self, predictions, references):
         """Returns the scores"""
         assert len(predictions) == len(
             references
         ), "predictions and references should have the same number of sentences."
-        if device is not None:
-            assert device in ["gpu", "cpu"], "device should be either gpu or cpu."
+        if self.config.device is not None:
+            assert self.config.device in ["gpu", "cpu"], "device should be either gpu or cpu."
+            device = self.config.device
         else:
             device = "gpu" if torch.cuda.is_available() else "cpu"
         training_args = TrainingArguments(
             "trainer",
             fp16=(device == "gpu"),
-            per_device_eval_batch_size=batch_size,
+            per_device_eval_batch_size=self.config.batch_size,
             report_to="all",
             no_cuda=(device == "cpu"),
             log_level="warning",
@@ -107,7 +119,7 @@ class FRUGALSCORE(evaluate.Metric):
 
         def tokenize_function(data):
             return self.tokenizer(
-                data["sentence1"], data["sentence2"], max_length=max_length, truncation=True, padding=True
+                data["sentence1"], data["sentence2"], max_length=self.config.max_length, truncation=True, padding=True
             )
 
         tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

requirements.txt

@@ -1,3 +1,3 @@
-git+https://github.com/huggingface/evaluate@80448674f5447a9682afe051db243c4a13bfe4ff
+git+https://github.com/huggingface/evaluate@e4a2724377909fe2aeb4357e3971e5a569673b39
 torch
 transformers