unite mup

app.py

@@ -0,0 +1,6 @@
+import evaluate
+from evaluate.utils import launch_gradio_widget
+
+
+module = evaluate.load("unite_mup")
+launch_gradio_widget(module)

requirements.txt

@@ -0,0 +1 @@
+unbabel-comet>=2.0.1

unitemup.py

@@ -0,0 +1,90 @@
+# 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.
+
+import os
+import pathlib
+
+import datasets
+import evaluate
+from huggingface_hub import snapshot_download
+
+from comet.models.multitask.unified_metric import UnifiedMetric
+
+
+_CITATION = """\
+@inproceedings{wan-etal-2022-unite,
+    title = "{U}ni{TE}: Unified Translation Evaluation",
+    author = "Wan, Yu  and
+      Liu, Dayiheng  and
+      Yang, Baosong  and
+      Zhang, Haibo  and
+      Chen, Boxing  and
+      Wong, Derek  and
+      Chao, Lidia",
+    booktitle = "Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
+    month = may,
+    year = "2022",
+    address = "Dublin, Ireland",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2022.acl-long.558",
+    doi = "10.18653/v1/2022.acl-long.558",
+    pages = "8117--8127",
+}
+"""
+
+
+_DESCRIPTION = """\
+From https://huggingface.co/Unbabel/unite-mup
+"""
+
+class UNITEMUP(evaluate.Metric):
+    def _info(self):
+        return evaluate.MetricInfo(
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            features=datasets.Features(
+                {
+                    "predictions": datasets.Value("string"),
+                    "references": datasets.Value("string"),
+                }
+            ),
+        )
+    
+    def _download_and_prepare(self, dl_manager):
+        try:
+            model_checkpoint_path = next(pathlib.Path('./models--Unbabel--unite-mup/').rglob('*.ckpt'))
+            self.model = UnifiedMetric.load_from_checkpoint(model_checkpoint_path)
+        except:
+            model_path = snapshot_download(repo_id="Unbabel/unite-mup", cache_dir=os.path.abspath(os.path.dirname('.')))
+            model_checkpoint_path = f"{model_path}/checkpoints/model.ckpt"
+            self.model = UnifiedMetric.load_from_checkpoint(model_checkpoint_path)
+
+
+    def _compute(
+        self,
+        predictions,
+        references,
+        data_keys=None,
+    ):  # Allows user to use either source inputs or reference translations as ground truth.
+        data = [{data_keys[0]: p, data_keys[1]: r} for p, r in zip(predictions, references)]
+        return {"scores": self.model.predict(data, batch_size=8).scores}
+
+
+    def compute_triplet(
+        self,
+        predictions,
+        references,
+        sources,
+    ):  # Unified scores, uses sources, hypotheses and references.
+        data = [{"src": s, "mt": p, "ref": r} for s, p, r in zip(sources, predictions, references)]
+        return {"scores": self.model.predict(data, batch_size=8).metadata.unified_scores}
+