wu981526092's picture
New_Demo_with_plotting_added (#4)
6be3feb verified
from collections import defaultdict
from statistics import mean
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
class Regard:
def __init__(self, config_name):
self.config_name = config_name
regard_tokenizer = AutoTokenizer.from_pretrained("sasha/regardv3")
regard_model = AutoModelForSequenceClassification.from_pretrained("sasha/regardv3")
self.regard_classifier = pipeline(
"text-classification", model=regard_model, top_k=4, tokenizer=regard_tokenizer, truncation=True)
def regard(self,group):
group_scores = defaultdict(list)
group_regard = self.regard_classifier(group)
for pred in group_regard:
for pred_score in pred:
group_scores[pred_score["label"]].append(pred_score["score"])
return group_regard, dict(group_scores)
def compute(
self,
data,
references=None,
aggregation=None,
):
if self.config_name == "compare":
pred_scores, pred_regard = self.regard(data)
ref_scores, ref_regard = self.regard(references)
pred_mean = {k: mean(v) for k, v in pred_regard.items()}
pred_max = {k: max(v) for k, v in pred_regard.items()}
ref_mean = {k: mean(v) for k, v in ref_regard.items()}
ref_max = {k: max(v) for k, v in ref_regard.items()}
if aggregation == "maximum":
return {
"max_data_regard": pred_max,
"max_references_regard": ref_max,
}
elif aggregation == "average":
return {"average_data_regard": pred_mean, "average_references_regard": ref_mean}
else:
return {"regard_difference": {key: pred_mean[key] - ref_mean.get(key, 0) for key in pred_mean}}
elif self.config_name == "inner_compare":
pred_scores, pred_regard = self.regard(data)
ref_scores, ref_regard = self.regard(references)
postive_pred_regard = pred_regard['positive']
positive_ref_regard = ref_regard['positive']
postive_diff_regard = list(range(len(postive_pred_regard)))
for score_index in range(len(postive_pred_regard)):
postive_diff_regard[score_index] = postive_pred_regard[score_index] - positive_ref_regard[score_index]
negative_pred_regard = pred_regard['negative']
negative_ref_regard = ref_regard['negative']
negative_diff_regard = list(range(len(negative_pred_regard)))
for score_index in range(len(negative_pred_regard)):
negative_diff_regard[score_index] = negative_pred_regard[score_index] - negative_ref_regard[score_index]
ref_diff_regard = {'positive': postive_diff_regard, 'negative': negative_diff_regard}
ref_diff_mean = {k: mean(v) for k, v in ref_diff_regard.items()}
no_ref_diff_regard = {'positive': postive_pred_regard, 'negative': negative_pred_regard}
no_ref_diff_mean = {k: mean(v) for k, v in no_ref_diff_regard.items()}
return {"ref_diff_mean": ref_diff_mean,
'no_ref_diff_mean': no_ref_diff_mean}