codebleu.py

# 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
from .bleu import *
from .weighted_ngram_match import *
from .syntax_match import *
from .dataflow_match import *
from tree_sitter import Language, Parser
import os


# 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 = """\
This new module is designed to solve this great ML task and is crafted with a lot of care.
"""


# 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}
"""

# TODO: Define external resources urls if needed
BAD_WORDS_URL = "http://url/to/external/resource/bad_words.txt"


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class CodeBLEU(evaluate.Metric):
    """TODO: Short description of my evaluation module."""

    def _info(self):
        # TODO: Specifies the evaluate.EvaluationModuleInfo object
        return evaluate.MetricInfo(
            # This is the description that will appear on the modules page.
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            # This defines the format of each prediction and reference
            features=datasets.Features({
                'predictions': datasets.Value('string'),
                'references': datasets.Sequence(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
      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
      }