Repository: obss/jury (language: python)

Issue:
New Metric: Add BARTScore
BARTScore was published in NIPS'21. We can add it to jury with minimum requirement set.

Repo: https://github.com/neulab/BARTScore
Paper: https://arxiv.org/abs/2106.11520


Discussion / hints:

Fix patch:
diff --git a/jury/metrics/__init__.py b/jury/metrics/__init__.py
index 54b01fa..a5a5858 100644
--- a/jury/metrics/__init__.py
+++ b/jury/metrics/__init__.py
@@ -13,10 +13,10 @@
     load_metric,
 )
 from jury.metrics.accuracy import Accuracy
+from jury.metrics.bartscore import Bartscore
 from jury.metrics.bertscore import Bertscore
 from jury.metrics.bleu import Bleu
 from jury.metrics.bleurt import Bleurt
-from jury.metrics.cer import CER
 from jury.metrics.f1 import F1
 from jury.metrics.meteor import Meteor
 from jury.metrics.precision import Precision
diff --git a/jury/metrics/bartscore/__init__.py b/jury/metrics/bartscore/__init__.py
new file mode 100644
index 0000000..5e8e708
--- /dev/null
+++ b/jury/metrics/bartscore/__init__.py
@@ -0,0 +1,1 @@
+from jury.metrics.bartscore.bartscore import Bartscore
diff --git a/jury/metrics/bartscore/bartscore.py b/jury/metrics/bartscore/bartscore.py
new file mode 100644
index 0000000..1576b20
--- /dev/null
+++ b/jury/metrics/bartscore/bartscore.py
@@ -0,0 +1,8 @@
+from jury.metrics._core import MetricAlias
+from jury.metrics.bartscore.bartscore_for_language_generation import BartscoreForLanguageGeneration
+
+__main_class__ = "Bartscore"
+
+
+class Bartscore(MetricAlias):
+    _SUBCLASS = BartscoreForLanguageGeneration
diff --git a/jury/metrics/bartscore/bartscore_for_language_generation.py b/jury/metrics/bartscore/bartscore_for_language_generation.py
new file mode 100644
index 0000000..0e80a31
--- /dev/null
+++ b/jury/metrics/bartscore/bartscore_for_language_generation.py
@@ -0,0 +1,304 @@
+# coding=utf-8
+# Copyright 2021 Open Business Software Solutions, The HuggingFace Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Bartscore metric. The part of this file is adapted from metric implementations
+of datasets package. See
+https://github.com/huggingface/datasets/blob/master/metrics/ """
+from typing import Callable, Dict, List
+
+import datasets
+import numpy as np
+
+from jury.metrics import LanguageGenerationInstance
+from jury.metrics._core import MetricForLanguageGeneration
+
+_CITATION = """
+@misc{yuan2021bartscore,
+      title={BARTScore: Evaluating Generated Text as Text Generation}, 
+      author={Weizhe Yuan and Graham Neubig and Pengfei Liu},
+      year={2021},
+      eprint={2106.11520},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+"""
+
+_DESCRIPTION = """
+BARTScore formulates evaluating generated text as a text generation task from pre-trained language models. It
+operationalizes this idea using BART, an encoder-decoder based pre-trained model. BARTScore is conceptually
+simple and empirically effective.
+
+See the `README.md` file at [https://github.com/neulab/BARTScore](https://github.com/neulab/BARTScore) for more
+information.
+"""
+
+_KWARGS_DESCRIPTION = """
+BARTScore metric arguments.
+
+Construction Args:
+    model_checkpoint (str): BARTScore checkpoint. Will default to bartscore-large-cnn.
+    model_weights (str): Optional BARTScore weights, overrides the checkpoint weights.
+    device (str): On which the contextual embedding model will be allocated on.
+        If this argument is None, the model lives on cuda:0 if cuda is available.
+    
+Computation Args:
+    predictions (list of str): Prediction/candidate sentences.
+    references (list of str or list of list of str): Reference sentences.
+    batch_size (int): BARTScore processing batch size.
+    segment_scores (bool): If True, then score for each instance are returned separately. Otherwise,
+        average score is returned.
+
+Returns:
+    'score': BARTScore loss.
+
+Examples:
+
+    >>> bartscore = jury.load_metric("bartscore")
+    >>> predictions = [
+        ["the cat is on the mat", "There is cat playing on mat"],
+        ["Look! what a wonderful day, today.", "Today is a very wonderful day"],
+    ]
+    >>> references = [
+        ["the cat is playing on the mat.", "The cat plays on the mat."],
+        ["Today is a wonderful day", "The weather outside is wonderful."],
+    ]
+    >>> results = bartscore.compute(predictions=predictions, references=references)
+    >>> print(results)
+    {
+        'bartscore': {
+            'score': -1.8933048248291016,
+            'model_checkpoint': 'bartscore-large-cnn',
+            'model_weights': None,
+            'segment_scores': False
+        }
+    }
+"""
+
+_LICENSE = """Copyright 2022 NeuLab.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License."""
+
+CHECKPOINT_URLS = {
+    "bartscore-large-cnn": {
+        "model_checkpoint": "facebook/bart-large-cnn",
+        "model_weights": {
+            "parabank2": "https://drive.google.com/uc?export=download&id=1_7JfF7KOInb7ZrxKHIigTMR4ChVET01m&confirm=t"
+        },
+    }
+}
+
+
+@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
+class BartscoreForLanguageGeneration(MetricForLanguageGeneration):
+    def __init__(
+        self,
+        resulting_name: str = None,
+        compute_kwargs: Dict = None,
+        model_checkpoint: str = "bartscore-large-cnn",
+        model_weights: str = None,
+        max_length: int = 1024,
+        device: str = None,
+        **kwargs,
+    ):
+        self.model_checkpoint = model_checkpoint
+        self.model_weights = model_weights
+        self.max_length = max_length
+        self.device = device
+        self.model_path = None
+        self.scorer = None
+        super().__init__(resulting_name=resulting_name, compute_kwargs=compute_kwargs, **kwargs)
+
+    def _download_and_prepare(self, dl_manager) -> None:
+        """
+        Downloads and import the computation of BARTScore score from the implementation
+        of BARTScore computation from neulab/BARTScore. The code is sourced from a specific
+        commit on the master branch, in order to keep things stable. See
+        https://github.com/neulab/BARTScore/blob/47b8341854e1b8be965b65480ce236b0c2f7543b/bart_score.py
+        """
+
+        if self.model_checkpoint.lower() in CHECKPOINT_URLS:
+            checkpoint_name = self.model_checkpoint.lower()
+        else:
+            raise KeyError(
+                f"{self.model_checkpoint} checkpoint not found. You should supply the name of a model checkpoint for BARTScore in {CHECKPOINT_URLS.keys()}"
+            )
+        model_checkpoint = CHECKPOINT_URLS[checkpoint_name]["model_checkpoint"]
+
+        if self.model_weights:
+            if self.model_weights.lower() in CHECKPOINT_URLS[checkpoint_name]["model_weights"]:
+                weights_name = self.model_weights.lower()
+            else:
+                raise KeyError(
+                    f"Weights named '{self.model_weights}' not found. You should supply the name of a model weight for BARTScore in "
+                    + str(list(CHECKPOINT_URLS[checkpoint_name]["model_weights"].keys()))
+                )
+            model_path = CHECKPOINT_URLS[checkpoint_name]["model_weights"][weights_name]
+        else:
+            model_path = None
+
+        bartscore_source = (
+            "https://raw.githubusercontent.com/neulab/BARTScore/47b8341854e1b8be965b65480ce236b0c2f7543b/bart_score.py"
+        )
+        self.external_module_path = dl_manager.download(bartscore_source)
+        BARTScorer = self._get_external_resource("bart_score", attr="BARTScorer")
+
+        self.scorer = BARTScorer(device=self.device, max_length=self.max_length, checkpoint=model_checkpoint)
+
+        if model_path is not None:
+            model_dest = dl_manager.download(model_path)
+            self.scorer.load(path=model_dest)
+
+    def _info(self):
+        return datasets.MetricInfo(
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            homepage="https://github.com/neulab/BARTScore",
+            inputs_description=_KWARGS_DESCRIPTION,
+            features=self._default_features,
+            codebase_urls=["https://github.com/neulab/BARTScore"],
+            reference_urls=[
+                "https://github.com/neulab/BARTScore",
+                "https://arxiv.org/abs/2106.11520",
+            ],
+            license=_LICENSE,
+        )
+
+    def _compute_single_pred_single_ref(
+        self,
+        predictions: LanguageGenerationInstance,
+        references: LanguageGenerationInstance,
+        reduce_fn: Callable = None,
+        batch_size: int = 4,
+        segment_scores: bool = False,
+        **kwargs,
+    ):
+        score = self.scorer.score(predictions, references, batch_size=batch_size, **kwargs)
+        if not segment_scores:
+            score = float(np.mean(score))
+
+        return {
+            "score": score,
+            "model_checkpoint": self.model_checkpoint,
+            "model_weights": self.model_weights,
+            "segment_scores": segment_scores,
+        }
+
+    def _compute_single_pred_multi_ref(
+        self,
+        predictions: LanguageGenerationInstance,
+        references: LanguageGenerationInstance,
+        reduce_fn: Callable = None,
+        batch_size: int = 4,
+        segment_scores: bool = False,
+        **kwargs,
+    ):
+        """
+        Ideally we would use the already-implemented multi-ref function in BARTScore,
+        but unfortunately it only supports two reduction functions and doesn't support
+        inconsistent number of reference per predictions. So instead we pre-combine our
+        prediction/reference pairs and use their single prediction/reference function
+        in order to utilize the batch_size argument.
+        """
+        flat_predictions: List[str] = []
+        flat_references: List[str] = []
+        ranges = []
+        cursor = 0
+        for pred, refs in zip(predictions, references):
+            ref_count = len(refs)
+            flat_predictions += [pred] * ref_count
+            flat_references += refs
+            new_cursor = cursor + ref_count
+            ranges.append((cursor, new_cursor))
+            cursor = new_cursor
+
+        flat_scores = self.scorer.score(flat_predictions, flat_references, batch_size=batch_size, **kwargs)
+        score = []
+        for start, end in ranges:
+            reduced_score = float(reduce_fn(flat_scores[start:end]))
+            score.append(reduced_score)
+
+        if not segment_scores:
+            score = float(np.mean(score))
+
+        return {
+            "score": score,
+            "model_checkpoint": self.model_checkpoint,
+            "model_weights": self.model_weights,
+            "segment_scores": segment_scores,
+        }
+
+    def _compute_multi_pred_multi_ref(
+        self,
+        predictions: LanguageGenerationInstance,
+        references: LanguageGenerationInstance,
+        reduce_fn: Callable = None,
+        batch_size: int = 4,
+        segment_scores: bool = False,
+        **kwargs,
+    ):
+        """
+        Like Single Pred/Multi Ref, we pre-combine all possible prediction/reference
+        pairs into a list of single prediction/single reference pairs in order to utilize
+        the batch_size argument.
+        """
+        flat_predictions: List[str] = []
+        flat_references: List[str] = []
+        ranges = []
+        shapes = []
+        cursor = 0
+        for preds, refs in zip(predictions, references):
+            pred_count = len(preds)
+            ref_count = len(refs)
+            for pred in preds:
+                flat_predictions += [pred] * ref_count
+                flat_references += refs
+            new_cursor = cursor + (pred_count * ref_count)
+            ranges.append((cursor, new_cursor))
+            shapes.append((pred_count, ref_count))
+            cursor = new_cursor
+
+        flat_scores = self.scorer.score(flat_predictions, flat_references, batch_size=batch_size, **kwargs)
+        score = []
+        for pair_range, pair_shape in zip(ranges, shapes):
+            pair_start, pair_end = pair_range
+            pred_count, ref_count = pair_shape
+            pair_scores = []
+            for i in range(pred_count):
+                pred_start = pair_start + i * ref_count
+                pred_end = pred_start + ref_count
+                pred_score = flat_scores[pred_start:pred_end]
+                reduced_pred_score = float(reduce_fn(pred_score))
+                pair_scores.append(reduced_pred_score)
+            reduced_score = float(reduce_fn(pair_scores))
+            score.append(reduced_score)
+
+        if not segment_scores:
+            score = float(np.mean(score))
+
+        return {
+            "score": score,
+            "model_checkpoint": self.model_checkpoint,
+            "model_weights": self.model_weights,
+            "segment_scores": segment_scores,
+        }
diff --git a/pyproject.toml b/pyproject.toml
index 1cac48e..527b93c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ exclude = '''
 
 [tool.pytest.ini_options]
 minversion = "6.2"
-timeout = 300
+timeout = 400
 timeout_method = "thread"
 testpaths = "tests"
 addopts = "-vvv"
diff --git a/setup.py b/setup.py
index f083c20..ef01e67 100644
--- a/setup.py
+++ b/setup.py
@@ -39,9 +39,9 @@ def add_pywin(reqs: List[str]) -> None:
     "deepdiff==5.5.0",
     "flake8==3.9.2",
     "isort==5.9.2",
-    "pytest>=6.2.4",
-    "pytest-cov>=2.12.1",
-    "pytest-timeout>=1.4.2",
+    "pytest>=7.0.1",
+    "pytest-cov>=3.0.0",
+    "pytest-timeout>=2.1.0",
 ]
 
 _PRISM_REQUIREMENTS = ["fairseq==0.9.0", "validators"]


Test patch:
diff --git a/tests/jury/metrics/test_bartscore.py b/tests/jury/metrics/test_bartscore.py
new file mode 100644
index 0000000..edf616e
--- /dev/null
+++ b/tests/jury/metrics/test_bartscore.py
@@ -0,0 +1,64 @@
+import pytest
+
+from jury import Jury
+from jury.metrics import AutoMetric
+from tests.jury.conftest import get_expected_output
+from tests.utils import assert_almost_equal_dict
+
+
+@pytest.fixture(scope="module")
+def jury_bartscore():
+    metric = AutoMetric.load("bartscore", device="cpu")
+    return Jury(metrics=metric)
+
+
+@pytest.fixture(scope="module")
+def jury_bartscore_segment():
+    metric = AutoMetric.load("bartscore", device="cpu", compute_kwargs={"segment_scores": True})
+    return Jury(metrics=metric)
+
+
+@pytest.fixture
+@get_expected_output(prefix="metrics")
+def output_basic():
+    return output_basic.output
+
+
+@pytest.fixture
+@get_expected_output(prefix="metrics")
+def output_basic_segment():
+    return output_basic_segment.output
+
+
+@pytest.fixture
+@get_expected_output(prefix="metrics")
+def output_multiple_ref():
+    return output_multiple_ref.output
+
+
+@pytest.fixture
+@get_expected_output(prefix="metrics")
+def output_multiple_pred_multiple_ref():
+    return output_multiple_pred_multiple_ref.output
+
+
+def test_basic(predictions, references, jury_bartscore, output_basic):
+    scores = jury_bartscore(predictions=predictions, references=references)
+    assert_almost_equal_dict(actual=scores, desired=output_basic)
+
+
+def test_basic_segment(predictions, references, jury_bartscore_segment, output_basic_segment):
+    scores = jury_bartscore_segment(predictions=predictions, references=references)
+    assert_almost_equal_dict(actual=scores, desired=output_basic_segment)
+
+
+def test_multiple_ref(predictions, multiple_references, jury_bartscore, output_multiple_ref):
+    scores = jury_bartscore(predictions=predictions, references=multiple_references)
+    assert_almost_equal_dict(actual=scores, desired=output_multiple_ref)
+
+
+def test_multiple_pred_multiple_ref(
+    multiple_predictions, multiple_references, jury_bartscore, output_multiple_pred_multiple_ref
+):
+    scores = jury_bartscore(predictions=multiple_predictions, references=multiple_references)
+    assert_almost_equal_dict(actual=scores, desired=output_multiple_pred_multiple_ref)
diff --git a/tests/test_data/expected_outputs/metrics/test_bartscore.json b/tests/test_data/expected_outputs/metrics/test_bartscore.json
new file mode 100644
index 0000000..cd5aaf1
--- /dev/null
+++ b/tests/test_data/expected_outputs/metrics/test_bartscore.json
@@ -0,0 +1,42 @@
+{
+    "basic": {
+        "total_items": 2,
+        "empty_items": 0,
+        "bartscore": {
+          "score": -2.3962643146514893,
+          "model_checkpoint": "bartscore-large-cnn",
+          "model_weights": null,
+          "segment_scores": false
+        }
+      },
+    "basic_segment": {
+        "total_items": 2,
+        "empty_items": 0,
+        "bartscore": {
+          "score": [-1.740734338760376, -3.0517942905426025],
+          "model_checkpoint": "bartscore-large-cnn",
+          "model_weights": null,
+          "segment_scores": true
+        }
+      },
+    "multiple_ref": {
+        "total_items": 2,
+        "empty_items": 0,
+        "bartscore": {
+          "score": -2.4193317890167236,
+          "model_checkpoint": "bartscore-large-cnn",
+          "model_weights": null,
+          "segment_scores": false
+        }
+      },
+    "multiple_pred_multiple_ref": {
+        "total_items": 2,
+        "empty_items": 0,
+        "bartscore": {
+          "score": -1.8933048248291016,
+          "model_checkpoint": "bartscore-large-cnn",
+          "model_weights": null,
+          "segment_scores": false
+        }
+      }
+}
diff --git a/tests/test_data/expected_outputs/metrics/test_cer.json b/tests/test_data/expected_outputs/metrics/test_cer.json
index ec3eed6..9dc0497 100644
--- a/tests/test_data/expected_outputs/metrics/test_cer.json
+++ b/tests/test_data/expected_outputs/metrics/test_cer.json
@@ -51,4 +51,4 @@
       }
     }
   }
-}
+}
\ No newline at end of file