Update codebleu.py

codebleu.py

@@ -15,11 +15,7 @@
 
 import evaluate
 import datasets
-from .bleu import *
-from .weighted_ngram_match import *
-from .syntax_match import *
-from .dataflow_match import *
-from tree_sitter import Language, Parser
+from .calc_code_bleu import calculate
 import os
 
 
@@ -92,79 +88,10 @@ class CodeBLEU(evaluate.Metric):
       # TODO: Download external resources if needed
       if self.config_name == "python":
         Language.build_library('./parser/my-languages.so',['tree-sitter-python'])
-      elif self.config_name == "go":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-go'])
-      elif self.config_name == "javascript":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-javascript'])
-      elif self.config_name == "php":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-php'])
-      elif self.config_name == "java":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-java'])
-      elif self.config_name == "ruby":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-ruby'])
-      elif self.config_name == "c-sharp":
-        Language.build_library('./parser/my-languages.so',['tree-sitter-c-sharp'])
       elif self.config_name == "cpp":
         Language.build_library('./parser/my-languages.so',['tree-sitter-cpp'])
 
 
     def _compute(self, predictions, references, language="python", alpha=0.25, beta=0.25, gamma=0.25, theta=0.25):
-
-      # preprocess inputs
-      pre_references = [[s.strip() for s in my_list] for my_list in references]
-      #pre_references = [[x.strip() for x in open(file, 'r', encoding='utf-8').readlines()] for file in references]
-      hypothesis = [s.strip() for s in predictions]
-      #hypothesis = [x.strip() for x in open(predictions, 'r', encoding='utf-8').readlines()]
-
-      for i in range(len(pre_references)):
-          assert len(hypothesis) == len(pre_references[i])
-
-      references = []
-      for i in range(len(hypothesis)):
-          ref_for_instance = []
-          for j in range(len(pre_references)):
-              ref_for_instance.append(pre_references[j][i])
-          references.append(ref_for_instance)
-      assert len(references) == len(pre_references)*len(hypothesis)
-
-
-      # calculate ngram match (BLEU)
-      tokenized_hyps = [x.split() for x in hypothesis]
-      tokenized_refs = [[x.split() for x in reference] for reference in references]
-
-      ngram_match_score = corpus_bleu(tokenized_refs,tokenized_hyps)
-
-      # calculate weighted ngram match
-      # from os import listdir
-      # from os.path import isfile, join
-      # onlyfiles = [f for f in listdir("./keywords") if isfile(join("keywords", f))]
-      # print(onlyfiles)
-      curr_path = os.path.dirname(os.path.abspath(__file__))
-      keywords = [x.strip() for x in open(curr_path + language +'.txt', 'r', encoding='utf-8').readlines()]
-      def make_weights(reference_tokens, key_word_list):
-          return {token:1 if token in key_word_list else 0.2 \
-                  for token in reference_tokens}
-      tokenized_refs_with_weights = [[[reference_tokens, make_weights(reference_tokens, keywords)]\
-                  for reference_tokens in reference] for reference in tokenized_refs]
-
-      weighted_ngram_match_score = corpus_weighted_ngram_match(tokenized_refs_with_weights,tokenized_hyps)
-
-      # calculate syntax match
-      syntax_match_score = corpus_syntax_match(references, hypothesis, language)
-
-      # calculate dataflow match
-      dataflow_match_score = corpus_dataflow_match(references, hypothesis, language)
-
-
-
-      code_bleu_score = alpha*ngram_match_score\
-                      + beta*weighted_ngram_match_score\
-                      + gamma*syntax_match_score\
-                      + theta*dataflow_match_score
-      return {
-        "ngram_match_score": ngram_match_score,
-        "weighted_ngram_match_score": weighted_ngram_match_score,
-        "syntax_match_score": syntax_match_score,
-        "dataflow_match_score": dataflow_match_score,
-        "code_bleu_score": code_bleu_score
-      }
+      
+      return calculate(predictions, references, language, alpha, beta, gamma, theta)