idsedykh/codebleu

1# Copyright (c) Microsoft Corporation.2# Licensed under the MIT license.3 4# -*- coding:utf-8 -*-5import os6import logging7from . import bleu8from . import weighted_ngram_match9from . import syntax_match10from . import dataflow_match11 12 13def calc_codebleu(predictions, references, lang, tokenizer=None, params='0.25,0.25,0.25,0.25'):14    """_summary_15 16    Args:17        predictions (list[str]): list of predictions18        references (list[str]): list of lists with references19        lang (str): ['java','js','c_sharp','php','go','python','ruby']20        tokenizer (callable): tokenizer function, Defaults to lambda s: s.split()21        params (str, optional): Defaults to '0.25,0.25,0.25,0.25'.22    """23 24    alpha, beta, gamma, theta = [float(x) for x in params.split(',')]25 26    # preprocess inputs27    references = [[x.strip() for x in ref] for ref in references]28    hypothesis = [x.strip() for x in predictions]29 30    if not len(references) == len(hypothesis):31        raise ValueError32 33    # calculate ngram match (BLEU)34    if tokenizer is None:35        tokenizer = lambda s: s.split()36 37    tokenized_hyps = [tokenizer(x) for x in hypothesis]38    tokenized_refs = [[tokenizer(x) for x in reference]39                      for reference in references]40 41    ngram_match_score = bleu.corpus_bleu(tokenized_refs, tokenized_hyps)42 43    # calculate weighted ngram match44    keywords = [x.strip() for x in open(os.path.abspath(os.path.dirname(__file__)) + '/keywords/' + lang +45                                        '.txt', 'r', encoding='utf-8').readlines()]46 47    def make_weights(reference_tokens, key_word_list):48        return {token: 1 if token in key_word_list else 0.249                for token in reference_tokens}50    tokenized_refs_with_weights = [[[reference_tokens, make_weights(reference_tokens, keywords)]51                                    for reference_tokens in reference] for reference in tokenized_refs]52 53    weighted_ngram_match_score = weighted_ngram_match.corpus_bleu(54        tokenized_refs_with_weights, tokenized_hyps)55 56    # calculate syntax match57    syntax_match_score = syntax_match.corpus_syntax_match(58        references, hypothesis, lang)59 60    # calculate dataflow match61    dataflow_match_score = dataflow_match.corpus_dataflow_match(62        references, hypothesis, lang)63 64    # print('ngram match: {0}, weighted ngram match: {1}, syntax_match: {2}, dataflow_match: {3}'.65        #   format(ngram_match_score, weighted_ngram_match_score, syntax_match_score, dataflow_match_score))66 67    code_bleu_score = alpha*ngram_match_score\68        + beta*weighted_ngram_match_score\69        + gamma*syntax_match_score\70        + theta*dataflow_match_score71 72    # print('CodeBLEU score: ', code_bleu_score)73 74    return {75        'CodeBLEU': code_bleu_score,76        'ngram_match_score': ngram_match_score,77        'weighted_ngram_match_score': weighted_ngram_match_score,78        'syntax_match_score': syntax_match_score,79        'dataflow_match_score': dataflow_match_score80    }81