Create sescore_english_mt.py

sescore_english_mt.py

@@ -0,0 +1,139 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
+#
+# 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.
+"""SEScore: a text generation evaluation metric """
+
+import evaluate
+import datasets
+
+import comet
+from typing import Dict
+import torch
+from comet.encoders.base import Encoder
+from comet.encoders.bert import BERTEncoder
+from transformers import AutoModel, AutoTokenizer
+
+class robertaEncoder(BERTEncoder):
+    def __init__(self, pretrained_model: str) -> None:
+        super(Encoder, self).__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model)
+        self.model = AutoModel.from_pretrained(
+            pretrained_model, add_pooling_layer=False
+        )
+        self.model.encoder.output_hidden_states = True
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model: str) -> Encoder:
+        return robertaEncoder(pretrained_model)
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> Dict[str, torch.Tensor]:
+        last_hidden_states, _, all_layers = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=False,
+        )
+        return {
+            "sentemb": last_hidden_states[:, 0, :],
+            "wordemb": last_hidden_states,
+            "all_layers": all_layers,
+            "attention_mask": attention_mask,
+        }
+
+
+# TODO: Add BibTeX citation
+_CITATION = """\
+@inproceedings{xu-etal-2022-not,
+  title={Not All Errors are Equal: Learning Text Generation Metrics using Stratified Error Synthesis},
+  author={Xu, Wenda and Tuan, Yi-lin and Lu, Yujie and Saxon, Michael and Li, Lei and Wang, William Yang},
+  booktitle ={Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing},
+  month={dec},
+  year={2022},
+  url={https://arxiv.org/abs/2210.05035}
+}
+"""
+
+_DESCRIPTION = """\
+SEScore is an evaluation metric that trys to compute an overall score to measure text generation quality.
+"""
+
+_KWARGS_DESCRIPTION = """
+Calculates how good are predictions given some references
+Args:
+    predictions: list of candidate outputs
+    references: list of references 
+Returns:
+    {"mean_score": mean_score, "scores": scores}
+    
+Examples:
+    >>> import evaluate
+    >>> sescore = evaluate.load("xu1998hz/sescore")
+    >>> score = sescore.compute(
+            references=['sescore is a simple but effective next-generation text evaluation metric'],
+            predictions=['sescore is simple effective text evaluation metric for next generation']
+        )
+"""
+
+# TODO: Define external resources urls if needed
+BAD_WORDS_URL = "http://url/to/external/resource/bad_words.txt"
+
+
+@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
+class SEScore(evaluate.Metric):
+    """SEScore"""
+
+    def _info(self):
+        # TODO: Specifies the evaluate.EvaluationModuleInfo object
+        return evaluate.MetricInfo(
+            # This is the description that will appear on the modules page.
+            module_type="metric",
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            inputs_description=_KWARGS_DESCRIPTION,
+            # This defines the format of each prediction and reference
+            features=datasets.Features({
+                'predictions': datasets.Value("string", id="sequence"),
+                'references': datasets.Value("string", id="sequence"),
+            }),
+            # Homepage of the module for documentation
+            homepage="http://module.homepage",
+            # Additional links to the codebase or references
+            codebase_urls=["http://github.com/path/to/codebase/of/new_module"],
+            reference_urls=["http://path.to.reference.url/new_module"]
+        )
+
+    def _download_and_prepare(self, dl_manager):
+        """download SEScore checkpoints to compute the scores"""
+        # Download SEScore checkpoint
+        from comet import load_from_checkpoint
+        import os
+        from huggingface_hub import snapshot_download
+        # initialize roberta into str2encoder
+        comet.encoders.str2encoder['RoBERTa'] = robertaEncoder
+        print("config name: ", self.config_name)
+        if self.config_name == "default":
+            destination = snapshot_download(repo_id="xu1998hz/sescore_english_mt", revision="main")
+            self.scorer = load_from_checkpoint(f'{destination}/checkpoint/sescore_english_mt.ckpt')
+        else:
+            print("Config name is not supported!")
+
+    def _compute(self, predictions, references, gpus=None, progress_bar=False):
+        if gpus is None:
+            gpus = 1 if torch.cuda.is_available() else 0
+
+        data = {"src": references, "mt": predictions}
+        data = [dict(zip(data, t)) for t in zip(*data.values())]
+        scores, mean_score = self.scorer.predict(data, gpus=gpus, progress_bar=progress_bar)
+        return {"mean_score": mean_score, "scores": scores}