JetBrains-Research/commit-message-editing-visualization

1import evaluate2from rapidfuzz.distance.Levenshtein import distance, normalized_similarity3 4import config5 6BLEU = evaluate.load("saridormi/b_norm", cache_dir=config.CACHE_DIR)7 8 9def bleu_fn(pred, ref, **kwargs):10    if "refs" in kwargs:11        return BLEU.compute(predictions=[pred] * len(kwargs["refs"]), references=kwargs["refs"])["b_norm"]12    return BLEU.compute(predictions=[pred], references=[ref])["b_norm"]13 14 15METEOR = evaluate.load("meteor", cache_dir=config.CACHE_DIR)16 17 18def meteor_fn(pred, ref, **kwargs):19    if "refs" in kwargs:20        return METEOR.compute(predictions=[pred] * len(kwargs["refs"]), references=kwargs["refs"])["meteor"]21    return METEOR.compute(predictions=[pred], references=[ref])["meteor"]22 23 24ROUGE = evaluate.load("rouge", cache_dir=config.CACHE_DIR)25 26 27def rouge1_fn(pred, ref, **kwargs):28    if "refs" in kwargs:29        return ROUGE.compute(predictions=[pred] * len(kwargs["refs"]), references=kwargs["refs"])["rouge1"]30    return ROUGE.compute(predictions=[pred], references=[ref])["rouge1"]31 32 33def rouge2_fn(pred, ref, **kwargs):34    if "refs" in kwargs:35        return ROUGE.compute(predictions=[pred] * len(kwargs["refs"]), references=kwargs["refs"])["rouge2"]36    return ROUGE.compute(predictions=[pred], references=[ref])["rouge2"]37 38 39def rougeL_fn(pred, ref, **kwargs):40    if "refs" in kwargs:41        return ROUGE.compute(predictions=[pred] * len(kwargs["refs"]), references=kwargs["refs"])["rougeL"]42    return ROUGE.compute(predictions=[pred], references=[ref])["rougeL"]43 44 45BERTSCORE = evaluate.load("bertscore", cache_dir=config.CACHE_DIR)46 47 48def bertscore_fn(pred, ref, **kwargs):49    if "refs" in kwargs:50        return BERTSCORE.compute(predictions=[pred], references=[kwargs["refs"]], model_type="distilbert-base-uncased")[51            "f1"52        ][0]53    return BERTSCORE.compute(predictions=[pred], references=[ref], model_type="distilbert-base-uncased")["f1"][0]54 55 56CHRF = evaluate.load("chrf")57 58 59def chrf_fn(pred, ref, **kwargs):60    if "refs" in kwargs:61        return CHRF.compute(predictions=[pred], references=[kwargs["refs"]])["score"]62    return CHRF.compute(predictions=[pred], references=[[ref]])["score"]63 64 65def edit_distance_fn(pred, ref, **kwargs):66    if "refs" in kwargs:67        scores = [distance(pred, ref) for ref in kwargs["refs"]]68        return sum(scores) / len(scores)69    return distance(pred, ref)70 71 72def edit_distance_norm_fn(pred, ref, **kwargs):73    if "refs" in kwargs:74        scores = [normalized_similarity(pred, ref) for ref in kwargs["refs"]]75        return sum(scores) / len(scores)76    return normalized_similarity(pred, ref)77 78 79AGGR_METRICS = {80    "editdist": edit_distance_fn,81    "editsim": edit_distance_norm_fn,82    "bleu": bleu_fn,83    "meteor": meteor_fn,84    "rouge1": rouge1_fn,85    "rouge2": rouge2_fn,86    "rougeL": rougeL_fn,87    "bertscore": bertscore_fn,88    "chrF": chrf_fn,89}90 91 92REL_METRICS = {93    "editdist": edit_distance_fn,94}95