# 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.
"""TODO: Add a description here."""
import evaluate
import datasets
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import getpass
import pdb
import os
import torch
from rouge_score import scoring
from contextlib import contextmanager
# TODO: Add BibTeX citation
_CITATION = """\
@InProceedings{huggingface:module,
title = {A great new module},
authors={huggingface, Inc.},
year={2020}
}
"""
# TODO: Add description of the module here
_DESCRIPTION = """\
local coherecence with classifier trained on the shuffle task, window=3 sentences
"""
# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates how good are predictions given some references, using certain scores
Args:
predictions: list of predictions to score. Each predictions
should be a string with tokens separated by spaces.
references: list of reference for each prediction. Each
reference should be a string with tokens separated by spaces.
Returns:
accuracy: description of the first score,
another_score: description of the second score,
Examples:
Examples should be written in doctest format, and should illustrate how
to use the function.
>>> my_new_module = evaluate.load("my_new_module")
>>> results = my_new_module.compute(references=[0, 1], predictions=[0, 1])
>>> print(results)
{'accuracy': 1.0}
"""
WINDOW_SIZE = 3
@contextmanager
def filter_logging_context():
def filter_log(record):
return False if "This IS expected if you are initializing" in record.msg else True
logger = datasets.utils.logging.get_logger("transformers.modeling_utils")
logger.addFilter(filter_log)
try:
yield
finally:
logger.removeFilter(filter_log)
class Scorer:
def __init__(
self,
model_type=None,
batch_size=64,
device=None,
use_fast_tokenizer=False):
if device is not None:
# assert device in ["gpu", "cpu", "cuda"], "device should be either gpu or cpu."
if device == "gpu":
device = "cuda"
else:
device = "cuda" if torch.cuda.is_available() else "cpu"
self.device = device
self.model_type = model_type
self.batch_size = batch_size
self._tokenizer = AutoTokenizer.from_pretrained("roberta-large")
self._model = AutoModelForSequenceClassification.from_pretrained(f"ronaldahmed/ccl_win-{model_type}")
self._model.to(device)
self._model.eval()
@property
def hash(self):
return self.model_type
def preprocess_adjacent_window(self,preds):
pred_list = []
lens = []
for pred in preds:
sents = pred.split("\n")
ns = len(sents)
if ns <= WINDOW_SIZE:
pred_list.append(pred)
lens.append(1)
else:
llen = 0
for i in range(0,ns-WINDOW_SIZE+1):
sss = sents[i:i+WINDOW_SIZE]
ss = "\n".join(sss)
pred_list.append(ss)
llen += 1
lens.append(llen)
#
return pred_list,lens
def score(self,predictions):
sent_lens = [len(x.split("\n")) for x in predictions]
pred_list,len_by_sample = self.preprocess_adjacent_window(predictions)
scores = []
n_preds = len(pred_list)
with torch.no_grad():
for b in range(0,n_preds,self.batch_size):
strides = [x.lower() for x in pred_list[b:b+self.batch_size]]
tinput = self._tokenizer(strides,padding=True,truncation=True,max_length=512,return_tensors="pt")
tinput = {k:v.to(self.device) for k,v in tinput.items()}
output = self._model(**tinput)
probs = torch.softmax(output.logits,dim=-1).detach().cpu().numpy()
scores.extend(probs[:,0].tolist())
#
results = []
offset = 0
for i,_len in enumerate(len_by_sample):
score = float(np.mean(scores[offset:offset+_len])) if sent_lens[i]>1 else 0.
results.append(score)
offset += _len
#
return results
@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class ccl_win(evaluate.Measurement):
"""TODO: Short description of my evaluation module."""
def _info(self):
# TODO: Specifies the evaluate.EvaluationModuleInfo object
return evaluate.MeasurementInfo(
# This is the description that will appear on the modules page.
module_type="measurement",
description=_DESCRIPTION,
citation=_CITATION,
inputs_description=_KWARGS_DESCRIPTION,
# This defines the format of each prediction and reference
features=datasets.Features({
'predictions': datasets.Value('string'),
}),
# 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):
"""Optional: download external resources useful to compute the scores"""
# TODO: Download external resources if needed
pass
def _compute(self, predictions, dataset="arxiv", batch_size: int = 16, device=None, use_aggregator=True):
"""Returns the scores"""
hashcode = dataset
with filter_logging_context():
if not hasattr(self, "cached_scorer") or self.cached_scorer.hash != hashcode:
self.cached_scorer = Scorer(
model_type=dataset,
batch_size=batch_size,
device=device,
)
results = self.cached_scorer.score(predictions)
outres = {}
aggregator = None
if use_aggregator:
np.random.seed(42)
aggregator = scoring.BootstrapAggregator()
for score in results:
aggregator.add_scores({"loc_coh_ccl": score})
#
res = aggregator.aggregate()
for k in res: outres[k] = res[k].mid
else:
outres = {"loc_coh_ccl": results}
return outres