Spaces:
Build error
Build error
| # 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 | |
| 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" | |
| 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 | |
| destination = snapshot_download(repo_id="xu1998hz/sescore_english_mt", revision="main") | |
| self.scorer = load_from_checkpoint(f'{destination}/checkpoint/sescore_english_mt.ckpt') | |
| 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} | |