Upload model from GitHub.

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+codet5.gif filter=lfs diff=lfs merge=lfs -text
+evaluator/CodeBLEU/parser/my-languages.so filter=lfs diff=lfs merge=lfs -text

CODEOWNERS

@@ -0,0 +1,2 @@
+# Comment line immediately above ownership line is reserved for related gus information. Please be careful while editing.
+#ECCN:Open Source

CODE_OF_CONDUCT.md

@@ -0,0 +1,105 @@
+# Salesforce Open Source Community Code of Conduct
+
+## About the Code of Conduct
+
+Equality is a core value at Salesforce. We believe a diverse and inclusive
+community fosters innovation and creativity, and are committed to building a
+culture where everyone feels included.
+
+Salesforce open-source projects are committed to providing a friendly, safe, and
+welcoming environment for all, regardless of gender identity and expression,
+sexual orientation, disability, physical appearance, body size, ethnicity, nationality, 
+race, age, religion, level of experience, education, socioeconomic status, or 
+other similar personal characteristics.
+
+The goal of this code of conduct is to specify a baseline standard of behavior so
+that people with different social values and communication styles can work
+together effectively, productively, and respectfully in our open source community.
+It also establishes a mechanism for reporting issues and resolving conflicts.
+
+All questions and reports of abusive, harassing, or otherwise unacceptable behavior
+in a Salesforce open-source project may be reported by contacting the Salesforce
+Open Source Conduct Committee at [email protected].
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of gender 
+identity and expression, sexual orientation, disability, physical appearance, 
+body size, ethnicity, nationality, race, age, religion, level of experience, education, 
+socioeconomic status, or other similar personal characteristics.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy toward other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+advances
+* Personal attacks, insulting/derogatory comments, or trolling
+* Public or private harassment
+* Publishing, or threatening to publish, others' private information—such as
+a physical or electronic address—without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+professional setting
+* Advocating for or encouraging any of the above behaviors
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned with this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project email
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the Salesforce Open Source Conduct Committee 
+at [email protected]. All complaints will be reviewed and investigated 
+and will result in a response that is deemed necessary and appropriate to the 
+circumstances. The committee is obligated to maintain confidentiality with 
+regard to the reporter of an incident. Further details of specific enforcement 
+policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership and the Salesforce Open Source Conduct 
+Committee.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][contributor-covenant-home],
+version 1.4, available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html. 
+It includes adaptions and additions from [Go Community Code of Conduct][golang-coc], 
+[CNCF Code of Conduct][cncf-coc], and [Microsoft Open Source Code of Conduct][microsoft-coc].
+
+This Code of Conduct is licensed under the [Creative Commons Attribution 3.0 License][cc-by-3-us].
+
+[contributor-covenant-home]: https://www.contributor-covenant.org (https://www.contributor-covenant.org/)
+[golang-coc]: https://golang.org/conduct
+[cncf-coc]: https://github.com/cncf/foundation/blob/master/code-of-conduct.md
+[microsoft-coc]: https://opensource.microsoft.com/codeofconduct/
+[cc-by-3-us]: https://creativecommons.org/licenses/by/3.0/us/

LICENSE.txt

@@ -0,0 +1,12 @@
+Copyright (c) 2021, Salesforce.com, Inc.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+* Neither the name of Salesforce.com nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,3 +1,247 @@
----
-license: openrail
----
+# CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation
+
+This is the official PyTorch implementation for the following EMNLP 2021 paper from Salesforce Research:
+
+**Title**: [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://arxiv.org/pdf/2109.00859.pdf)
+
+**Authors**: [Yue Wang](https://yuewang-cuhk.github.io/), [Weishi Wang](https://www.linkedin.com/in/weishi-wang/)
+, [Shafiq Joty](https://raihanjoty.github.io/), and [Steven C.H. Hoi](https://sites.google.com/view/stevenhoi/home)
+
+![CodeT5 demo](codet5.gif)
+
+## Updates
+
+**July 06, 2022**
+
+We release two large-sized CodeT5 checkpoints at Hugging Face: [Salesforce/codet5-large](https://huggingface.co/Salesforce/codet5-large) and [Salesforce/codet5-large-ntp-py](https://huggingface.co/Salesforce/codet5-large-ntp-py), which are introduced by the paper: [CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning](https://arxiv.org/pdf/2207.01780.pdf) by Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, Steven C.H. Hoi.
+
+* CodeT5-large was pretrained using Masked Span Prediction (MSP) objective on CodeSearchNet and achieve new SOTA results on several CodeXGLUE benchmarks. The finetuned checkpoints are released at [here](https://console.cloud.google.com/storage/browser/sfr-codet5-data-research/finetuned_models). See Appendix A.1 of the [paper](https://arxiv.org/pdf/2207.01780.pdf) for more details.
+
+* CodeT5-large-ntp-py was first pretrained using Masked Span Prediction (MSP) objective on CodeSearchNet and GCPY (the Python split of [Github Code](https://huggingface.co/datasets/codeparrot/github-code) data), followed by another 10 epochs on GCPY using Next Token Prediction (NTP) objective. 
+
+CodeT5-large-ntp-py is especially optimized for Python code generation tasks and employed as the foundation model for our [CodeRL](https://github.com/salesforce/CodeRL), yielding new SOTA results on the APPS Python competition-level program synthesis benchmark. See the [paper](https://arxiv.org/pdf/2207.01780.pdf) for more details.
+
+**Oct 29, 2021**
+
+We release [fine-tuned checkpoints](https://console.cloud.google.com/storage/browser/sfr-codet5-data-research/finetuned_models)
+for all the downstream tasks covered in the paper.
+
+**Oct 25, 2021**
+
+We release a CodeT5-base fine-tuned
+checkpoint ([Salesforce/codet5-base-multi-sum](https://huggingface.co/Salesforce/codet5-base-multi-sum)) for
+multilingual code summarzation. Below is how to use this model:
+
+```python
+from transformers import RobertaTokenizer, T5ForConditionalGeneration
+
+if __name__ == '__main__':
+    tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
+    model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')
+
+    text = """def svg_to_image(string, size=None):
+    if isinstance(string, unicode):
+        string = string.encode('utf-8')
+        renderer = QtSvg.QSvgRenderer(QtCore.QByteArray(string))
+    if not renderer.isValid():
+        raise ValueError('Invalid SVG data.')
+    if size is None:
+        size = renderer.defaultSize()
+        image = QtGui.QImage(size, QtGui.QImage.Format_ARGB32)
+        painter = QtGui.QPainter(image)
+        renderer.render(painter)
+    return image"""
+
+    input_ids = tokenizer(text, return_tensors="pt").input_ids
+
+    generated_ids = model.generate(input_ids, max_length=20)
+    print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+    # this prints: "Convert a SVG string to a QImage."
+```
+
+**Oct 18, 2021**
+
+We add a [model card](https://github.com/salesforce/CodeT5/blob/main/CodeT5_model_card.pdf) for CodeT5! Please reach out
+if you have any questions about it.
+
+**Sep 24, 2021**
+
+CodeT5 is now in [hugginface](https://huggingface.co/)!
+
+You can simply load the model ([CodeT5-small](https://huggingface.co/Salesforce/codet5-small)
+and [CodeT5-base](https://huggingface.co/Salesforce/codet5-base)) and do the inference:
+
+```python
+from transformers import RobertaTokenizer, T5ForConditionalGeneration
+
+tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
+model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base')
+
+text = "def greet(user): print(f'hello <extra_id_0>!')"
+input_ids = tokenizer(text, return_tensors="pt").input_ids
+
+# simply generate one code span
+generated_ids = model.generate(input_ids, max_length=8)
+print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+# this prints "{user.username}"
+```
+
+## Introduction
+
+This repo provides the code for reproducing the experiments
+in [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://arxiv.org/pdf/2109.00859.pdf)
+. CodeT5 is a new pre-trained encoder-decoder model for programming languages, which is pre-trained on **8.35M**
+functions in 8 programming languages (Python, Java, JavaScript, PHP, Ruby, Go, C, and C#). In total, it achieves
+state-of-the-art results on **14 sub-tasks** in a code intelligence benchmark - [CodeXGLUE](https://github.com/microsoft/CodeXGLUE).
+
+Paper link: https://arxiv.org/abs/2109.00859
+
+Blog link: https://blog.salesforceairesearch.com/codet5/
+
+The code currently includes two pre-trained checkpoints ([CodeT5-small](https://huggingface.co/Salesforce/codet5-small)
+and [CodeT5-base](https://huggingface.co/Salesforce/codet5-base)) and scripts to fine-tune them on 4 generation tasks (
+code summarization, code generation, translation, and refinement) plus 2 understanding tasks (code defect detection and
+clone detection) in CodeXGLUE. We also provide their fine-tuned checkpoints to facilitate the easy replication
+of our paper.
+
+In practice, CodeT5 can be deployed as an AI-powered coding assistant to boost the productivity of software developers.
+At Salesforce, we build an [AI coding assistant demo](https://github.com/salesforce/CodeT5/raw/main/codet5.gif) using
+CodeT5 as a VS Code plugin to provide three capabilities for Apex developers:
+
+- **Text-to-code generation**: generate code based on the natural language description.
+- **Code autocompletion**: complete the whole function of code given the target function name.
+- **Code summarization**: generate the summary of a function in natural language description.
+
+## Table of Contents
+
+1. [Citation](#citation)
+2. [License](#license)
+3. [Dependency](#dependency)
+4. [Download](#download)
+5. [Fine-tuning](#fine-tuning)
+6. [Get Involved](#get-involved)
+
+## Citation
+
+If you find this code to be useful for your research, please consider citing:
+
+```
+@inproceedings{
+    wang2021codet5,
+    title={CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation}, 
+    author={Yue Wang, Weishi Wang, Shafiq Joty, Steven C.H. Hoi},
+    booktitle={Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021},
+    year={2021},
+}
+
+@article{coderl2022,
+  title={CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning},
+  author={Le, Hung and Wang, Yue and Gotmare, Akhilesh Deepak and Savarese, Silvio and Hoi, Steven C. H.},
+  journal={arXiv preprint arXiv:2207.01780},
+  year={2022}
+}
+```
+
+## License
+
+The code is released under the BSD-3 License (see `LICENSE.txt` for details), but we also ask that users respect the
+following:
+
+This software should not be used to promote or profit from:
+
+violence, hate, and division,
+
+environmental destruction,
+
+abuse of human rights, or
+
+the destruction of people's physical and mental health.
+
+We encourage users of this software to tell us about the applications in which they are putting it to use by emailing
+[email protected], and to
+use [appropriate](https://arxiv.org/abs/1810.03993) [documentation](https://www.partnershiponai.org/about-ml/) when
+developing high-stakes applications of this model.
+
+## Dependency
+
+- Pytorch 1.7.1
+- tensorboard 2.4.1
+- transformers 4.6.1
+- tree-sitter 0.2.2
+
+## Download
+
+* [Pre-trained checkpoints](https://console.cloud.google.com/storage/browser/sfr-codet5-data-research/pretrained_models)
+* [Fine-tuning data](https://console.cloud.google.com/storage/browser/sfr-codet5-data-research/data)
+* [Fine-tuned checkpoints](https://console.cloud.google.com/storage/browser/sfr-codet5-data-research/finetuned_models)
+
+Instructions to download:
+
+```
+# pip install gsutil
+cd your-cloned-codet5-path
+
+gsutil -m cp -r "gs://sfr-codet5-data-research/pretrained_models" .
+gsutil -m cp -r "gs://sfr-codet5-data-research/data" .
+gsutil -m cp -r "gs://sfr-codet5-data-research/finetuned_models" .
+```
+
+## Fine-tuning
+
+Go to `sh` folder, set the `WORKDIR` in `exp_with_args.sh` to be your cloned CodeT5 repository path.
+
+You can use `run_exp.py` to run a broad set of experiments by simply passing the `model_tag`, `task`, and `sub_task`
+arguments. In total, we support five models (i.e., ['roberta', 'codebert', 'bart_base', 'codet5_small', 'codet5_base'])
+and six tasks (i.e., ['summarize', 'concode', 'translate', 'refine', 'defect', 'clone']). For each task, we use
+the `sub_task` to specify which specific datasets to fine-tne on. Below is the full list:
+
+| \--task   | \--sub\_task                       | Description                                                                                                                      |
+| --------- | ---------------------------------- | -------------------------------------------------------------------------------------------------------------------------------- |
+| summarize | ruby/javascript/go/python/java/php | code summarization task on [CodeSearchNet](https://arxiv.org/abs/1909.09436) data with six PLs                                   |
+| concode   | none                               | text-to-code generation on [Concode](https://aclanthology.org/D18-1192.pdf) data                                                 |
+| translate | java-cs/cs-java                    | code-to-code translation between [Java and C#](https://arxiv.org/pdf/2102.04664.pdf)                                             |
+| refine    | small/medium                       | code refinement on [code repair data](https://arxiv.org/pdf/1812.08693.pdf) with small/medium functions                          |
+| defect    | none                               | code defect detection in [C/C++ data](https://proceedings.neurips.cc/paper/2019/file/49265d2447bc3bbfe9e76306ce40a31f-Paper.pdf) |
+| clone     | none                               | code clone detection in [Java data](https://arxiv.org/pdf/2002.08653.pdf)                                                        |
+
+For example, if you want to run CodeT5-base model on the code summarization task for Python, you can simply run:
+
+```
+python run_exp.py --model_tag codet5_base --task summarize --sub_task python
+```
+
+For multi-task training, you can type:
+
+```
+python run_exp.py --model_tag codet5_base --task multi_task --sub_task none
+```
+
+Besides, you can specify:
+
+```
+model_dir: where to save fine-tuning checkpoints
+res_dir: where to save the performance results 
+summary_dir: where to save the training curves
+data_num: how many data instances to use, the default -1 is for using the full data
+gpu: the index of the GPU to use in the cluster
+``` 
+
+You can also revise the suggested
+arguments [here](https://github.com/salesforce/CodeT5/blob/0bf3c0c43e92fcf54d9df68c793ac22f2b60aad4/sh/run_exp.py#L14) or directly customize the [exp_with_args.sh](https://github.com/salesforce/CodeT5/blob/main/sh/exp_with_args.sh) bash file.
+Please refer to the argument flags in [configs.py](https://github.com/salesforce/CodeT5/blob/main/configs.py) for the full
+available options. The saved training curves in `summary_dir` can be visualized using [tensorboard](https://pypi.org/project/tensorboard/).
+Note that we employ one A100 GPU for all fine-tuning experiments.
+
+### How to reproduce the results using the released finetuned checkpoints?
+
+* Remove the `--do_train --do_eval --do_eval_bleu` and reserve only `--do_test` at [here](https://github.com/salesforce/CodeT5/blob/5b37c34f4bbbfcfd972c24a9dd1f45716568ecb5/sh/exp_with_args.sh#L84). 
+* Pass the path of your downloaded finetuned checkpoint to load at [here](https://github.com/salesforce/CodeT5/blob/5b37c34f4bbbfcfd972c24a9dd1f45716568ecb5/run_gen.py#L366), e.g., `file = "CodeT5/finetuned_models/summarize_python_codet5_base.bin"`
+* Run the program: `python run_exp.py --model_tag codet5_base --task summarize --sub_task python`
+
+### How to fine-tune on your own task and dataset?
+If you want to fine-tune on your dataset, you can add your own task and sub_task in `configs.py` ([here](https://github.com/salesforce/CodeT5/blob/d27512d23ba6130e089e571d8c3e399760db1c31/configs.py#L11)) and add your data path and the function to read in `utils.py` ([here](https://github.com/salesforce/CodeT5/blob/5bb41e21b07fee73f310476a91ded00e385290d7/utils.py#L103) and [here](https://github.com/salesforce/CodeT5/blob/5bb41e21b07fee73f310476a91ded00e385290d7/utils.py#L149)). The read function can be implemented in `_utils.py` similar to [this one](https://github.com/salesforce/CodeT5/blob/aaf9c4a920c4986abfd54a74f5456b056b6409e0/_utils.py#L213). If your task to add is a generation task, you can simply reuse or customize the `run_gen.py`. For understanding tasks, please refer to `run_defect.py` and `run_clone.py`.
+
+## Get Involved
+
+Please create a GitHub issue if you have any questions, suggestions, requests or bug-reports. We welcome PRs!
+

SECURITY.md

@@ -0,0 +1,7 @@
+## Security
+
+Please report any security issue to [[email protected]](mailto:[email protected])
+as soon as it is discovered. This library limits its runtime dependencies in
+order to reduce the total cost of ownership as much as can be, but all consumers
+should remain vigilant and have their security stakeholders review all third-party
+products (3PP) like this one and their dependencies.

_utils.py

@@ -0,0 +1,306 @@
+import json
+
+
+def add_lang_by_task(target_str, task, sub_task):
+    if task == 'summarize':
+        target_str = '<en> ' + target_str
+    elif task == 'refine':
+        target_str = '<java> ' + target_str
+    elif task == 'translate':
+        if sub_task == 'java-cs':
+            target_str = '<c_sharp> ' + target_str
+        else:
+            target_str = '<java> ' + target_str
+    elif task == 'concode':
+        target_str = '<java> ' + target_str
+    elif task == 'defect':
+        target_str = target_str
+    return target_str
+
+
+def convert_examples_to_features(item):
+    example, example_index, tokenizer, args, stage = item
+
+    if args.model_type in ['t5', 'codet5'] and args.add_task_prefix:
+        if args.sub_task != 'none':
+            source_str = "{} {}: {}".format(args.task, args.sub_task, example.source)
+        else:
+            source_str = "{}: {}".format(args.task, example.source)
+    else:
+        source_str = example.source
+
+    source_str = source_str.replace('</s>', '<unk>')
+    source_ids = tokenizer.encode(source_str, max_length=args.max_source_length, padding='max_length', truncation=True)
+    assert source_ids.count(tokenizer.eos_token_id) == 1
+    if stage == 'test':
+        target_ids = []
+    else:
+        target_str = example.target
+        if args.add_lang_ids:
+            target_str = add_lang_by_task(example.target, args.task, args.sub_task)
+        if args.task in ['defect', 'clone']:
+            if target_str == 0:
+                target_str = 'false'
+            elif target_str == 1:
+                target_str = 'true'
+            else:
+                raise NameError
+        target_str = target_str.replace('</s>', '<unk>')
+        target_ids = tokenizer.encode(target_str, max_length=args.max_target_length, padding='max_length',
+                                      truncation=True)
+        assert target_ids.count(tokenizer.eos_token_id) == 1
+
+    return InputFeatures(
+        example_index,
+        source_ids,
+        target_ids,
+        url=example.url
+    )
+
+
+def convert_clone_examples_to_features(item):
+    example, example_index, tokenizer, args = item
+    if args.model_type in ['t5', 'codet5'] and args.add_task_prefix:
+        source_str = "{}: {}".format(args.task, example.source)
+        target_str = "{}: {}".format(args.task, example.target)
+    else:
+        source_str = example.source
+        target_str = example.target
+    code1 = tokenizer.encode(source_str, max_length=args.max_source_length, padding='max_length', truncation=True)
+    code2 = tokenizer.encode(target_str, max_length=args.max_source_length, padding='max_length', truncation=True)
+    source_ids = code1 + code2
+    return CloneInputFeatures(example_index, source_ids, example.label, example.url1, example.url2)
+
+
+def convert_defect_examples_to_features(item):
+    example, example_index, tokenizer, args = item
+    if args.model_type in ['t5', 'codet5'] and args.add_task_prefix:
+        source_str = "{}: {}".format(args.task, example.source)
+    else:
+        source_str = example.source
+    code = tokenizer.encode(source_str, max_length=args.max_source_length, padding='max_length', truncation=True)
+    return DefectInputFeatures(example_index, code, example.target)
+
+
+class CloneInputFeatures(object):
+    """A single training/test features for a example."""
+
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 label,
+                 url1,
+                 url2
+                 ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.label = label
+        self.url1 = url1
+        self.url2 = url2
+
+
+class DefectInputFeatures(object):
+    """A single training/test features for a example."""
+
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 label
+                 ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.label = label
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+                 url=None
+                 ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids
+        self.url = url
+
+
+class Example(object):
+    """A single training/test example."""
+
+    def __init__(self,
+                 idx,
+                 source,
+                 target,
+                 url=None,
+                 task='',
+                 sub_task=''
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+        self.url = url
+        self.task = task
+        self.sub_task = sub_task
+
+
+class CloneExample(object):
+    """A single training/test example."""
+
+    def __init__(self,
+                 code1,
+                 code2,
+                 label,
+                 url1,
+                 url2
+                 ):
+        self.source = code1
+        self.target = code2
+        self.label = label
+        self.url1 = url1
+        self.url2 = url2
+
+
+def read_translate_examples(filename, data_num):
+    """Read examples from filename."""
+    examples = []
+    assert len(filename.split(',')) == 2
+    src_filename = filename.split(',')[0]
+    trg_filename = filename.split(',')[1]
+    idx = 0
+    with open(src_filename) as f1, open(trg_filename) as f2:
+        for line1, line2 in zip(f1, f2):
+            src = line1.strip()
+            trg = line2.strip()
+            examples.append(
+                Example(
+                    idx=idx,
+                    source=src,
+                    target=trg,
+                )
+            )
+            idx += 1
+            if idx == data_num:
+                break
+    return examples
+
+
+def read_refine_examples(filename, data_num):
+    """Read examples from filename."""
+    examples = []
+    assert len(filename.split(',')) == 2
+    src_filename = filename.split(',')[0]
+    trg_filename = filename.split(',')[1]
+    idx = 0
+
+    with open(src_filename) as f1, open(trg_filename) as f2:
+        for line1, line2 in zip(f1, f2):
+            examples.append(
+                Example(
+                    idx=idx,
+                    source=line1.strip(),
+                    target=line2.strip(),
+                )
+            )
+            idx += 1
+            if idx == data_num:
+                break
+    return examples
+
+
+def read_concode_examples(filename, data_num):
+    """Read examples from filename."""
+    examples = []
+
+    with open(filename) as f:
+        for idx, line in enumerate(f):
+            x = json.loads(line)
+            examples.append(
+                Example(
+                    idx=idx,
+                    source=x["nl"].strip(),
+                    target=x["code"].strip()
+                )
+            )
+            idx += 1
+            if idx == data_num:
+                break
+    return examples
+
+
+def read_summarize_examples(filename, data_num):
+    """Read examples from filename."""
+    examples = []
+    with open(filename, encoding="utf-8") as f:
+        for idx, line in enumerate(f):
+            line = line.strip()
+            js = json.loads(line)
+            if 'idx' not in js:
+                js['idx'] = idx
+            code = ' '.join(js['code_tokens']).replace('\n', ' ')
+            code = ' '.join(code.strip().split())
+            nl = ' '.join(js['docstring_tokens']).replace('\n', '')
+            nl = ' '.join(nl.strip().split())
+            examples.append(
+                Example(
+                    idx=idx,
+                    source=code,
+                    target=nl,
+                )
+            )
+            if idx + 1 == data_num:
+                break
+    return examples
+
+
+def read_defect_examples(filename, data_num):
+    """Read examples from filename."""
+    examples = []
+    with open(filename, encoding="utf-8") as f:
+        for idx, line in enumerate(f):
+            line = line.strip()
+            js = json.loads(line)
+
+            code = ' '.join(js['func'].split())
+            examples.append(
+                Example(
+                    idx=js['idx'],
+                    source=code,
+                    target=js['target']
+                )
+            )
+            if idx + 1 == data_num:
+                break
+    return examples
+
+
+def read_clone_examples(filename, data_num):
+    """Read examples from filename."""
+    index_filename = filename
+    url_to_code = {}
+    with open('/'.join(index_filename.split('/')[:-1]) + '/data.jsonl') as f:
+        for line in f:
+            line = line.strip()
+            js = json.loads(line)
+            code = ' '.join(js['func'].split())
+            url_to_code[js['idx']] = code
+
+    data = []
+    with open(index_filename) as f:
+        idx = 0
+        for line in f:
+            line = line.strip()
+            url1, url2, label = line.split('\t')
+            if url1 not in url_to_code or url2 not in url_to_code:
+                continue
+            if label == '0':
+                label = 0
+            else:
+                label = 1
+            data.append(CloneExample(url_to_code[url1], url_to_code[url2], label, url1, url2))
+            idx += 1
+            if idx == data_num:
+                break
+    return data

configs.py

@@ -0,0 +1,137 @@
+import random
+import torch
+import logging
+import multiprocessing
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def add_args(parser):
+    parser.add_argument("--task", type=str, required=True,
+                        choices=['summarize', 'concode', 'translate', 'refine', 'defect', 'clone', 'multi_task'])
+    parser.add_argument("--sub_task", type=str, default='')
+    parser.add_argument("--lang", type=str, default='')
+    parser.add_argument("--eval_task", type=str, default='')
+    parser.add_argument("--model_type", default="codet5", type=str, choices=['roberta', 'bart', 'codet5'])
+    parser.add_argument("--add_lang_ids", action='store_true')
+    parser.add_argument("--data_num", default=-1, type=int)
+    parser.add_argument("--start_epoch", default=0, type=int)
+    parser.add_argument("--num_train_epochs", default=100, type=int)
+    parser.add_argument("--patience", default=5, type=int)
+    parser.add_argument("--cache_path", type=str, required=True)
+    parser.add_argument("--summary_dir", type=str, required=True)
+    parser.add_argument("--data_dir", type=str, required=True)
+    parser.add_argument("--res_dir", type=str, required=True)
+    parser.add_argument("--res_fn", type=str, default='')
+    parser.add_argument("--add_task_prefix", action='store_true', help="Whether to add task prefix for t5 and codet5")
+    parser.add_argument("--save_last_checkpoints", action='store_true')
+    parser.add_argument("--always_save_model", action='store_true')
+    parser.add_argument("--do_eval_bleu", action='store_true', help="Whether to evaluate bleu on dev set.")
+
+    ## Required parameters
+    parser.add_argument("--model_name_or_path", default="roberta-base", type=str,
+                        help="Path to pre-trained model: e.g. roberta-base")
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model: Should contain the .bin files")
+    ## Other parameters
+    parser.add_argument("--train_filename", default=None, type=str,
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default=None, type=str,
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default=None, type=str,
+                        help="The test filename. Should contain the .jsonl files for this task.")
+
+    parser.add_argument("--config_name", default="", type=str,
+                        help="Pretrained config name or path if not the same as model_name")
+    parser.add_argument("--tokenizer_name", default="roberta-base", type=str,
+                        help="Pretrained tokenizer name or path if not the same as model_name")
+    parser.add_argument("--max_source_length", default=64, type=int,
+                        help="The maximum total source sequence length after tokenization. Sequences longer "
+                             "than this will be truncated, sequences shorter will be padded.")
+    parser.add_argument("--max_target_length", default=32, type=int,
+                        help="The maximum total target sequence length after tokenization. Sequences longer "
+                             "than this will be truncated, sequences shorter will be padded.")
+
+    parser.add_argument("--do_train", action='store_true',
+                        help="Whether to run eval on the train set.")
+    parser.add_argument("--do_eval", action='store_true',
+                        help="Whether to run eval on the dev set.")
+    parser.add_argument("--do_test", action='store_true',
+                        help="Whether to run eval on the dev set.")
+    parser.add_argument("--do_lower_case", action='store_true',
+                        help="Set this flag if you are using an uncased model.")
+    parser.add_argument("--no_cuda", action='store_true',
+                        help="Avoid using CUDA when available")
+
+    parser.add_argument("--train_batch_size", default=8, type=int,
+                        help="Batch size per GPU/CPU for training.")
+    parser.add_argument("--eval_batch_size", default=8, type=int,
+                        help="Batch size per GPU/CPU for evaluation.")
+    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
+                        help="Number of updates steps to accumulate before performing a backward/update pass.")
+    parser.add_argument("--learning_rate", default=5e-5, type=float,
+                        help="The initial learning rate for Adam.")
+    parser.add_argument("--beam_size", default=10, type=int,
+                        help="beam size for beam search")
+    parser.add_argument("--weight_decay", default=0.0, type=float,
+                        help="Weight deay if we apply some.")
+    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
+                        help="Epsilon for Adam optimizer.")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float,
+                        help="Max gradient norm.")
+
+    parser.add_argument("--save_steps", default=-1, type=int, )
+    parser.add_argument("--log_steps", default=-1, type=int, )
+    parser.add_argument("--max_steps", default=-1, type=int,
+                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
+    parser.add_argument("--eval_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=100, type=int,
+                        help="Linear warmup over warmup_steps.")
+    parser.add_argument("--local_rank", type=int, default=-1,
+                        help="For distributed training: local_rank")
+    parser.add_argument('--seed', type=int, default=1234,
+                        help="random seed for initialization")
+    args = parser.parse_args()
+
+    if args.task in ['summarize']:
+        args.lang = args.sub_task
+    elif args.task in ['refine', 'concode', 'clone']:
+        args.lang = 'java'
+    elif args.task == 'defect':
+        args.lang = 'c'
+    elif args.task == 'translate':
+        args.lang = 'c_sharp' if args.sub_task == 'java-cs' else 'java'
+    return args
+
+
+def set_dist(args):
+    # Setup CUDA, GPU & distributed training
+    if args.local_rank == -1 or args.no_cuda:
+        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
+        args.n_gpu = torch.cuda.device_count()
+    else:
+        # Setup for distributed data parallel
+        torch.cuda.set_device(args.local_rank)
+        device = torch.device("cuda", args.local_rank)
+        torch.distributed.init_process_group(backend='nccl')
+        args.n_gpu = 1
+    cpu_cont = multiprocessing.cpu_count()
+    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, cpu count: %d",
+                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), cpu_cont)
+    args.device = device
+    args.cpu_cont = cpu_cont
+
+
+def set_seed(args):
+    """set random seed."""
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if args.n_gpu > 0:
+        torch.cuda.manual_seed_all(args.seed)

evaluator/CodeBLEU/bleu.py

@@ -0,0 +1,589 @@
+# -*- coding: utf-8 -*-
+# Natural Language Toolkit: BLEU Score
+#
+# Copyright (C) 2001-2020 NLTK Project
+# Authors: Chin Yee Lee, Hengfeng Li, Ruxin Hou, Calvin Tanujaya Lim
+# Contributors: Björn Mattsson, Dmitrijs Milajevs, Liling Tan
+# URL: <http://nltk.org/>
+# For license information, see LICENSE.TXT
+
+"""BLEU score implementation."""
+
+import math
+import sys
+from fractions import Fraction
+import warnings
+from collections import Counter
+
+from evaluator.CodeBLEU.utils import ngrams
+
+
+def sentence_bleu(
+    references,
+    hypothesis,
+    weights=(0.25, 0.25, 0.25, 0.25),
+    smoothing_function=None,
+    auto_reweigh=False,
+):
+    """
+    Calculate BLEU score (Bilingual Evaluation Understudy) from
+    Papineni, Kishore, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002.
+    "BLEU: a method for automatic evaluation of machine translation."
+    In Proceedings of ACL. http://www.aclweb.org/anthology/P02-1040.pdf
+    >>> hypothesis1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
+    ...               'ensures', 'that', 'the', 'military', 'always',
+    ...               'obeys', 'the', 'commands', 'of', 'the', 'party']
+    >>> hypothesis2 = ['It', 'is', 'to', 'insure', 'the', 'troops',
+    ...               'forever', 'hearing', 'the', 'activity', 'guidebook',
+    ...               'that', 'party', 'direct']
+    >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+    ...               'ensures', 'that', 'the', 'military', 'will', 'forever',
+    ...               'heed', 'Party', 'commands']
+    >>> reference2 = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+    ...               'guarantees', 'the', 'military', 'forces', 'always',
+    ...               'being', 'under', 'the', 'command', 'of', 'the',
+    ...               'Party']
+    >>> reference3 = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+    ...               'army', 'always', 'to', 'heed', 'the', 'directions',
+    ...               'of', 'the', 'party']
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis1) # doctest: +ELLIPSIS
+    0.5045...
+    If there is no ngrams overlap for any order of n-grams, BLEU returns the
+    value 0. This is because the precision for the order of n-grams without
+    overlap is 0, and the geometric mean in the final BLEU score computation
+    multiplies the 0 with the precision of other n-grams. This results in 0
+    (independently of the precision of the othe n-gram orders). The following
+    example has zero 3-gram and 4-gram overlaps:
+    >>> round(sentence_bleu([reference1, reference2, reference3], hypothesis2),4) # doctest: +ELLIPSIS
+    0.0
+    To avoid this harsh behaviour when no ngram overlaps are found a smoothing
+    function can be used.
+    >>> chencherry = SmoothingFunction()
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis2,
+    ...     smoothing_function=chencherry.method1) # doctest: +ELLIPSIS
+    0.0370...
+    The default BLEU calculates a score for up to 4-grams using uniform
+    weights (this is called BLEU-4). To evaluate your translations with
+    higher/lower order ngrams, use customized weights. E.g. when accounting
+    for up to 5-grams with uniform weights (this is called BLEU-5) use:
+    >>> weights = (1./5., 1./5., 1./5., 1./5., 1./5.)
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis1, weights) # doctest: +ELLIPSIS
+    0.3920...
+    :param references: reference sentences
+    :type references: list(list(str))
+    :param hypothesis: a hypothesis sentence
+    :type hypothesis: list(str)
+    :param weights: weights for unigrams, bigrams, trigrams and so on
+    :type weights: list(float)
+    :param smoothing_function:
+    :type smoothing_function: SmoothingFunction
+    :param auto_reweigh: Option to re-normalize the weights uniformly.
+    :type auto_reweigh: bool
+    :return: The sentence-level BLEU score.
+    :rtype: float
+    """
+    return corpus_bleu(
+        [references], [hypothesis], weights, smoothing_function, auto_reweigh
+    )
+
+
+def corpus_bleu(
+    list_of_references,
+    hypotheses,
+    weights=(0.25, 0.25, 0.25, 0.25),
+    smoothing_function=None,
+    auto_reweigh=False,
+):
+    """
+    Calculate a single corpus-level BLEU score (aka. system-level BLEU) for all
+    the hypotheses and their respective references.
+    Instead of averaging the sentence level BLEU scores (i.e. marco-average
+    precision), the original BLEU metric (Papineni et al. 2002) accounts for
+    the micro-average precision (i.e. summing the numerators and denominators
+    for each hypothesis-reference(s) pairs before the division).
+    >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
+    ...         'ensures', 'that', 'the', 'military', 'always',
+    ...         'obeys', 'the', 'commands', 'of', 'the', 'party']
+    >>> ref1a = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+    ...          'ensures', 'that', 'the', 'military', 'will', 'forever',
+    ...          'heed', 'Party', 'commands']
+    >>> ref1b = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+    ...          'guarantees', 'the', 'military', 'forces', 'always',
+    ...          'being', 'under', 'the', 'command', 'of', 'the', 'Party']
+    >>> ref1c = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+    ...          'army', 'always', 'to', 'heed', 'the', 'directions',
+    ...          'of', 'the', 'party']
+    >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
+    ...         'interested', 'in', 'world', 'history']
+    >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
+    ...          'because', 'he', 'read', 'the', 'book']
+    >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]
+    >>> hypotheses = [hyp1, hyp2]
+    >>> corpus_bleu(list_of_references, hypotheses) # doctest: +ELLIPSIS
+    0.5920...
+    The example below show that corpus_bleu() is different from averaging
+    sentence_bleu() for hypotheses
+    >>> score1 = sentence_bleu([ref1a, ref1b, ref1c], hyp1)
+    >>> score2 = sentence_bleu([ref2a], hyp2)
+    >>> (score1 + score2) / 2 # doctest: +ELLIPSIS
+    0.6223...
+    :param list_of_references: a corpus of lists of reference sentences, w.r.t. hypotheses
+    :type list_of_references: list(list(list(str)))
+    :param hypotheses: a list of hypothesis sentences
+    :type hypotheses: list(list(str))
+    :param weights: weights for unigrams, bigrams, trigrams and so on
+    :type weights: list(float)
+    :param smoothing_function:
+    :type smoothing_function: SmoothingFunction
+    :param auto_reweigh: Option to re-normalize the weights uniformly.
+    :type auto_reweigh: bool
+    :return: The corpus-level BLEU score.
+    :rtype: float
+    """
+    # Before proceeding to compute BLEU, perform sanity checks.
+
+    p_numerators = Counter()  # Key = ngram order, and value = no. of ngram matches.
+    p_denominators = Counter()  # Key = ngram order, and value = no. of ngram in ref.
+    hyp_lengths, ref_lengths = 0, 0
+
+    assert len(list_of_references) == len(hypotheses), (
+        "The number of hypotheses and their reference(s) should be the " "same "
+    )
+
+    # Iterate through each hypothesis and their corresponding references.
+    for references, hypothesis in zip(list_of_references, hypotheses):
+        # For each order of ngram, calculate the numerator and
+        # denominator for the corpus-level modified precision.
+        for i, _ in enumerate(weights, start=1):
+            p_i = modified_precision(references, hypothesis, i)
+            p_numerators[i] += p_i.numerator
+            p_denominators[i] += p_i.denominator
+
+        # Calculate the hypothesis length and the closest reference length.
+        # Adds them to the corpus-level hypothesis and reference counts.
+        hyp_len = len(hypothesis)
+        hyp_lengths += hyp_len
+        ref_lengths += closest_ref_length(references, hyp_len)
+
+    # Calculate corpus-level brevity penalty.
+    bp = brevity_penalty(ref_lengths, hyp_lengths)
+
+    # Uniformly re-weighting based on maximum hypothesis lengths if largest
+    # order of n-grams < 4 and weights is set at default.
+    if auto_reweigh:
+        if hyp_lengths < 4 and weights == (0.25, 0.25, 0.25, 0.25):
+            weights = (1 / hyp_lengths,) * hyp_lengths
+
+    # Collects the various precision values for the different ngram orders.
+    p_n = [
+        Fraction(p_numerators[i], p_denominators[i], _normalize=False)
+        for i, _ in enumerate(weights, start=1)
+    ]
+
+    # Returns 0 if there's no matching n-grams
+    # We only need to check for p_numerators[1] == 0, since if there's
+    # no unigrams, there won't be any higher order ngrams.
+    if p_numerators[1] == 0:
+        return 0
+
+    # If there's no smoothing, set use method0 from SmoothinFunction class.
+    if not smoothing_function:
+        smoothing_function = SmoothingFunction().method1
+    # Smoothen the modified precision.
+    # Note: smoothing_function() may convert values into floats;
+    #       it tries to retain the Fraction object as much as the
+    #       smoothing method allows.
+    p_n = smoothing_function(
+        p_n, references=references, hypothesis=hypothesis, hyp_len=hyp_lengths
+    )
+    s = (w_i * math.log(p_i) for w_i, p_i in zip(weights, p_n))
+    s = bp * math.exp(math.fsum(s))
+    return s
+
+
+def modified_precision(references, hypothesis, n):
+    """
+    Calculate modified ngram precision.
+    The normal precision method may lead to some wrong translations with
+    high-precision, e.g., the translation, in which a word of reference
+    repeats several times, has very high precision.
+    This function only returns the Fraction object that contains the numerator
+    and denominator necessary to calculate the corpus-level precision.
+    To calculate the modified precision for a single pair of hypothesis and
+    references, cast the Fraction object into a float.
+    The famous "the the the ... " example shows that you can get BLEU precision
+    by duplicating high frequency words.
+        >>> reference1 = 'the cat is on the mat'.split()
+        >>> reference2 = 'there is a cat on the mat'.split()
+        >>> hypothesis1 = 'the the the the the the the'.split()
+        >>> references = [reference1, reference2]
+        >>> float(modified_precision(references, hypothesis1, n=1)) # doctest: +ELLIPSIS
+        0.2857...
+    In the modified n-gram precision, a reference word will be considered
+    exhausted after a matching hypothesis word is identified, e.g.
+        >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+        ...               'ensures', 'that', 'the', 'military', 'will',
+        ...               'forever', 'heed', 'Party', 'commands']
+        >>> reference2 = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+        ...               'guarantees', 'the', 'military', 'forces', 'always',
+        ...               'being', 'under', 'the', 'command', 'of', 'the',
+        ...               'Party']
+        >>> reference3 = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+        ...               'army', 'always', 'to', 'heed', 'the', 'directions',
+        ...               'of', 'the', 'party']
+        >>> hypothesis = 'of the'.split()
+        >>> references = [reference1, reference2, reference3]
+        >>> float(modified_precision(references, hypothesis, n=1))
+        1.0
+        >>> float(modified_precision(references, hypothesis, n=2))
+        1.0
+    An example of a normal machine translation hypothesis:
+        >>> hypothesis1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
+        ...               'ensures', 'that', 'the', 'military', 'always',
+        ...               'obeys', 'the', 'commands', 'of', 'the', 'party']
+        >>> hypothesis2 = ['It', 'is', 'to', 'insure', 'the', 'troops',
+        ...               'forever', 'hearing', 'the', 'activity', 'guidebook',
+        ...               'that', 'party', 'direct']
+        >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+        ...               'ensures', 'that', 'the', 'military', 'will',
+        ...               'forever', 'heed', 'Party', 'commands']
+        >>> reference2 = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+        ...               'guarantees', 'the', 'military', 'forces', 'always',
+        ...               'being', 'under', 'the', 'command', 'of', 'the',
+        ...               'Party']
+        >>> reference3 = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+        ...               'army', 'always', 'to', 'heed', 'the', 'directions',
+        ...               'of', 'the', 'party']
+        >>> references = [reference1, reference2, reference3]
+        >>> float(modified_precision(references, hypothesis1, n=1)) # doctest: +ELLIPSIS
+        0.9444...
+        >>> float(modified_precision(references, hypothesis2, n=1)) # doctest: +ELLIPSIS
+        0.5714...
+        >>> float(modified_precision(references, hypothesis1, n=2)) # doctest: +ELLIPSIS
+        0.5882352941176471
+        >>> float(modified_precision(references, hypothesis2, n=2)) # doctest: +ELLIPSIS
+        0.07692...
+    :param references: A list of reference translations.
+    :type references: list(list(str))
+    :param hypothesis: A hypothesis translation.
+    :type hypothesis: list(str)
+    :param n: The ngram order.
+    :type n: int
+    :return: BLEU's modified precision for the nth order ngram.
+    :rtype: Fraction
+    """
+    # Extracts all ngrams in hypothesis
+    # Set an empty Counter if hypothesis is empty.
+
+    counts = Counter(ngrams(hypothesis, n)) if len(hypothesis) >= n else Counter()
+    # Extract a union of references' counts.
+    # max_counts = reduce(or_, [Counter(ngrams(ref, n)) for ref in references])
+    max_counts = {}
+    for reference in references:
+        reference_counts = (
+            Counter(ngrams(reference, n)) if len(reference) >= n else Counter()
+        )
+        for ngram in counts:
+            max_counts[ngram] = max(max_counts.get(ngram, 0), reference_counts[ngram])
+
+    # Assigns the intersection between hypothesis and references' counts.
+    clipped_counts = {
+        ngram: min(count, max_counts[ngram]) for ngram, count in counts.items()
+    }
+
+    numerator = sum(clipped_counts.values())
+    # Ensures that denominator is minimum 1 to avoid ZeroDivisionError.
+    # Usually this happens when the ngram order is > len(reference).
+    denominator = max(1, sum(counts.values()))
+
+    return Fraction(numerator, denominator, _normalize=False)
+
+
+def closest_ref_length(references, hyp_len):
+    """
+    This function finds the reference that is the closest length to the
+    hypothesis. The closest reference length is referred to as *r* variable
+    from the brevity penalty formula in Papineni et. al. (2002)
+    :param references: A list of reference translations.
+    :type references: list(list(str))
+    :param hyp_len: The length of the hypothesis.
+    :type hyp_len: int
+    :return: The length of the reference that's closest to the hypothesis.
+    :rtype: int
+    """
+    ref_lens = (len(reference) for reference in references)
+    closest_ref_len = min(
+        ref_lens, key=lambda ref_len: (abs(ref_len - hyp_len), ref_len)
+    )
+    return closest_ref_len
+
+
+def brevity_penalty(closest_ref_len, hyp_len):
+    """
+    Calculate brevity penalty.
+    As the modified n-gram precision still has the problem from the short
+    length sentence, brevity penalty is used to modify the overall BLEU
+    score according to length.
+    An example from the paper. There are three references with length 12, 15
+    and 17. And a concise hypothesis of the length 12. The brevity penalty is 1.
+        >>> reference1 = list('aaaaaaaaaaaa')      # i.e. ['a'] * 12
+        >>> reference2 = list('aaaaaaaaaaaaaaa')   # i.e. ['a'] * 15
+        >>> reference3 = list('aaaaaaaaaaaaaaaaa') # i.e. ['a'] * 17
+        >>> hypothesis = list('aaaaaaaaaaaa')      # i.e. ['a'] * 12
+        >>> references = [reference1, reference2, reference3]
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        1.0
+    In case a hypothesis translation is shorter than the references, penalty is
+    applied.
+        >>> references = [['a'] * 28, ['a'] * 28]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        0.2635971381157267
+    The length of the closest reference is used to compute the penalty. If the
+    length of a hypothesis is 12, and the reference lengths are 13 and 2, the
+    penalty is applied because the hypothesis length (12) is less then the
+    closest reference length (13).
+        >>> references = [['a'] * 13, ['a'] * 2]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len) # doctest: +ELLIPSIS
+        0.9200...
+    The brevity penalty doesn't depend on reference order. More importantly,
+    when two reference sentences are at the same distance, the shortest
+    reference sentence length is used.
+        >>> references = [['a'] * 13, ['a'] * 11]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> bp1 = brevity_penalty(closest_ref_len, hyp_len)
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(reversed(references), hyp_len)
+        >>> bp2 = brevity_penalty(closest_ref_len, hyp_len)
+        >>> bp1 == bp2 == 1
+        True
+    A test example from mteval-v13a.pl (starting from the line 705):
+        >>> references = [['a'] * 11, ['a'] * 8]
+        >>> hypothesis = ['a'] * 7
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len) # doctest: +ELLIPSIS
+        0.8668...
+        >>> references = [['a'] * 11, ['a'] * 8, ['a'] * 6, ['a'] * 7]
+        >>> hypothesis = ['a'] * 7
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        1.0
+    :param hyp_len: The length of the hypothesis for a single sentence OR the
+    sum of all the hypotheses' lengths for a corpus
+    :type hyp_len: int
+    :param closest_ref_len: The length of the closest reference for a single
+    hypothesis OR the sum of all the closest references for every hypotheses.
+    :type closest_ref_len: int
+    :return: BLEU's brevity penalty.
+    :rtype: float
+    """
+    if hyp_len > closest_ref_len:
+        return 1
+    # If hypothesis is empty, brevity penalty = 0 should result in BLEU = 0.0
+    elif hyp_len == 0:
+        return 0
+    else:
+        return math.exp(1 - closest_ref_len / hyp_len)
+
+
+class SmoothingFunction:
+    """
+    This is an implementation of the smoothing techniques
+    for segment-level BLEU scores that was presented in
+    Boxing Chen and Collin Cherry (2014) A Systematic Comparison of
+    Smoothing Techniques for Sentence-Level BLEU. In WMT14.
+    http://acl2014.org/acl2014/W14-33/pdf/W14-3346.pdf
+    """
+
+    def __init__(self, epsilon=0.1, alpha=5, k=5):
+        """
+        This will initialize the parameters required for the various smoothing
+        techniques, the default values are set to the numbers used in the
+        experiments from Chen and Cherry (2014).
+        >>> hypothesis1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which', 'ensures',
+        ...                 'that', 'the', 'military', 'always', 'obeys', 'the',
+        ...                 'commands', 'of', 'the', 'party']
+        >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that', 'ensures',
+        ...               'that', 'the', 'military', 'will', 'forever', 'heed',
+        ...               'Party', 'commands']
+        >>> chencherry = SmoothingFunction()
+        >>> print(sentence_bleu([reference1], hypothesis1)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method0)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method1)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method2)) # doctest: +ELLIPSIS
+        0.4489...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method3)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method4)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method5)) # doctest: +ELLIPSIS
+        0.4905...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method6)) # doctest: +ELLIPSIS
+        0.4135...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method7)) # doctest: +ELLIPSIS
+        0.4905...
+        :param epsilon: the epsilon value use in method 1
+        :type epsilon: float
+        :param alpha: the alpha value use in method 6
+        :type alpha: int
+        :param k: the k value use in method 4
+        :type k: int
+        """
+        self.epsilon = epsilon
+        self.alpha = alpha
+        self.k = k
+
+    def method0(self, p_n, *args, **kwargs):
+        """
+        No smoothing.
+        """
+        p_n_new = []
+        for i, p_i in enumerate(p_n):
+            if p_i.numerator != 0:
+                p_n_new.append(p_i)
+            else:
+                _msg = str(
+                    "\nThe hypothesis contains 0 counts of {}-gram overlaps.\n"
+                    "Therefore the BLEU score evaluates to 0, independently of\n"
+                    "how many N-gram overlaps of lower order it contains.\n"
+                    "Consider using lower n-gram order or use "
+                    "SmoothingFunction()"
+                ).format(i + 1)
+                warnings.warn(_msg)
+                # When numerator==0 where denonminator==0 or !=0, the result
+                # for the precision score should be equal to 0 or undefined.
+                # Due to BLEU geometric mean computation in logarithm space,
+                # we we need to take the return sys.float_info.min such that
+                # math.log(sys.float_info.min) returns a 0 precision score.
+                p_n_new.append(sys.float_info.min)
+        return p_n_new
+
+    def method1(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 1: Add *epsilon* counts to precision with 0 counts.
+        """
+        return [
+            (p_i.numerator + self.epsilon) / p_i.denominator
+            if p_i.numerator == 0
+            else p_i
+            for p_i in p_n
+        ]
+
+    def method2(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 2: Add 1 to both numerator and denominator from
+        Chin-Yew Lin and Franz Josef Och (2004) Automatic evaluation of
+        machine translation quality using longest common subsequence and
+        skip-bigram statistics. In ACL04.
+        """
+        return [
+            Fraction(p_i.numerator + 1, p_i.denominator + 1, _normalize=False)
+            for p_i in p_n
+        ]
+
+    def method3(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 3: NIST geometric sequence smoothing
+        The smoothing is computed by taking 1 / ( 2^k ), instead of 0, for each
+        precision score whose matching n-gram count is null.
+        k is 1 for the first 'n' value for which the n-gram match count is null/
+        For example, if the text contains:
+         - one 2-gram match
+         - and (consequently) two 1-gram matches
+        the n-gram count for each individual precision score would be:
+         - n=1  =>  prec_count = 2     (two unigrams)
+         - n=2  =>  prec_count = 1     (one bigram)
+         - n=3  =>  prec_count = 1/2   (no trigram,  taking 'smoothed' value of 1 / ( 2^k ), with k=1)
+         - n=4  =>  prec_count = 1/4   (no fourgram, taking 'smoothed' value of 1 / ( 2^k ), with k=2)
+        """
+        incvnt = 1  # From the mteval-v13a.pl, it's referred to as k.
+        for i, p_i in enumerate(p_n):
+            if p_i.numerator == 0:
+                p_n[i] = 1 / (2 ** incvnt * p_i.denominator)
+                incvnt += 1
+        return p_n
+
+    def method4(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 4:
+        Shorter translations may have inflated precision values due to having
+        smaller denominators; therefore, we give them proportionally
+        smaller smoothed counts. Instead of scaling to 1/(2^k), Chen and Cherry
+        suggests dividing by 1/ln(len(T)), where T is the length of the translation.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        for i, p_i in enumerate(p_n):
+            if p_i.numerator == 0 and hyp_len != 0:
+                incvnt = i + 1 * self.k / math.log(
+                    hyp_len
+                )  # Note that this K is different from the K from NIST.
+                p_n[i] = incvnt / p_i.denominator
+        return p_n
+
+    def method5(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 5:
+        The matched counts for similar values of n should be similar. To a
+        calculate the n-gram matched count, it averages the n−1, n and n+1 gram
+        matched counts.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        m = {}
+        # Requires an precision value for an addition ngram order.
+        p_n_plus1 = p_n + [modified_precision(references, hypothesis, 5)]
+        m[-1] = p_n[0] + 1
+        for i, p_i in enumerate(p_n):
+            p_n[i] = (m[i - 1] + p_i + p_n_plus1[i + 1]) / 3
+            m[i] = p_n[i]
+        return p_n
+
+    def method6(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 6:
+        Interpolates the maximum likelihood estimate of the precision *p_n* with
+        a prior estimate *pi0*. The prior is estimated by assuming that the ratio
+        between pn and pn−1 will be the same as that between pn−1 and pn−2; from
+        Gao and He (2013) Training MRF-Based Phrase Translation Models using
+        Gradient Ascent. In NAACL.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        # This smoothing only works when p_1 and p_2 is non-zero.
+        # Raise an error with an appropriate message when the input is too short
+        # to use this smoothing technique.
+        assert p_n[2], "This smoothing method requires non-zero precision for bigrams."
+        for i, p_i in enumerate(p_n):
+            if i in [0, 1]:  # Skips the first 2 orders of ngrams.
+                continue
+            else:
+                pi0 = 0 if p_n[i - 2] == 0 else p_n[i - 1] ** 2 / p_n[i - 2]
+                # No. of ngrams in translation that matches the reference.
+                m = p_i.numerator
+                # No. of ngrams in translation.
+                l = sum(1 for _ in ngrams(hypothesis, i + 1))
+                # Calculates the interpolated precision.
+                p_n[i] = (m + self.alpha * pi0) / (l + self.alpha)
+        return p_n
+
+    def method7(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 7:
+        Interpolates methods 4 and 5.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        p_n = self.method4(p_n, references, hypothesis, hyp_len)
+        p_n = self.method5(p_n, references, hypothesis, hyp_len)
+        return p_n

evaluator/CodeBLEU/calc_code_bleu.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+# https://github.com/microsoft/CodeXGLUE/tree/main/Code-Code/code-to-code-trans/evaluator/CodeBLEU
+
+# -*- coding:utf-8 -*-
+import argparse
+import os
+from evaluator.CodeBLEU import bleu, weighted_ngram_match, syntax_match, dataflow_match
+
+
+def get_codebleu(refs, hyp, lang, params='0.25,0.25,0.25,0.25'):
+    if not isinstance(refs, list):
+        refs = [refs]
+    alpha, beta, gamma, theta = [float(x) for x in params.split(',')]
+
+    # preprocess inputs
+    pre_references = [[x.strip() for x in open(file, 'r', encoding='utf-8').readlines()] for file in refs]
+    hypothesis = [x.strip() for x in open(hyp, '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 = bleu.corpus_bleu(tokenized_refs, tokenized_hyps)
+
+    # calculate weighted ngram match
+    root_dir = os.path.dirname(__file__)
+    keywords = [x.strip() for x in open(root_dir + '/keywords/' + lang + '.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 = weighted_ngram_match.corpus_bleu(tokenized_refs_with_weights, tokenized_hyps)
+
+    # calculate syntax match
+    syntax_match_score = syntax_match.corpus_syntax_match(references, hypothesis, lang)
+
+    # calculate dataflow match
+    dataflow_match_score = dataflow_match.corpus_dataflow_match(references, hypothesis, lang)
+
+    print('ngram match: {0}, weighted ngram match: {1}, syntax_match: {2}, dataflow_match: {3}'. \
+          format(ngram_match_score, weighted_ngram_match_score, syntax_match_score, dataflow_match_score))
+
+    code_bleu_score = alpha * ngram_match_score \
+                      + beta * weighted_ngram_match_score \
+                      + gamma * syntax_match_score \
+                      + theta * dataflow_match_score
+
+    return code_bleu_score
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--refs', type=str, nargs='+', required=True,
+                        help='reference files')
+    parser.add_argument('--hyp', type=str, required=True,
+                        help='hypothesis file')
+    parser.add_argument('--lang', type=str, required=True,
+                        choices=['java', 'js', 'c_sharp', 'php', 'go', 'python', 'ruby'],
+                        help='programming language')
+    parser.add_argument('--params', type=str, default='0.25,0.25,0.25,0.25',
+                        help='alpha, beta and gamma')
+
+    args = parser.parse_args()
+    code_bleu_score = get_codebleu(args.refs, args.hyp, args.lang, args.params)
+    print('CodeBLEU score: ', code_bleu_score)
+

evaluator/CodeBLEU/dataflow_match.py

@@ -0,0 +1,149 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from evaluator.CodeBLEU.parser import DFG_python, DFG_java, DFG_ruby, DFG_go, DFG_php, DFG_javascript, DFG_csharp
+from evaluator.CodeBLEU.parser import (remove_comments_and_docstrings,
+                                       tree_to_token_index,
+                                       index_to_code_token,
+                                       tree_to_variable_index)
+from tree_sitter import Language, Parser
+import os
+
+root_dir = os.path.dirname(__file__)
+
+dfg_function = {
+    'python': DFG_python,
+    'java': DFG_java,
+    'ruby': DFG_ruby,
+    'go': DFG_go,
+    'php': DFG_php,
+    'javascript': DFG_javascript,
+    'c_sharp': DFG_csharp,
+}
+
+
+def calc_dataflow_match(references, candidate, lang):
+    return corpus_dataflow_match([references], [candidate], lang)
+
+
+def corpus_dataflow_match(references, candidates, lang):
+    LANGUAGE = Language(root_dir + '/parser/my-languages.so', lang)
+    parser = Parser()
+    parser.set_language(LANGUAGE)
+    parser = [parser, dfg_function[lang]]
+    match_count = 0
+    total_count = 0
+
+    for i in range(len(candidates)):
+        references_sample = references[i]
+        candidate = candidates[i]
+        for reference in references_sample:
+            try:
+                candidate = remove_comments_and_docstrings(candidate, 'java')
+            except:
+                pass
+            try:
+                reference = remove_comments_and_docstrings(reference, 'java')
+            except:
+                pass
+
+            cand_dfg = get_data_flow(candidate, parser)
+            ref_dfg = get_data_flow(reference, parser)
+
+            normalized_cand_dfg = normalize_dataflow(cand_dfg)
+            normalized_ref_dfg = normalize_dataflow(ref_dfg)
+
+            if len(normalized_ref_dfg) > 0:
+                total_count += len(normalized_ref_dfg)
+                for dataflow in normalized_ref_dfg:
+                    if dataflow in normalized_cand_dfg:
+                        match_count += 1
+                        normalized_cand_dfg.remove(dataflow)
+    if total_count == 0:
+        print(
+            "WARNING: There is no reference data-flows extracted from the whole corpus, and the data-flow match score degenerates to 0. Please consider ignoring this score.")
+        return 0
+    score = match_count / total_count
+    return score
+
+
+def get_data_flow(code, parser):
+    try:
+        tree = parser[0].parse(bytes(code, 'utf8'))
+        root_node = tree.root_node
+        tokens_index = tree_to_token_index(root_node)
+        code = code.split('\n')
+        code_tokens = [index_to_code_token(x, code) for x in tokens_index]
+        index_to_code = {}
+        for idx, (index, code) in enumerate(zip(tokens_index, code_tokens)):
+            index_to_code[index] = (idx, code)
+        try:
+            DFG, _ = parser[1](root_node, index_to_code, {})
+        except:
+            DFG = []
+        DFG = sorted(DFG, key=lambda x: x[1])
+        indexs = set()
+        for d in DFG:
+            if len(d[-1]) != 0:
+                indexs.add(d[1])
+            for x in d[-1]:
+                indexs.add(x)
+        new_DFG = []
+        for d in DFG:
+            if d[1] in indexs:
+                new_DFG.append(d)
+        codes = code_tokens
+        dfg = new_DFG
+    except:
+        codes = code.split()
+        dfg = []
+    # merge nodes
+    dic = {}
+    for d in dfg:
+        if d[1] not in dic:
+            dic[d[1]] = d
+        else:
+            dic[d[1]] = (d[0], d[1], d[2], list(set(dic[d[1]][3] + d[3])), list(set(dic[d[1]][4] + d[4])))
+    DFG = []
+    for d in dic:
+        DFG.append(dic[d])
+    dfg = DFG
+    return dfg
+
+
+def normalize_dataflow_item(dataflow_item):
+    var_name = dataflow_item[0]
+    var_pos = dataflow_item[1]
+    relationship = dataflow_item[2]
+    par_vars_name_list = dataflow_item[3]
+    par_vars_pos_list = dataflow_item[4]
+
+    var_names = list(set(par_vars_name_list + [var_name]))
+    norm_names = {}
+    for i in range(len(var_names)):
+        norm_names[var_names[i]] = 'var_' + str(i)
+
+    norm_var_name = norm_names[var_name]
+    relationship = dataflow_item[2]
+    norm_par_vars_name_list = [norm_names[x] for x in par_vars_name_list]
+
+    return (norm_var_name, relationship, norm_par_vars_name_list)
+
+
+def normalize_dataflow(dataflow):
+    var_dict = {}
+    i = 0
+    normalized_dataflow = []
+    for item in dataflow:
+        var_name = item[0]
+        relationship = item[2]
+        par_vars_name_list = item[3]
+        for name in par_vars_name_list:
+            if name not in var_dict:
+                var_dict[name] = 'var_' + str(i)
+                i += 1
+        if var_name not in var_dict:
+            var_dict[var_name] = 'var_' + str(i)
+            i += 1
+        normalized_dataflow.append((var_dict[var_name], relationship, [var_dict[x] for x in par_vars_name_list]))
+    return normalized_dataflow

evaluator/CodeBLEU/keywords/c_sharp.txt

@@ -0,0 +1,107 @@
+abstract
+as
+base
+bool
+break
+byte
+case
+catch
+char
+checked
+class
+const
+continue
+decimal
+default
+delegate
+do
+double
+else
+enum
+event
+explicit
+extern
+false
+finally
+fixed
+float
+for
+foreach
+goto
+if
+implicit
+in
+int
+interface
+internal
+is
+lock
+long
+namespace
+new
+null
+object
+operator
+out
+override
+params
+private
+protected
+public
+readonly
+ref
+return
+sbyte
+sealed
+short
+sizeof
+stackalloc
+static
+string
+struct
+switch
+this
+throw
+true
+try
+typeof
+uint
+ulong
+unchecked
+unsafe
+ushort
+using
+virtual
+void
+volatile
+while
+add
+alias
+ascending
+async
+await
+by
+descending
+dynamic
+equals
+from
+get
+global
+group
+into
+join
+let
+nameof
+notnull
+on
+orderby
+partial
+remove
+select
+set
+unmanaged
+value
+var
+when
+where
+yield

evaluator/CodeBLEU/keywords/java.txt

@@ -0,0 +1,50 @@
+abstract
+assert
+boolean
+break
+byte
+case
+catch
+char
+class
+const
+continue
+default
+do
+double
+else
+enum
+extends
+final
+finally
+float
+for
+goto
+if
+implements
+import
+instanceof
+int
+interface
+long
+native
+new
+package
+private
+protected
+public
+return
+short
+static
+strictfp
+super
+switch
+synchronized
+this
+throw
+throws
+transient
+try
+void
+volatile
+while

evaluator/CodeBLEU/parser/DFG.py

@@ -0,0 +1,1184 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from tree_sitter import Language, Parser
+from .utils import (remove_comments_and_docstrings,
+                   tree_to_token_index,
+                   index_to_code_token,
+                   tree_to_variable_index)
+
+
+def DFG_python(root_node,index_to_code,states):
+    assignment=['assignment','augmented_assignment','for_in_clause']
+    if_statement=['if_statement']
+    for_statement=['for_statement']
+    while_statement=['while_statement']
+    do_first_statement=['for_in_clause'] 
+    def_statement=['default_parameter']
+    states=states.copy() 
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':        
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_python(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states        
+    elif root_node.type in assignment:
+        if root_node.type=='for_in_clause':
+            right_nodes=[root_node.children[-1]]
+            left_nodes=[root_node.child_by_field_name('left')]
+        else:
+            if root_node.child_by_field_name('right') is None:
+                return [],states
+            left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+            right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+            if len(right_nodes)!=len(left_nodes):
+                left_nodes=[root_node.child_by_field_name('left')]
+                right_nodes=[root_node.child_by_field_name('right')]
+            if len(left_nodes)==0:
+                left_nodes=[root_node.child_by_field_name('left')]
+            if len(right_nodes)==0:
+                right_nodes=[root_node.child_by_field_name('right')]
+        DFG=[]
+        for node in right_nodes:
+            temp,states=DFG_python(node,index_to_code,states)
+            DFG+=temp
+            
+        for left_node,right_node in zip(left_nodes,right_nodes):
+            left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+            right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+            temp=[]
+            for token1_index in left_tokens_index:
+                idx1,code1=index_to_code[token1_index]
+                temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                             [index_to_code[x][0] for x in right_tokens_index]))
+                states[code1]=[idx1]
+            DFG+=temp        
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in ['elif_clause','else_clause']:
+                temp,current_states=DFG_python(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                temp,new_states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for i in range(2):
+            right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+            left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+            if len(right_nodes)!=len(left_nodes):
+                left_nodes=[root_node.child_by_field_name('left')]
+                right_nodes=[root_node.child_by_field_name('right')]
+            if len(left_nodes)==0:
+                left_nodes=[root_node.child_by_field_name('left')]
+            if len(right_nodes)==0:
+                right_nodes=[root_node.child_by_field_name('right')]
+            for node in right_nodes:
+                temp,states=DFG_python(node,index_to_code,states)
+                DFG+=temp
+            for left_node,right_node in zip(left_nodes,right_nodes):
+                left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+                right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+                temp=[]
+                for token1_index in left_tokens_index:
+                    idx1,code1=index_to_code[token1_index]
+                    temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                                 [index_to_code[x][0] for x in right_tokens_index]))
+                    states[code1]=[idx1]
+                DFG+=temp   
+            if  root_node.children[-1].type=="block":
+                temp,states=DFG_python(root_node.children[-1],index_to_code,states)
+                DFG+=temp 
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+        
+
+def DFG_java(root_node,index_to_code,states):
+    assignment=['assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=['enhanced_for_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_java(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_java(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_java(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_java(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="local_variable_declaration":
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_java(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_java(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+def DFG_csharp(root_node,index_to_code,states):
+    assignment=['assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['postfix_unary_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=['for_each_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        if len(root_node.children)==2:
+            name=root_node.children[0]
+            value=root_node.children[1]
+        else:
+            name=root_node.children[0]
+            value=None
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_csharp(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_csharp(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_csharp(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_csharp(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="local_variable_declaration":
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=root_node.child_by_field_name('left')
+        value=root_node.child_by_field_name('right')
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_csharp(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_csharp(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+    
+    
+def DFG_ruby(root_node,index_to_code,states):
+    assignment=['assignment','operator_assignment']
+    if_statement=['if','elsif','else','unless','when']
+    for_statement=['for']
+    while_statement=['while_modifier','until']
+    do_first_statement=[] 
+    def_statement=['keyword_parameter']
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        states=states.copy()
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_ruby(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states        
+    elif root_node.type in assignment:
+        left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+        right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+        if len(right_nodes)!=len(left_nodes):
+            left_nodes=[root_node.child_by_field_name('left')]
+            right_nodes=[root_node.child_by_field_name('right')]
+        if len(left_nodes)==0:
+            left_nodes=[root_node.child_by_field_name('left')]
+        if len(right_nodes)==0:
+            right_nodes=[root_node.child_by_field_name('right')]
+        if root_node.type=="operator_assignment":
+            left_nodes=[root_node.children[0]]
+            right_nodes=[root_node.children[-1]]
+
+        DFG=[]
+        for node in right_nodes:
+            temp,states=DFG_ruby(node,index_to_code,states)
+            DFG+=temp
+            
+        for left_node,right_node in zip(left_nodes,right_nodes):
+            left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+            right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+            temp=[]
+            for token1_index in left_tokens_index:
+                idx1,code1=index_to_code[token1_index]
+                temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                             [index_to_code[x][0] for x in right_tokens_index]))
+                states[code1]=[idx1]
+            DFG+=temp        
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement:
+                temp,current_states=DFG_ruby(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                temp,new_states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for i in range(2):
+            left_nodes=[root_node.child_by_field_name('pattern')]
+            right_nodes=[root_node.child_by_field_name('value')]
+            assert len(right_nodes)==len(left_nodes)
+            for node in right_nodes:
+                temp,states=DFG_ruby(node,index_to_code,states)
+                DFG+=temp
+            for left_node,right_node in zip(left_nodes,right_nodes):
+                left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+                right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+                temp=[]
+                for token1_index in left_tokens_index:
+                    idx1,code1=index_to_code[token1_index]
+                    temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                                 [index_to_code[x][0] for x in right_tokens_index]))
+                    states[code1]=[idx1]
+                DFG+=temp 
+            temp,states=DFG_ruby(root_node.child_by_field_name('body'),index_to_code,states)
+            DFG+=temp 
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+def DFG_go(root_node,index_to_code,states):
+    assignment=['assignment_statement',]
+    def_statement=['var_spec']
+    increment_statement=['inc_statement']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=[]
+    while_statement=[]
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_go(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_go(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_go(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_go(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="for_clause":
+                if child.child_by_field_name('update') is not None:
+                    temp,states=DFG_go(child.child_by_field_name('update'),index_to_code,states)
+                    DFG+=temp                 
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+    
+    
+
+def DFG_php(root_node,index_to_code,states):
+    assignment=['assignment_expression','augmented_assignment_expression']
+    def_statement=['simple_parameter']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else_clause']
+    for_statement=['for_statement']
+    enhanced_for_statement=['foreach_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('default_value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_php(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_php(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_php(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_php(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="assignment_expression":               
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=None
+        value=None
+        for child in root_node.children:
+            if child.type=='variable_name' and value is None:
+                value=child
+            elif child.type=='variable_name' and name is None:
+                name=child
+                break
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_php(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_php(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+def DFG_javascript(root_node,index_to_code,states):
+    assignment=['assignment_pattern','augmented_assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=[]
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type in ['string_literal','string','character_literal']) and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_javascript(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_javascript(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_javascript(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)        
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_javascript(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="variable_declaration":               
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states    
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+     

evaluator/CodeBLEU/parser/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from .utils import (remove_comments_and_docstrings,
+                   tree_to_token_index,
+                   index_to_code_token,
+                   tree_to_variable_index)
+from .DFG import DFG_python,DFG_java,DFG_ruby,DFG_go,DFG_php,DFG_javascript,DFG_csharp

evaluator/CodeBLEU/parser/build.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from tree_sitter import Language, Parser
+
+Language.build_library(
+  # Store the library in the `build` directory
+  'my-languages.so',
+
+  # Include one or more languages
+  [
+    'tree-sitter-go',
+    'tree-sitter-javascript',
+    'tree-sitter-python',
+    'tree-sitter-php',
+    'tree-sitter-java',
+    'tree-sitter-ruby',
+    'tree-sitter-c-sharp',
+  ]
+)
+

evaluator/CodeBLEU/parser/build.sh

@@ -0,0 +1,8 @@
+git clone https://github.com/tree-sitter/tree-sitter-go
+git clone https://github.com/tree-sitter/tree-sitter-javascript
+git clone https://github.com/tree-sitter/tree-sitter-python
+git clone https://github.com/tree-sitter/tree-sitter-ruby
+git clone https://github.com/tree-sitter/tree-sitter-php
+git clone https://github.com/tree-sitter/tree-sitter-java
+git clone https://github.com/tree-sitter/tree-sitter-c-sharp
+python build.py

evaluator/CodeBLEU/parser/my-languages.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66d01dcb2f38f3ff418839a10b856d4a5e2ef38f472c21ad7c6fb4bd14fc307d
+size 3000336

evaluator/CodeBLEU/parser/utils.py

@@ -0,0 +1,108 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import re
+from io import StringIO
+import tokenize
+
+
+def remove_comments_and_docstrings(source, lang):
+    if lang in ['python']:
+        """
+        Returns 'source' minus comments and docstrings.
+        """
+        io_obj = StringIO(source)
+        out = ""
+        prev_toktype = tokenize.INDENT
+        last_lineno = -1
+        last_col = 0
+        for tok in tokenize.generate_tokens(io_obj.readline):
+            token_type = tok[0]
+            token_string = tok[1]
+            start_line, start_col = tok[2]
+            end_line, end_col = tok[3]
+            ltext = tok[4]
+            if start_line > last_lineno:
+                last_col = 0
+            if start_col > last_col:
+                out += (" " * (start_col - last_col))
+            # Remove comments:
+            if token_type == tokenize.COMMENT:
+                pass
+            # This series of conditionals removes docstrings:
+            elif token_type == tokenize.STRING:
+                if prev_toktype != tokenize.INDENT:
+                    # This is likely a docstring; double-check we're not inside an operator:
+                    if prev_toktype != tokenize.NEWLINE:
+                        if start_col > 0:
+                            out += token_string
+            else:
+                out += token_string
+            prev_toktype = token_type
+            last_col = end_col
+            last_lineno = end_line
+        temp = []
+        for x in out.split('\n'):
+            if x.strip() != "":
+                temp.append(x)
+        return '\n'.join(temp)
+    elif lang in ['ruby']:
+        return source
+    else:
+        def replacer(match):
+            s = match.group(0)
+            if s.startswith('/'):
+                return " "  # note: a space and not an empty string
+            else:
+                return s
+
+        pattern = re.compile(
+            r'//.*?$|/\*.*?\*/|\'(?:\\.|[^\\\'])*\'|"(?:\\.|[^\\"])*"',
+            re.DOTALL | re.MULTILINE
+        )
+        temp = []
+        for x in re.sub(pattern, replacer, source).split('\n'):
+            if x.strip() != "":
+                temp.append(x)
+        return '\n'.join(temp)
+
+
+def tree_to_token_index(root_node):
+    if (len(root_node.children) == 0 or root_node.type in ['string_literal', 'string',
+                                                           'character_literal']) and root_node.type != 'comment':
+        return [(root_node.start_point, root_node.end_point)]
+    else:
+        code_tokens = []
+        for child in root_node.children:
+            code_tokens += tree_to_token_index(child)
+        return code_tokens
+
+
+def tree_to_variable_index(root_node, index_to_code):
+    if (len(root_node.children) == 0 or root_node.type in ['string_literal', 'string',
+                                                           'character_literal']) and root_node.type != 'comment':
+        index = (root_node.start_point, root_node.end_point)
+        _, code = index_to_code[index]
+        if root_node.type != code:
+            return [(root_node.start_point, root_node.end_point)]
+        else:
+            return []
+    else:
+        code_tokens = []
+        for child in root_node.children:
+            code_tokens += tree_to_variable_index(child, index_to_code)
+        return code_tokens
+
+
+def index_to_code_token(index, code):
+    start_point = index[0]
+    end_point = index[1]
+    if start_point[0] == end_point[0]:
+        s = code[start_point[0]][start_point[1]:end_point[1]]
+    else:
+        s = ""
+        s += code[start_point[0]][start_point[1]:]
+        for i in range(start_point[0] + 1, end_point[0]):
+            s += code[i]
+        s += code[end_point[0]][:end_point[1]]
+    return s

evaluator/CodeBLEU/readme.txt

@@ -0,0 +1 @@
+python calc_code_bleu.py --refs reference_files --hyp candidate_file --language java ( or c_sharp) --params 0.25,0.25,0.25,0.25(default)

evaluator/CodeBLEU/syntax_match.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from evaluator.CodeBLEU.parser import DFG_python, DFG_java, DFG_ruby, DFG_go, DFG_php, DFG_javascript, DFG_csharp
+from evaluator.CodeBLEU.parser import (remove_comments_and_docstrings,
+                                       tree_to_token_index,
+                                       index_to_code_token,
+                                       tree_to_variable_index)
+from tree_sitter import Language, Parser
+import os
+
+root_dir = os.path.dirname(__file__)
+dfg_function = {
+    'python': DFG_python,
+    'java': DFG_java,
+    'ruby': DFG_ruby,
+    'go': DFG_go,
+    'php': DFG_php,
+    'javascript': DFG_javascript,
+    'c_sharp': DFG_csharp,
+}
+
+
+def calc_syntax_match(references, candidate, lang):
+    return corpus_syntax_match([references], [candidate], lang)
+
+
+def corpus_syntax_match(references, candidates, lang):
+    JAVA_LANGUAGE = Language(root_dir + '/parser/my-languages.so', lang)
+    parser = Parser()
+    parser.set_language(JAVA_LANGUAGE)
+    match_count = 0
+    total_count = 0
+
+    for i in range(len(candidates)):
+        references_sample = references[i]
+        candidate = candidates[i]
+        for reference in references_sample:
+            try:
+                candidate = remove_comments_and_docstrings(candidate, 'java')
+            except:
+                pass
+            try:
+                reference = remove_comments_and_docstrings(reference, 'java')
+            except:
+                pass
+
+            candidate_tree = parser.parse(bytes(candidate, 'utf8')).root_node
+
+            reference_tree = parser.parse(bytes(reference, 'utf8')).root_node
+
+            def get_all_sub_trees(root_node):
+                node_stack = []
+                sub_tree_sexp_list = []
+                depth = 1
+                node_stack.append([root_node, depth])
+                while len(node_stack) != 0:
+                    cur_node, cur_depth = node_stack.pop()
+                    sub_tree_sexp_list.append([cur_node.sexp(), cur_depth])
+                    for child_node in cur_node.children:
+                        if len(child_node.children) != 0:
+                            depth = cur_depth + 1
+                            node_stack.append([child_node, depth])
+                return sub_tree_sexp_list
+
+            cand_sexps = [x[0] for x in get_all_sub_trees(candidate_tree)]
+            ref_sexps = get_all_sub_trees(reference_tree)
+
+            # print(cand_sexps)
+            # print(ref_sexps)
+
+            for sub_tree, depth in ref_sexps:
+                if sub_tree in cand_sexps:
+                    match_count += 1
+            total_count += len(ref_sexps)
+
+    score = match_count / total_count
+    return score

evaluator/CodeBLEU/utils.py

@@ -0,0 +1,106 @@
+# Natural Language Toolkit: Utility functions
+#
+# Copyright (C) 2001-2020 NLTK Project
+# Author: Steven Bird <[email protected]>
+# URL: <http://nltk.org/>
+# For license information, see LICENSE.TXT
+
+from itertools import chain
+
+def pad_sequence(
+    sequence,
+    n,
+    pad_left=False,
+    pad_right=False,
+    left_pad_symbol=None,
+    right_pad_symbol=None,
+):
+    """
+    Returns a padded sequence of items before ngram extraction.
+        >>> list(pad_sequence([1,2,3,4,5], 2, pad_left=True, pad_right=True, left_pad_symbol='<s>', right_pad_symbol='</s>'))
+        ['<s>', 1, 2, 3, 4, 5, '</s>']
+        >>> list(pad_sequence([1,2,3,4,5], 2, pad_left=True, left_pad_symbol='<s>'))
+        ['<s>', 1, 2, 3, 4, 5]
+        >>> list(pad_sequence([1,2,3,4,5], 2, pad_right=True, right_pad_symbol='</s>'))
+        [1, 2, 3, 4, 5, '</s>']
+    :param sequence: the source data to be padded
+    :type sequence: sequence or iter
+    :param n: the degree of the ngrams
+    :type n: int
+    :param pad_left: whether the ngrams should be left-padded
+    :type pad_left: bool
+    :param pad_right: whether the ngrams should be right-padded
+    :type pad_right: bool
+    :param left_pad_symbol: the symbol to use for left padding (default is None)
+    :type left_pad_symbol: any
+    :param right_pad_symbol: the symbol to use for right padding (default is None)
+    :type right_pad_symbol: any
+    :rtype: sequence or iter
+    """
+    sequence = iter(sequence)
+    if pad_left:
+        sequence = chain((left_pad_symbol,) * (n - 1), sequence)
+    if pad_right:
+        sequence = chain(sequence, (right_pad_symbol,) * (n - 1))
+    return sequence
+
+
+# add a flag to pad the sequence so we get peripheral ngrams?
+
+
+def ngrams(
+    sequence,
+    n,
+    pad_left=False,
+    pad_right=False,
+    left_pad_symbol=None,
+    right_pad_symbol=None,
+):
+    """
+    Return the ngrams generated from a sequence of items, as an iterator.
+    For example:
+        >>> from nltk.util import ngrams
+        >>> list(ngrams([1,2,3,4,5], 3))
+        [(1, 2, 3), (2, 3, 4), (3, 4, 5)]
+    Wrap with list for a list version of this function.  Set pad_left
+    or pad_right to true in order to get additional ngrams:
+        >>> list(ngrams([1,2,3,4,5], 2, pad_right=True))
+        [(1, 2), (2, 3), (3, 4), (4, 5), (5, None)]
+        >>> list(ngrams([1,2,3,4,5], 2, pad_right=True, right_pad_symbol='</s>'))
+        [(1, 2), (2, 3), (3, 4), (4, 5), (5, '</s>')]
+        >>> list(ngrams([1,2,3,4,5], 2, pad_left=True, left_pad_symbol='<s>'))
+        [('<s>', 1), (1, 2), (2, 3), (3, 4), (4, 5)]
+        >>> list(ngrams([1,2,3,4,5], 2, pad_left=True, pad_right=True, left_pad_symbol='<s>', right_pad_symbol='</s>'))
+        [('<s>', 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, '</s>')]
+    :param sequence: the source data to be converted into ngrams
+    :type sequence: sequence or iter
+    :param n: the degree of the ngrams
+    :type n: int
+    :param pad_left: whether the ngrams should be left-padded
+    :type pad_left: bool
+    :param pad_right: whether the ngrams should be right-padded
+    :type pad_right: bool
+    :param left_pad_symbol: the symbol to use for left padding (default is None)
+    :type left_pad_symbol: any
+    :param right_pad_symbol: the symbol to use for right padding (default is None)
+    :type right_pad_symbol: any
+    :rtype: sequence or iter
+    """
+    sequence = pad_sequence(
+        sequence, n, pad_left, pad_right, left_pad_symbol, right_pad_symbol
+    )
+
+    history = []
+    while n > 1:
+        # PEP 479, prevent RuntimeError from being raised when StopIteration bubbles out of generator
+        try:
+            next_item = next(sequence)
+        except StopIteration:
+            # no more data, terminate the generator
+            return
+        history.append(next_item)
+        n -= 1
+    for item in sequence:
+        history.append(item)
+        yield tuple(history)
+        del history[0]

evaluator/CodeBLEU/weighted_ngram_match.py

@@ -0,0 +1,558 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+# Natural Language Toolkit: BLEU Score
+#
+# Copyright (C) 2001-2020 NLTK Project
+# Authors: Chin Yee Lee, Hengfeng Li, Ruxin Hou, Calvin Tanujaya Lim
+# Contributors: Björn Mattsson, Dmitrijs Milajevs, Liling Tan
+# URL: <http://nltk.org/>
+# For license information, see LICENSE.TXT
+
+"""BLEU score implementation."""
+
+import math
+import sys
+from fractions import Fraction
+import warnings
+from collections import Counter
+
+from evaluator.CodeBLEU.utils import ngrams
+import pdb
+
+
+def sentence_bleu(
+    references,
+    hypothesis,
+    weights=(0.25, 0.25, 0.25, 0.25),
+    smoothing_function=None,
+    auto_reweigh=False,
+):
+    """
+    Calculate BLEU score (Bilingual Evaluation Understudy) from
+    Papineni, Kishore, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002.
+    "BLEU: a method for automatic evaluation of machine translation."
+    In Proceedings of ACL. http://www.aclweb.org/anthology/P02-1040.pdf
+    >>> hypothesis1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
+    ...               'ensures', 'that', 'the', 'military', 'always',
+    ...               'obeys', 'the', 'commands', 'of', 'the', 'party']
+    >>> hypothesis2 = ['It', 'is', 'to', 'insure', 'the', 'troops',
+    ...               'forever', 'hearing', 'the', 'activity', 'guidebook',
+    ...               'that', 'party', 'direct']
+    >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+    ...               'ensures', 'that', 'the', 'military', 'will', 'forever',
+    ...               'heed', 'Party', 'commands']
+    >>> reference2 = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+    ...               'guarantees', 'the', 'military', 'forces', 'always',
+    ...               'being', 'under', 'the', 'command', 'of', 'the',
+    ...               'Party']
+    >>> reference3 = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+    ...               'army', 'always', 'to', 'heed', 'the', 'directions',
+    ...               'of', 'the', 'party']
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis1) # doctest: +ELLIPSIS
+    0.5045...
+    If there is no ngrams overlap for any order of n-grams, BLEU returns the
+    value 0. This is because the precision for the order of n-grams without
+    overlap is 0, and the geometric mean in the final BLEU score computation
+    multiplies the 0 with the precision of other n-grams. This results in 0
+    (independently of the precision of the othe n-gram orders). The following
+    example has zero 3-gram and 4-gram overlaps:
+    >>> round(sentence_bleu([reference1, reference2, reference3], hypothesis2),4) # doctest: +ELLIPSIS
+    0.0
+    To avoid this harsh behaviour when no ngram overlaps are found a smoothing
+    function can be used.
+    >>> chencherry = SmoothingFunction()
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis2,
+    ...     smoothing_function=chencherry.method1) # doctest: +ELLIPSIS
+    0.0370...
+    The default BLEU calculates a score for up to 4-grams using uniform
+    weights (this is called BLEU-4). To evaluate your translations with
+    higher/lower order ngrams, use customized weights. E.g. when accounting
+    for up to 5-grams with uniform weights (this is called BLEU-5) use:
+    >>> weights = (1./5., 1./5., 1./5., 1./5., 1./5.)
+    >>> sentence_bleu([reference1, reference2, reference3], hypothesis1, weights) # doctest: +ELLIPSIS
+    0.3920...
+    :param references: reference sentences
+    :type references: list(list(str))
+    :param hypothesis: a hypothesis sentence
+    :type hypothesis: list(str)
+    :param weights: weights for unigrams, bigrams, trigrams and so on
+    :type weights: list(float)
+    :param smoothing_function:
+    :type smoothing_function: SmoothingFunction
+    :param auto_reweigh: Option to re-normalize the weights uniformly.
+    :type auto_reweigh: bool
+    :return: The sentence-level BLEU score.
+    :rtype: float
+    """
+    return corpus_bleu(
+        [references], [hypothesis], weights, smoothing_function, auto_reweigh
+    )
+
+
+def corpus_bleu(
+    list_of_references,
+    hypotheses,
+    weights=(0.25, 0.25, 0.25, 0.25),
+    smoothing_function=None,
+    auto_reweigh=False,
+):
+    """
+    Calculate a single corpus-level BLEU score (aka. system-level BLEU) for all
+    the hypotheses and their respective references.
+    Instead of averaging the sentence level BLEU scores (i.e. marco-average
+    precision), the original BLEU metric (Papineni et al. 2002) accounts for
+    the micro-average precision (i.e. summing the numerators and denominators
+    for each hypothesis-reference(s) pairs before the division).
+    >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
+    ...         'ensures', 'that', 'the', 'military', 'always',
+    ...         'obeys', 'the', 'commands', 'of', 'the', 'party']
+    >>> ref1a = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
+    ...          'ensures', 'that', 'the', 'military', 'will', 'forever',
+    ...          'heed', 'Party', 'commands']
+    >>> ref1b = ['It', 'is', 'the', 'guiding', 'principle', 'which',
+    ...          'guarantees', 'the', 'military', 'forces', 'always',
+    ...          'being', 'under', 'the', 'command', 'of', 'the', 'Party']
+    >>> ref1c = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
+    ...          'army', 'always', 'to', 'heed', 'the', 'directions',
+    ...          'of', 'the', 'party']
+    >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
+    ...         'interested', 'in', 'world', 'history']
+    >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
+    ...          'because', 'he', 'read', 'the', 'book']
+    >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]
+    >>> hypotheses = [hyp1, hyp2]
+    >>> corpus_bleu(list_of_references, hypotheses) # doctest: +ELLIPSIS
+    0.5920...
+    The example below show that corpus_bleu() is different from averaging
+    sentence_bleu() for hypotheses
+    >>> score1 = sentence_bleu([ref1a, ref1b, ref1c], hyp1)
+    >>> score2 = sentence_bleu([ref2a], hyp2)
+    >>> (score1 + score2) / 2 # doctest: +ELLIPSIS
+    0.6223...
+    :param list_of_references: a corpus of lists of reference sentences, w.r.t. hypotheses
+    :type list_of_references: list(list(list(str)))
+    :param hypotheses: a list of hypothesis sentences
+    :type hypotheses: list(list(str))
+    :param weights: weights for unigrams, bigrams, trigrams and so on
+    :type weights: list(float)
+    :param smoothing_function:
+    :type smoothing_function: SmoothingFunction
+    :param auto_reweigh: Option to re-normalize the weights uniformly.
+    :type auto_reweigh: bool
+    :return: The corpus-level BLEU score.
+    :rtype: float
+    """
+    # Before proceeding to compute BLEU, perform sanity checks.
+
+    p_numerators = Counter()  # Key = ngram order, and value = no. of ngram matches.
+    p_denominators = Counter()  # Key = ngram order, and value = no. of ngram in ref.
+    hyp_lengths, ref_lengths = 0, 0
+
+    assert len(list_of_references) == len(hypotheses), (
+        "The number of hypotheses and their reference(s) should be the " "same "
+    )
+
+    # Iterate through each hypothesis and their corresponding references.
+    for references, hypothesis in zip(list_of_references, hypotheses):
+        # For each order of ngram, calculate the numerator and
+        # denominator for the corpus-level modified precision.
+        for i, _ in enumerate(weights, start=1):
+            p_i_numeraotr, p_i_denominator = modified_recall(references, hypothesis, i)
+            p_numerators[i] += p_i_numeraotr
+            p_denominators[i] += p_i_denominator
+
+        # Calculate the hypothesis length and the closest reference length.
+        # Adds them to the corpus-level hypothesis and reference counts.
+        hyp_len = len(hypothesis)
+        hyp_lengths += hyp_len
+        ref_lengths += closest_ref_length(references, hyp_len)
+
+    # Calculate corpus-level brevity penalty.
+    bp = brevity_penalty(ref_lengths, hyp_lengths)
+
+    # Uniformly re-weighting based on maximum hypothesis lengths if largest
+    # order of n-grams < 4 and weights is set at default.
+    if auto_reweigh:
+        if hyp_lengths < 4 and weights == (0.25, 0.25, 0.25, 0.25):
+            weights = (1 / hyp_lengths,) * hyp_lengths
+
+    # Collects the various recall values for the different ngram orders.
+    p_n = [
+        (p_numerators[i], p_denominators[i])
+        for i, _ in enumerate(weights, start=1)
+    ]
+
+    # Returns 0 if there's no matching n-grams
+    # We only need to check for p_numerators[1] == 0, since if there's
+    # no unigrams, there won't be any higher order ngrams.
+    if p_numerators[1] == 0:
+        return 0
+
+    # If there's no smoothing, set use method0 from SmoothinFunction class.
+    if not smoothing_function:
+        smoothing_function = SmoothingFunction().method1
+    # Smoothen the modified precision.
+    # Note: smoothing_function() may convert values into floats;
+    #       it tries to retain the Fraction object as much as the
+    #       smoothing method allows.
+    p_n = smoothing_function(
+        p_n, references=references, hypothesis=hypothesis, hyp_len=hyp_lengths
+    )
+    # pdb.set_trace()
+    s = (w_i * math.log(p_i[0]/p_i[1]) for w_i, p_i in zip(weights, p_n))
+    s = bp * math.exp(math.fsum(s))
+    return s
+
+
+def modified_recall(references, hypothesis, n):
+    """
+    Calculate modified ngram recall.
+    :param references: A list of reference translations.
+    :type references: list(list(str))
+    :param hypothesis: A hypothesis translation.
+    :type hypothesis: list(str)
+    :param n: The ngram order.
+    :type n: int
+    :return: BLEU's modified precision for the nth order ngram.
+    :rtype: Fraction
+    """
+    # Extracts all ngrams in hypothesis
+    # Set an empty Counter if hypothesis is empty.
+    # pdb.set_trace()
+    numerator = 0
+    denominator = 0
+
+    counts = Counter(ngrams(hypothesis, n)) if len(hypothesis) >= n else Counter()
+    # Extract a union of references' counts.
+    # max_counts = reduce(or_, [Counter(ngrams(ref, n)) for ref in references])
+    max_counts = {}
+    for reference_and_weights in references:
+        reference = reference_and_weights[0]
+        weights = reference_and_weights[1]
+        reference_counts = (
+            Counter(ngrams(reference, n)) if len(reference) >= n else Counter()
+        )
+        # for ngram in reference_counts:
+        #     max_counts[ngram] = max(max_counts.get(ngram, 0), counts[ngram])
+        clipped_counts = {
+            ngram: min(count, counts[ngram]) for ngram, count in reference_counts.items()
+        }
+        # reweight
+        if n == 1 and len(weights) == len(reference_counts):
+            def weighted_sum(weights, counts):
+                sum_counts = 0
+                for ngram, count in counts.items():
+                    sum_counts += count * (weights[ngram[0]] if ngram[0] in weights else 1)
+                return sum_counts
+
+            numerator += weighted_sum(weights, clipped_counts)
+            denominator += max(1, weighted_sum(weights, reference_counts))
+
+        else:
+            numerator += sum(clipped_counts.values())
+            denominator += max(1, sum(reference_counts.values()))
+
+        # # Assigns the intersection between hypothesis and references' counts.
+        # clipped_counts = {
+        #     ngram: min(count, max_counts[ngram]) for ngram, count in counts.items()
+        # }
+
+        # numerator += sum(clipped_counts.values())
+        # # Ensures that denominator is minimum 1 to avoid ZeroDivisionError.
+        # # Usually this happens when the ngram order is > len(reference).
+        # denominator += max(1, sum(counts.values()))
+
+    #return Fraction(numerator, denominator, _normalize=False)
+    return numerator, denominator
+
+
+def closest_ref_length(references, hyp_len):
+    """
+    This function finds the reference that is the closest length to the
+    hypothesis. The closest reference length is referred to as *r* variable
+    from the brevity penalty formula in Papineni et. al. (2002)
+    :param references: A list of reference translations.
+    :type references: list(list(str))
+    :param hyp_len: The length of the hypothesis.
+    :type hyp_len: int
+    :return: The length of the reference that's closest to the hypothesis.
+    :rtype: int
+    """
+    ref_lens = (len(reference) for reference in references)
+    closest_ref_len = min(
+        ref_lens, key=lambda ref_len: (abs(ref_len - hyp_len), ref_len)
+    )
+    return closest_ref_len
+
+
+def brevity_penalty(closest_ref_len, hyp_len):
+    """
+    Calculate brevity penalty.
+    As the modified n-gram precision still has the problem from the short
+    length sentence, brevity penalty is used to modify the overall BLEU
+    score according to length.
+    An example from the paper. There are three references with length 12, 15
+    and 17. And a concise hypothesis of the length 12. The brevity penalty is 1.
+        >>> reference1 = list('aaaaaaaaaaaa')      # i.e. ['a'] * 12
+        >>> reference2 = list('aaaaaaaaaaaaaaa')   # i.e. ['a'] * 15
+        >>> reference3 = list('aaaaaaaaaaaaaaaaa') # i.e. ['a'] * 17
+        >>> hypothesis = list('aaaaaaaaaaaa')      # i.e. ['a'] * 12
+        >>> references = [reference1, reference2, reference3]
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        1.0
+    In case a hypothesis translation is shorter than the references, penalty is
+    applied.
+        >>> references = [['a'] * 28, ['a'] * 28]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        0.2635971381157267
+    The length of the closest reference is used to compute the penalty. If the
+    length of a hypothesis is 12, and the reference lengths are 13 and 2, the
+    penalty is applied because the hypothesis length (12) is less then the
+    closest reference length (13).
+        >>> references = [['a'] * 13, ['a'] * 2]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len) # doctest: +ELLIPSIS
+        0.9200...
+    The brevity penalty doesn't depend on reference order. More importantly,
+    when two reference sentences are at the same distance, the shortest
+    reference sentence length is used.
+        >>> references = [['a'] * 13, ['a'] * 11]
+        >>> hypothesis = ['a'] * 12
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> bp1 = brevity_penalty(closest_ref_len, hyp_len)
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(reversed(references), hyp_len)
+        >>> bp2 = brevity_penalty(closest_ref_len, hyp_len)
+        >>> bp1 == bp2 == 1
+        True
+    A test example from mteval-v13a.pl (starting from the line 705):
+        >>> references = [['a'] * 11, ['a'] * 8]
+        >>> hypothesis = ['a'] * 7
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len) # doctest: +ELLIPSIS
+        0.8668...
+        >>> references = [['a'] * 11, ['a'] * 8, ['a'] * 6, ['a'] * 7]
+        >>> hypothesis = ['a'] * 7
+        >>> hyp_len = len(hypothesis)
+        >>> closest_ref_len =  closest_ref_length(references, hyp_len)
+        >>> brevity_penalty(closest_ref_len, hyp_len)
+        1.0
+    :param hyp_len: The length of the hypothesis for a single sentence OR the
+    sum of all the hypotheses' lengths for a corpus
+    :type hyp_len: int
+    :param closest_ref_len: The length of the closest reference for a single
+    hypothesis OR the sum of all the closest references for every hypotheses.
+    :type closest_ref_len: int
+    :return: BLEU's brevity penalty.
+    :rtype: float
+    """
+    if hyp_len > closest_ref_len:
+        return 1
+    # If hypothesis is empty, brevity penalty = 0 should result in BLEU = 0.0
+    elif hyp_len == 0:
+        return 0
+    else:
+        return math.exp(1 - closest_ref_len / hyp_len)
+
+
+class SmoothingFunction:
+    """
+    This is an implementation of the smoothing techniques
+    for segment-level BLEU scores that was presented in
+    Boxing Chen and Collin Cherry (2014) A Systematic Comparison of
+    Smoothing Techniques for Sentence-Level BLEU. In WMT14.
+    http://acl2014.org/acl2014/W14-33/pdf/W14-3346.pdf
+    """
+
+    def __init__(self, epsilon=0.1, alpha=5, k=5):
+        """
+        This will initialize the parameters required for the various smoothing
+        techniques, the default values are set to the numbers used in the
+        experiments from Chen and Cherry (2014).
+        >>> hypothesis1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which', 'ensures',
+        ...                 'that', 'the', 'military', 'always', 'obeys', 'the',
+        ...                 'commands', 'of', 'the', 'party']
+        >>> reference1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'that', 'ensures',
+        ...               'that', 'the', 'military', 'will', 'forever', 'heed',
+        ...               'Party', 'commands']
+        >>> chencherry = SmoothingFunction()
+        >>> print(sentence_bleu([reference1], hypothesis1)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method0)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method1)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method2)) # doctest: +ELLIPSIS
+        0.4489...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method3)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method4)) # doctest: +ELLIPSIS
+        0.4118...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method5)) # doctest: +ELLIPSIS
+        0.4905...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method6)) # doctest: +ELLIPSIS
+        0.4135...
+        >>> print(sentence_bleu([reference1], hypothesis1, smoothing_function=chencherry.method7)) # doctest: +ELLIPSIS
+        0.4905...
+        :param epsilon: the epsilon value use in method 1
+        :type epsilon: float
+        :param alpha: the alpha value use in method 6
+        :type alpha: int
+        :param k: the k value use in method 4
+        :type k: int
+        """
+        self.epsilon = epsilon
+        self.alpha = alpha
+        self.k = k
+
+    def method0(self, p_n, *args, **kwargs):
+        """
+        No smoothing.
+        """
+        p_n_new = []
+        for i, p_i in enumerate(p_n):
+            if p_i[0] != 0:
+                p_n_new.append(p_i)
+            else:
+                _msg = str(
+                    "\nThe hypothesis contains 0 counts of {}-gram overlaps.\n"
+                    "Therefore the BLEU score evaluates to 0, independently of\n"
+                    "how many N-gram overlaps of lower order it contains.\n"
+                    "Consider using lower n-gram order or use "
+                    "SmoothingFunction()"
+                ).format(i + 1)
+                warnings.warn(_msg)
+                # When numerator==0 where denonminator==0 or !=0, the result
+                # for the precision score should be equal to 0 or undefined.
+                # Due to BLEU geometric mean computation in logarithm space,
+                # we we need to take the return sys.float_info.min such that
+                # math.log(sys.float_info.min) returns a 0 precision score.
+                p_n_new.append(sys.float_info.min)
+        return p_n_new
+
+    def method1(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 1: Add *epsilon* counts to precision with 0 counts.
+        """
+        return [
+            ((p_i[0] + self.epsilon),  p_i[1])
+            if p_i[0] == 0
+            else p_i
+            for p_i in p_n
+        ]
+
+    def method2(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 2: Add 1 to both numerator and denominator from
+        Chin-Yew Lin and Franz Josef Och (2004) Automatic evaluation of
+        machine translation quality using longest common subsequence and
+        skip-bigram statistics. In ACL04.
+        """
+        return [
+            (p_i[0] + 1, p_i[1] + 1)
+            for p_i in p_n
+        ]
+
+    def method3(self, p_n, *args, **kwargs):
+        """
+        Smoothing method 3: NIST geometric sequence smoothing
+        The smoothing is computed by taking 1 / ( 2^k ), instead of 0, for each
+        precision score whose matching n-gram count is null.
+        k is 1 for the first 'n' value for which the n-gram match count is null/
+        For example, if the text contains:
+         - one 2-gram match
+         - and (consequently) two 1-gram matches
+        the n-gram count for each individual precision score would be:
+         - n=1  =>  prec_count = 2     (two unigrams)
+         - n=2  =>  prec_count = 1     (one bigram)
+         - n=3  =>  prec_count = 1/2   (no trigram,  taking 'smoothed' value of 1 / ( 2^k ), with k=1)
+         - n=4  =>  prec_count = 1/4   (no fourgram, taking 'smoothed' value of 1 / ( 2^k ), with k=2)
+        """
+        incvnt = 1  # From the mteval-v13a.pl, it's referred to as k.
+        for i, p_i in enumerate(p_n):
+            if p_i.numerator == 0:
+                p_n[i] = 1 / (2 ** incvnt * p_i.denominator)
+                incvnt += 1
+        return p_n
+
+    def method4(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 4:
+        Shorter translations may have inflated precision values due to having
+        smaller denominators; therefore, we give them proportionally
+        smaller smoothed counts. Instead of scaling to 1/(2^k), Chen and Cherry
+        suggests dividing by 1/ln(len(T)), where T is the length of the translation.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        for i, p_i in enumerate(p_n):
+            if p_i.numerator == 0 and hyp_len != 0:
+                incvnt = i + 1 * self.k / math.log(
+                    hyp_len
+                )  # Note that this K is different from the K from NIST.
+                p_n[i] = incvnt / p_i.denominator
+        return p_n
+
+    def method5(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 5:
+        The matched counts for similar values of n should be similar. To a
+        calculate the n-gram matched count, it averages the n−1, n and n+1 gram
+        matched counts.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        m = {}
+        # Requires an precision value for an addition ngram order.
+        p_n_plus1 = p_n + [modified_precision(references, hypothesis, 5)]
+        m[-1] = p_n[0] + 1
+        for i, p_i in enumerate(p_n):
+            p_n[i] = (m[i - 1] + p_i + p_n_plus1[i + 1]) / 3
+            m[i] = p_n[i]
+        return p_n
+
+    def method6(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 6:
+        Interpolates the maximum likelihood estimate of the precision *p_n* with
+        a prior estimate *pi0*. The prior is estimated by assuming that the ratio
+        between pn and pn−1 will be the same as that between pn−1 and pn−2; from
+        Gao and He (2013) Training MRF-Based Phrase Translation Models using
+        Gradient Ascent. In NAACL.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        # This smoothing only works when p_1 and p_2 is non-zero.
+        # Raise an error with an appropriate message when the input is too short
+        # to use this smoothing technique.
+        assert p_n[2], "This smoothing method requires non-zero precision for bigrams."
+        for i, p_i in enumerate(p_n):
+            if i in [0, 1]:  # Skips the first 2 orders of ngrams.
+                continue
+            else:
+                pi0 = 0 if p_n[i - 2] == 0 else p_n[i - 1] ** 2 / p_n[i - 2]
+                # No. of ngrams in translation that matches the reference.
+                m = p_i.numerator
+                # No. of ngrams in translation.
+                l = sum(1 for _ in ngrams(hypothesis, i + 1))
+                # Calculates the interpolated precision.
+                p_n[i] = (m + self.alpha * pi0) / (l + self.alpha)
+        return p_n
+
+    def method7(self, p_n, references, hypothesis, hyp_len=None, *args, **kwargs):
+        """
+        Smoothing method 7:
+        Interpolates methods 4 and 5.
+        """
+        hyp_len = hyp_len if hyp_len else len(hypothesis)
+        p_n = self.method4(p_n, references, hypothesis, hyp_len)
+        p_n = self.method5(p_n, references, hypothesis, hyp_len)
+        return p_n

evaluator/bleu.py

@@ -0,0 +1,134 @@
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+"""Python implementation of BLEU and smooth-BLEU.
+
+This module provides a Python implementation of BLEU and smooth-BLEU.
+Smooth BLEU is computed following the method outlined in the paper:
+Chin-Yew Lin, Franz Josef Och. ORANGE: a method for evaluating automatic
+evaluation metrics for machine translation. COLING 2004.
+"""
+
+import collections
+import math
+
+
+def _get_ngrams(segment, max_order):
+  """Extracts all n-grams upto a given maximum order from an input segment.
+
+  Args:
+    segment: text segment from which n-grams will be extracted.
+    max_order: maximum length in tokens of the n-grams returned by this
+        methods.
+
+  Returns:
+    The Counter containing all n-grams upto max_order in segment
+    with a count of how many times each n-gram occurred.
+  """
+  ngram_counts = collections.Counter()
+  for order in range(1, max_order + 1):
+    for i in range(0, len(segment) - order + 1):
+      ngram = tuple(segment[i:i+order])
+      ngram_counts[ngram] += 1
+  return ngram_counts
+
+
+def compute_bleu(reference_corpus, translation_corpus, max_order=4,
+                 smooth=False):
+  """Computes BLEU score of translated segments against one or more references.
+
+  Args:
+    reference_corpus: list of lists of references for each translation. Each
+        reference should be tokenized into a list of tokens.
+    translation_corpus: list of translations to score. Each translation
+        should be tokenized into a list of tokens.
+    max_order: Maximum n-gram order to use when computing BLEU score.
+    smooth: Whether or not to apply Lin et al. 2004 smoothing.
+
+  Returns:
+    3-Tuple with the BLEU score, n-gram precisions, geometric mean of n-gram
+    precisions and brevity penalty.
+  """
+  matches_by_order = [0] * max_order
+  possible_matches_by_order = [0] * max_order
+  reference_length = 0
+  translation_length = 0
+  for (references, translation) in zip(reference_corpus,
+                                       translation_corpus):
+    reference_length += min(len(r) for r in references)
+    translation_length += len(translation)
+
+    merged_ref_ngram_counts = collections.Counter()
+    for reference in references:
+      merged_ref_ngram_counts |= _get_ngrams(reference, max_order)
+    translation_ngram_counts = _get_ngrams(translation, max_order)
+    overlap = translation_ngram_counts & merged_ref_ngram_counts
+    for ngram in overlap:
+      matches_by_order[len(ngram)-1] += overlap[ngram]
+    for order in range(1, max_order+1):
+      possible_matches = len(translation) - order + 1
+      if possible_matches > 0:
+        possible_matches_by_order[order-1] += possible_matches
+
+  precisions = [0] * max_order
+  for i in range(0, max_order):
+    if smooth:
+      precisions[i] = ((matches_by_order[i] + 1.) /
+                       (possible_matches_by_order[i] + 1.))
+    else:
+      if possible_matches_by_order[i] > 0:
+        precisions[i] = (float(matches_by_order[i]) /
+                         possible_matches_by_order[i])
+      else:
+        precisions[i] = 0.0
+
+  if min(precisions) > 0:
+    p_log_sum = sum((1. / max_order) * math.log(p) for p in precisions)
+    geo_mean = math.exp(p_log_sum)
+  else:
+    geo_mean = 0
+
+  ratio = float(translation_length) / reference_length
+
+  if ratio > 1.0:
+    bp = 1.
+  else:
+    bp = math.exp(1 - 1. / ratio)
+
+  bleu = geo_mean * bp
+
+  return (bleu, precisions, bp, ratio, translation_length, reference_length)
+
+
+def _bleu(ref_file, trans_file, subword_option=None):
+    max_order = 4
+    smooth = True
+    ref_files = [ref_file]
+    reference_text = []
+    for reference_filename in ref_files:
+        with open(reference_filename) as fh:
+            reference_text.append(fh.readlines())
+    per_segment_references = []
+    for references in zip(*reference_text):
+        reference_list = []
+        for reference in references:
+            reference_list.append(reference.strip().split())
+        per_segment_references.append(reference_list)
+    translations = []
+    with open(trans_file) as fh:
+        for line in fh:
+            translations.append(line.strip().split())
+    bleu_score, _, _, _, _, _ = compute_bleu(per_segment_references, translations, max_order, smooth)
+    return round(100 * bleu_score,2)

evaluator/smooth_bleu.py

@@ -0,0 +1,208 @@
+#!/usr/bin/python
+
+'''
+This script was adapted from the original version by hieuhoang1972 which is part of MOSES. 
+'''
+
+# $Id: bleu.py 1307 2007-03-14 22:22:36Z hieuhoang1972 $
+
+'''Provides:
+
+cook_refs(refs, n=4): Transform a list of reference sentences as strings into a form usable by cook_test().
+cook_test(test, refs, n=4): Transform a test sentence as a string (together with the cooked reference sentences) into a form usable by score_cooked().
+score_cooked(alltest, n=4): Score a list of cooked test sentences.
+
+score_set(s, testid, refids, n=4): Interface with dataset.py; calculate BLEU score of testid against refids.
+
+The reason for breaking the BLEU computation into three phases cook_refs(), cook_test(), and score_cooked() is to allow the caller to calculate BLEU scores for multiple test sets as efficiently as possible.
+'''
+
+import sys, math, re, xml.sax.saxutils
+import subprocess
+import os
+
+# Added to bypass NIST-style pre-processing of hyp and ref files -- wade
+nonorm = 0
+
+preserve_case = False
+eff_ref_len = "shortest"
+
+normalize1 = [
+    ('<skipped>', ''),  # strip "skipped" tags
+    (r'-\n', ''),  # strip end-of-line hyphenation and join lines
+    (r'\n', ' '),  # join lines
+    #    (r'(\d)\s+(?=\d)', r'\1'), # join digits
+]
+normalize1 = [(re.compile(pattern), replace) for (pattern, replace) in normalize1]
+
+normalize2 = [
+    (r'([\{-\~\[-\` -\&\(-\+\:-\@\/])', r' \1 '),  # tokenize punctuation. apostrophe is missing
+    (r'([^0-9])([\.,])', r'\1 \2 '),  # tokenize period and comma unless preceded by a digit
+    (r'([\.,])([^0-9])', r' \1 \2'),  # tokenize period and comma unless followed by a digit
+    (r'([0-9])(-)', r'\1 \2 ')  # tokenize dash when preceded by a digit
+]
+normalize2 = [(re.compile(pattern), replace) for (pattern, replace) in normalize2]
+
+
+def normalize(s):
+    '''Normalize and tokenize text. This is lifted from NIST mteval-v11a.pl.'''
+    # Added to bypass NIST-style pre-processing of hyp and ref files -- wade
+    if (nonorm):
+        return s.split()
+    if type(s) is not str:
+        s = " ".join(s)
+    # language-independent part:
+    for (pattern, replace) in normalize1:
+        s = re.sub(pattern, replace, s)
+    s = xml.sax.saxutils.unescape(s, {'&quot;': '"'})
+    # language-dependent part (assuming Western languages):
+    s = " %s " % s
+    if not preserve_case:
+        s = s.lower()  # this might not be identical to the original
+    for (pattern, replace) in normalize2:
+        s = re.sub(pattern, replace, s)
+    return s.split()
+
+
+def count_ngrams(words, n=4):
+    counts = {}
+    for k in range(1, n + 1):
+        for i in range(len(words) - k + 1):
+            ngram = tuple(words[i:i + k])
+            counts[ngram] = counts.get(ngram, 0) + 1
+    return counts
+
+
+def cook_refs(refs, n=4):
+    '''Takes a list of reference sentences for a single segment
+    and returns an object that encapsulates everything that BLEU
+    needs to know about them.'''
+
+    refs = [normalize(ref) for ref in refs]
+    maxcounts = {}
+    for ref in refs:
+        counts = count_ngrams(ref, n)
+        for (ngram, count) in counts.items():
+            maxcounts[ngram] = max(maxcounts.get(ngram, 0), count)
+    return ([len(ref) for ref in refs], maxcounts)
+
+
+def cook_test(test, item, n=4):
+    '''Takes a test sentence and returns an object that
+    encapsulates everything that BLEU needs to know about it.'''
+    (reflens, refmaxcounts) = item
+    test = normalize(test)
+    result = {}
+    result["testlen"] = len(test)
+
+    # Calculate effective reference sentence length.
+
+    if eff_ref_len == "shortest":
+        result["reflen"] = min(reflens)
+    elif eff_ref_len == "average":
+        result["reflen"] = float(sum(reflens)) / len(reflens)
+    elif eff_ref_len == "closest":
+        min_diff = None
+        for reflen in reflens:
+            if min_diff is None or abs(reflen - len(test)) < min_diff:
+                min_diff = abs(reflen - len(test))
+                result['reflen'] = reflen
+
+    result["guess"] = [max(len(test) - k + 1, 0) for k in range(1, n + 1)]
+
+    result['correct'] = [0] * n
+    counts = count_ngrams(test, n)
+    for (ngram, count) in counts.items():
+        result["correct"][len(ngram) - 1] += min(refmaxcounts.get(ngram, 0), count)
+
+    return result
+
+
+def score_cooked(allcomps, n=4, ground=0, smooth=1):
+    totalcomps = {'testlen': 0, 'reflen': 0, 'guess': [0] * n, 'correct': [0] * n}
+    for comps in allcomps:
+        for key in ['testlen', 'reflen']:
+            totalcomps[key] += comps[key]
+        for key in ['guess', 'correct']:
+            for k in range(n):
+                totalcomps[key][k] += comps[key][k]
+    logbleu = 0.0
+    all_bleus = []
+    for k in range(n):
+        correct = totalcomps['correct'][k]
+        guess = totalcomps['guess'][k]
+        addsmooth = 0
+        if smooth == 1 and k > 0:
+            addsmooth = 1
+        logbleu += math.log(correct + addsmooth + sys.float_info.min) - math.log(guess + addsmooth + sys.float_info.min)
+        if guess == 0:
+            all_bleus.append(-10000000)
+        else:
+            all_bleus.append(math.log(correct + sys.float_info.min) - math.log(guess))
+
+    logbleu /= float(n)
+    all_bleus.insert(0, logbleu)
+
+    brevPenalty = min(0, 1 - float(totalcomps['reflen'] + 1) / (totalcomps['testlen'] + 1))
+    for i in range(len(all_bleus)):
+        if i == 0:
+            all_bleus[i] += brevPenalty
+        all_bleus[i] = math.exp(all_bleus[i])
+    return all_bleus
+
+
+def bleu(refs, candidate, ground=0, smooth=1):
+    refs = cook_refs(refs)
+    test = cook_test(candidate, refs)
+    return score_cooked([test], ground=ground, smooth=smooth)
+
+
+def splitPuncts(line):
+    return ' '.join(re.findall(r"[\w]+|[^\s\w]", line))
+
+
+def computeMaps(predictions, goldfile):
+    predictionMap = {}
+    goldMap = {}
+    gf = open(goldfile, 'r')
+
+    for row in predictions:
+        cols = row.strip().split('\t')
+        if len(cols) == 1:
+            (rid, pred) = (cols[0], '')
+        else:
+            (rid, pred) = (cols[0], cols[1])
+        predictionMap[rid] = [splitPuncts(pred.strip().lower())]
+
+    for row in gf:
+        (rid, pred) = row.split('\t')
+        if rid in predictionMap:  # Only insert if the id exists for the method
+            if rid not in goldMap:
+                goldMap[rid] = []
+            goldMap[rid].append(splitPuncts(pred.strip().lower()))
+
+    sys.stderr.write('Total: ' + str(len(goldMap)) + '\n')
+    return (goldMap, predictionMap)
+
+
+# m1 is the reference map
+# m2 is the prediction map
+def bleuFromMaps(m1, m2):
+    score = [0] * 5
+    num = 0.0
+
+    for key in m1:
+        if key in m2:
+            bl = bleu(m1[key], m2[key][0])
+            score = [score[i] + bl[i] for i in range(0, len(bl))]
+            num += 1
+    return [s * 100.0 / num for s in score]
+
+
+if __name__ == '__main__':
+    reference_file = sys.argv[1]
+    predictions = []
+    for row in sys.stdin:
+        predictions.append(row)
+    (goldMap, predictionMap) = computeMaps(predictions, reference_file)
+    print(bleuFromMaps(goldMap, predictionMap)[0])

models.py

@@ -0,0 +1,398 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import (RobertaConfig, RobertaModel, RobertaTokenizer,
+                          BartConfig, BartForConditionalGeneration, BartTokenizer,
+                          T5Config, T5ForConditionalGeneration, T5Tokenizer)
+import logging
+
+logger = logging.getLogger(__name__)
+
+MODEL_CLASSES = {'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer),
+                 't5': (T5Config, T5ForConditionalGeneration, T5Tokenizer),
+                 'codet5': (T5Config, T5ForConditionalGeneration, RobertaTokenizer),
+                 'bart': (BartConfig, BartForConditionalGeneration, BartTokenizer)}
+
+
+def get_model_size(model):
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    model_size = sum([np.prod(p.size()) for p in model_parameters])
+    return "{}M".format(round(model_size / 1e+6))
+
+
+def build_or_load_gen_model(args):
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name)
+    if args.model_type == 'roberta':
+        encoder = model_class.from_pretrained(args.model_name_or_path, config=config)
+        decoder_layer = nn.TransformerDecoderLayer(d_model=config.hidden_size, nhead=config.num_attention_heads)
+        decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+        model = Seq2Seq(encoder=encoder, decoder=decoder, config=config,
+                        beam_size=args.beam_size, max_length=args.max_target_length,
+                        sos_id=tokenizer.cls_token_id, eos_id=tokenizer.sep_token_id)
+    else:
+        model = model_class.from_pretrained(args.model_name_or_path)
+
+    logger.info("Finish loading model [%s] from %s", get_model_size(model), args.model_name_or_path)
+
+    if args.load_model_path is not None:
+        logger.info("Reload model from {}".format(args.load_model_path))
+        model.load_state_dict(torch.load(args.load_model_path))
+
+    return config, model, tokenizer
+
+
+class RobertaClassificationHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size * 2, config.hidden_size)
+        self.out_proj = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, x, **kwargs):
+        x = x.reshape(-1, x.size(-1) * 2)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.out_proj(x)
+        return x
+
+
+class CloneModel(nn.Module):
+    def __init__(self, encoder, config, tokenizer, args):
+        super(CloneModel, self).__init__()
+        self.encoder = encoder
+        self.config = config
+        self.tokenizer = tokenizer
+        self.classifier = RobertaClassificationHead(config)
+        self.args = args
+
+    def get_t5_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        outputs = self.encoder(input_ids=source_ids, attention_mask=attention_mask,
+                               labels=source_ids, decoder_attention_mask=attention_mask, output_hidden_states=True)
+        hidden_states = outputs['decoder_hidden_states'][-1]
+        eos_mask = source_ids.eq(self.config.eos_token_id)
+
+        if len(torch.unique(eos_mask.sum(1))) > 1:
+            raise ValueError("All examples must have the same number of <eos> tokens.")
+        vec = hidden_states[eos_mask, :].view(hidden_states.size(0), -1,
+                                              hidden_states.size(-1))[:, -1, :]
+        return vec
+
+    def get_bart_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        outputs = self.encoder(input_ids=source_ids, attention_mask=attention_mask,
+                               labels=source_ids, decoder_attention_mask=attention_mask, output_hidden_states=True)
+        hidden_states = outputs['decoder_hidden_states'][-1]
+        eos_mask = source_ids.eq(self.config.eos_token_id)
+
+        if len(torch.unique(eos_mask.sum(1))) > 1:
+            raise ValueError("All examples must have the same number of <eos> tokens.")
+        vec = hidden_states[eos_mask, :].view(hidden_states.size(0), -1,
+                                              hidden_states.size(-1))[:, -1, :]
+        return vec
+
+    def get_roberta_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        vec = self.encoder(input_ids=source_ids, attention_mask=attention_mask)[0][:, 0, :]
+        return vec
+
+    def forward(self, source_ids=None, labels=None):
+        source_ids = source_ids.view(-1, self.args.max_source_length)
+
+        if self.args.model_type == 'codet5':
+            vec = self.get_t5_vec(source_ids)
+        elif self.args.model_type == 'bart':
+            vec = self.get_bart_vec(source_ids)
+        elif self.args.model_type == 'roberta':
+            vec = self.get_roberta_vec(source_ids)
+
+        logits = self.classifier(vec)
+        prob = nn.functional.softmax(logits)
+
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+            return loss, prob
+        else:
+            return prob
+
+
+class DefectModel(nn.Module):
+    def __init__(self, encoder, config, tokenizer, args):
+        super(DefectModel, self).__init__()
+        self.encoder = encoder
+        self.config = config
+        self.tokenizer = tokenizer
+        self.classifier = nn.Linear(config.hidden_size, 2)
+        self.args = args
+
+    def get_t5_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        outputs = self.encoder(input_ids=source_ids, attention_mask=attention_mask,
+                               labels=source_ids, decoder_attention_mask=attention_mask, output_hidden_states=True)
+        hidden_states = outputs['decoder_hidden_states'][-1]
+        eos_mask = source_ids.eq(self.config.eos_token_id)
+
+        if len(torch.unique(eos_mask.sum(1))) > 1:
+            raise ValueError("All examples must have the same number of <eos> tokens.")
+        vec = hidden_states[eos_mask, :].view(hidden_states.size(0), -1,
+                                              hidden_states.size(-1))[:, -1, :]
+        return vec
+
+    def get_bart_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        outputs = self.encoder(input_ids=source_ids, attention_mask=attention_mask,
+                               labels=source_ids, decoder_attention_mask=attention_mask, output_hidden_states=True)
+        hidden_states = outputs['decoder_hidden_states'][-1]
+        eos_mask = source_ids.eq(self.config.eos_token_id)
+
+        if len(torch.unique(eos_mask.sum(1))) > 1:
+            raise ValueError("All examples must have the same number of <eos> tokens.")
+        vec = hidden_states[eos_mask, :].view(hidden_states.size(0), -1,
+                                              hidden_states.size(-1))[:, -1, :]
+        return vec
+
+    def get_roberta_vec(self, source_ids):
+        attention_mask = source_ids.ne(self.tokenizer.pad_token_id)
+        vec = self.encoder(input_ids=source_ids, attention_mask=attention_mask)[0][:, 0, :]
+        return vec
+
+    def forward(self, source_ids=None, labels=None):
+        source_ids = source_ids.view(-1, self.args.max_source_length)
+
+        if self.args.model_type == 'codet5':
+            vec = self.get_t5_vec(source_ids)
+        elif self.args.model_type == 'bart':
+            vec = self.get_bart_vec(source_ids)
+        elif self.args.model_type == 'roberta':
+            vec = self.get_roberta_vec(source_ids)
+
+        logits = self.classifier(vec)
+        prob = nn.functional.softmax(logits)
+
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+            return loss, prob
+        else:
+            return prob
+
+
+# https://github.com/microsoft/CodeBERT/blob/master/CodeBERT/code2nl/model.py
+class Seq2Seq(nn.Module):
+    """
+        Build Seqence-to-Sequence.
+
+        Parameters:
+
+        * `encoder`- encoder of seq2seq model. e.g. roberta
+        * `decoder`- decoder of seq2seq model. e.g. transformer
+        * `config`- configuration of encoder model.
+        * `beam_size`- beam size for beam search.
+        * `max_length`- max length of target for beam search.
+        * `sos_id`- start of symbol ids in target for beam search.
+        * `eos_id`- end of symbol ids in target for beam search.
+    """
+
+    def __init__(self, encoder, decoder, config, beam_size=None, max_length=None, sos_id=None, eos_id=None):
+        super(Seq2Seq, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.config = config
+        self.register_buffer("bias", torch.tril(torch.ones(2048, 2048)))
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lsm = nn.LogSoftmax(dim=-1)
+        self.tie_weights()
+
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id
+
+    def _tie_or_clone_weights(self, first_module, second_module):
+        """ Tie or clone module weights depending of weither we are using TorchScript or not
+        """
+        if self.config.torchscript:
+            first_module.weight = nn.Parameter(second_module.weight.clone())
+        else:
+            first_module.weight = second_module.weight
+
+    def tie_weights(self):
+        """ Make sure we are sharing the input and output embeddings.
+            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
+        """
+        self._tie_or_clone_weights(self.lm_head,
+                                   self.encoder.embeddings.word_embeddings)
+
+    def forward(self, source_ids=None, source_mask=None, target_ids=None, target_mask=None, args=None):
+        outputs = self.encoder(source_ids, attention_mask=source_mask)
+        encoder_output = outputs[0].permute([1, 0, 2]).contiguous()
+        if target_ids is not None:
+            attn_mask = -1e4 * (1 - self.bias[:target_ids.shape[1], :target_ids.shape[1]])
+            tgt_embeddings = self.encoder.embeddings(target_ids).permute([1, 0, 2]).contiguous()
+            out = self.decoder(tgt_embeddings, encoder_output, tgt_mask=attn_mask,
+                               memory_key_padding_mask=~source_mask)
+            # memory_key_padding_mask=(1 - source_mask).bool())
+            hidden_states = torch.tanh(self.dense(out)).permute([1, 0, 2]).contiguous()
+            lm_logits = self.lm_head(hidden_states)
+            # Shift so that tokens < n predict n
+            active_loss = target_mask[..., 1:].ne(0).view(-1) == 1
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = target_ids[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1))[active_loss],
+                            shift_labels.view(-1)[active_loss])
+
+            outputs = loss, loss * active_loss.sum(), active_loss.sum()
+            return outputs
+        else:
+            # Predict
+            preds = []
+            zero = torch.cuda.LongTensor(1).fill_(0)
+            for i in range(source_ids.shape[0]):
+                context = encoder_output[:, i:i + 1]
+                context_mask = source_mask[i:i + 1, :]
+                beam = Beam(self.beam_size, self.sos_id, self.eos_id)
+                input_ids = beam.getCurrentState()
+                context = context.repeat(1, self.beam_size, 1)
+                context_mask = context_mask.repeat(self.beam_size, 1)
+                for _ in range(self.max_length):
+                    if beam.done():
+                        break
+                    attn_mask = -1e4 * (1 - self.bias[:input_ids.shape[1], :input_ids.shape[1]])
+                    tgt_embeddings = self.encoder.embeddings(input_ids).permute([1, 0, 2]).contiguous()
+                    out = self.decoder(tgt_embeddings, context, tgt_mask=attn_mask,
+                                       memory_key_padding_mask=~context_mask)
+                    # memory_key_padding_mask=(1 - context_mask).bool())
+                    out = torch.tanh(self.dense(out))
+                    hidden_states = out.permute([1, 0, 2]).contiguous()[:, -1, :]
+                    out = self.lsm(self.lm_head(hidden_states)).data
+                    beam.advance(out)
+                    input_ids.data.copy_(input_ids.data.index_select(0, beam.getCurrentOrigin()))
+                    input_ids = torch.cat((input_ids, beam.getCurrentState()), -1)
+                hyp = beam.getHyp(beam.getFinal())
+                pred = beam.buildTargetTokens(hyp)[:self.beam_size]
+                pred = [torch.cat([x.view(-1) for x in p] + [zero] * (self.max_length - len(p))).view(1, -1) for p in
+                        pred]
+                preds.append(torch.cat(pred, 0).unsqueeze(0))
+
+            preds = torch.cat(preds, 0)
+            return preds
+
+
+class Beam(object):
+    def __init__(self, size, sos, eos):
+        self.size = size
+        self.tt = torch.cuda
+        # The score for each translation on the beam.
+        self.scores = self.tt.FloatTensor(size).zero_()
+        # The backpointers at each time-step.
+        self.prevKs = []
+        # The outputs at each time-step.
+        self.nextYs = [self.tt.LongTensor(size)
+                           .fill_(0)]
+        self.nextYs[0][0] = sos
+        # Has EOS topped the beam yet.
+        self._eos = eos
+        self.eosTop = False
+        # Time and k pair for finished.
+        self.finished = []
+
+    def getCurrentState(self):
+        "Get the outputs for the current timestep."
+        batch = self.tt.LongTensor(self.nextYs[-1]).view(-1, 1)
+        return batch
+
+    def getCurrentOrigin(self):
+        "Get the backpointers for the current timestep."
+        return self.prevKs[-1]
+
+    def advance(self, wordLk):
+        """
+        Given prob over words for every last beam `wordLk` and attention
+        `attnOut`: Compute and update the beam search.
+
+        Parameters:
+
+        * `wordLk`- probs of advancing from the last step (K x words)
+        * `attnOut`- attention at the last step
+
+        Returns: True if beam search is complete.
+        """
+        numWords = wordLk.size(1)
+
+        # Sum the previous scores.
+        if len(self.prevKs) > 0:
+            beamLk = wordLk + self.scores.unsqueeze(1).expand_as(wordLk)
+
+            # Don't let EOS have children.
+            for i in range(self.nextYs[-1].size(0)):
+                if self.nextYs[-1][i] == self._eos:
+                    beamLk[i] = -1e20
+        else:
+            beamLk = wordLk[0]
+        flatBeamLk = beamLk.view(-1)
+        bestScores, bestScoresId = flatBeamLk.topk(self.size, 0, True, True)
+
+        self.scores = bestScores
+
+        # bestScoresId is flattened beam x word array, so calculate which
+        # word and beam each score came from
+        prevK = bestScoresId // numWords
+        self.prevKs.append(prevK)
+        self.nextYs.append((bestScoresId - prevK * numWords))
+
+        for i in range(self.nextYs[-1].size(0)):
+            if self.nextYs[-1][i] == self._eos:
+                s = self.scores[i]
+                self.finished.append((s, len(self.nextYs) - 1, i))
+
+        # End condition is when top-of-beam is EOS and no global score.
+        if self.nextYs[-1][0] == self._eos:
+            self.eosTop = True
+
+    def done(self):
+        return self.eosTop and len(self.finished) >= self.size
+
+    def getFinal(self):
+        if len(self.finished) == 0:
+            self.finished.append((self.scores[0], len(self.nextYs) - 1, 0))
+        self.finished.sort(key=lambda a: -a[0])
+        if len(self.finished) != self.size:
+            unfinished = []
+            for i in range(self.nextYs[-1].size(0)):
+                if self.nextYs[-1][i] != self._eos:
+                    s = self.scores[i]
+                    unfinished.append((s, len(self.nextYs) - 1, i))
+            unfinished.sort(key=lambda a: -a[0])
+            self.finished += unfinished[:self.size - len(self.finished)]
+        return self.finished[:self.size]
+
+    def getHyp(self, beam_res):
+        """
+        Walk back to construct the full hypothesis.
+        """
+        hyps = []
+        for _, timestep, k in beam_res:
+            hyp = []
+            for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):
+                hyp.append(self.nextYs[j + 1][k])
+                k = self.prevKs[j][k]
+            hyps.append(hyp[::-1])
+        return hyps
+
+    def buildTargetTokens(self, preds):
+        sentence = []
+        for pred in preds:
+            tokens = []
+            for tok in pred:
+                if tok == self._eos:
+                    break
+                tokens.append(tok)
+            sentence.append(tokens)
+        return sentence

run_clone.py

@@ -0,0 +1,325 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+from __future__ import absolute_import
+import os
+import pdb
+
+from models import CloneModel
+import logging
+import argparse
+import math
+import numpy as np
+from io import open
+from tqdm import tqdm
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
+from torch.utils.data.distributed import DistributedSampler
+from transformers import (AdamW, get_linear_schedule_with_warmup,
+                          RobertaConfig, RobertaModel, RobertaTokenizer,
+                          BartConfig, BartForConditionalGeneration, BartTokenizer,
+                          T5Config, T5ForConditionalGeneration, T5Tokenizer)
+import multiprocessing
+from sklearn.metrics import recall_score, precision_score, f1_score
+import time
+
+from configs import add_args, set_seed
+from utils import get_filenames, get_elapse_time, load_and_cache_clone_data
+from models import get_model_size
+
+MODEL_CLASSES = {'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer),
+                 't5': (T5Config, T5ForConditionalGeneration, T5Tokenizer),
+                 'codet5': (T5Config, T5ForConditionalGeneration, RobertaTokenizer),
+                 'bart': (BartConfig, BartForConditionalGeneration, BartTokenizer)}
+
+cpu_cont = multiprocessing.cpu_count()
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def evaluate(args, model, eval_examples, eval_data, write_to_pred=False):
+    eval_sampler = SequentialSampler(eval_data)
+    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+    # Eval!
+    logger.info("***** Running evaluation  *****")
+    logger.info("  Num examples = %d", len(eval_examples))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+    eval_loss = 0.0
+    nb_eval_steps = 0
+    model.eval()
+    logits = []
+    y_trues = []
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Evaluating"):
+        inputs = batch[0].to(args.device)
+        labels = batch[1].to(args.device)
+        with torch.no_grad():
+            lm_loss, logit = model(inputs, labels)
+            eval_loss += lm_loss.mean().item()
+            logits.append(logit.cpu().numpy())
+            y_trues.append(labels.cpu().numpy())
+        nb_eval_steps += 1
+    logits = np.concatenate(logits, 0)
+    y_trues = np.concatenate(y_trues, 0)
+    best_threshold = 0.5
+
+    y_preds = logits[:, 1] > best_threshold
+    recall = recall_score(y_trues, y_preds)
+    precision = precision_score(y_trues, y_preds)
+    f1 = f1_score(y_trues, y_preds)
+    result = {
+        "eval_recall": float(recall),
+        "eval_precision": float(precision),
+        "eval_f1": float(f1),
+        "eval_threshold": best_threshold,
+    }
+
+    logger.info("***** Eval results *****")
+    for key in sorted(result.keys()):
+        logger.info("  %s = %s", key, str(round(result[key], 4)))
+    logger.info("  " + "*" * 20)
+
+    if write_to_pred:
+        with open(os.path.join(args.output_dir, "predictions.txt"), 'w') as f:
+            for example, pred in zip(eval_examples, y_preds):
+                if pred:
+                    f.write(example.url1 + '\t' + example.url2 + '\t' + '1' + '\n')
+                else:
+                    f.write(example.url1 + '\t' + example.url2 + '\t' + '0' + '\n')
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    t0 = time.time()
+    args = add_args(parser)
+    logger.info(args)
+
+    # Setup CUDA, GPU & distributed training
+    if args.local_rank == -1 or args.no_cuda:
+        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
+        args.n_gpu = torch.cuda.device_count()
+    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        torch.cuda.set_device(args.local_rank)
+        device = torch.device("cuda", args.local_rank)
+        torch.distributed.init_process_group(backend='nccl')
+        args.n_gpu = 1
+
+    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, cpu count: %d",
+                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), cpu_cont)
+    args.device = device
+    set_seed(args)
+
+    # Build model
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+    model = model_class.from_pretrained(args.model_name_or_path)
+    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name)
+    model.resize_token_embeddings(32000)
+
+    model = CloneModel(model, config, tokenizer, args)
+    logger.info("Finish loading model [%s] from %s", get_model_size(model), args.model_name_or_path)
+
+    if args.load_model_path is not None:
+        logger.info("Reload model from {}".format(args.load_model_path))
+        model.load_state_dict(torch.load(args.load_model_path))
+
+    model.to(device)
+
+    pool = multiprocessing.Pool(cpu_cont)
+    args.train_filename, args.dev_filename, args.test_filename = get_filenames(args.data_dir, args.task, args.sub_task)
+    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+')
+
+    if args.do_train:
+        if args.n_gpu > 1:
+            # multi-gpu training
+            model = torch.nn.DataParallel(model)
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            summary_fn = '{}/{}'.format(args.summary_dir, '/'.join(args.output_dir.split('/')[1:]))
+            tb_writer = SummaryWriter(summary_fn)
+
+        # Prepare training data loader
+        train_examples, train_data = load_and_cache_clone_data(args, args.train_filename, pool, tokenizer, 'train',
+                                                               is_sample=False)
+        if args.local_rank == -1:
+            train_sampler = RandomSampler(train_data)
+        else:
+            train_sampler = DistributedSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)
+
+        num_train_optimization_steps = args.num_train_epochs * len(train_dataloader)
+        save_steps = max(len(train_dataloader) // 5, 1)
+
+        # Prepare optimizer and schedule (linear warmup and decay)
+        no_decay = ['bias', 'LayerNorm.weight']
+        optimizer_grouped_parameters = [
+            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+             'weight_decay': args.weight_decay},
+            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+        ]
+        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
+
+        if args.warmup_steps < 1:
+            warmup_steps = num_train_optimization_steps * args.warmup_steps
+        else:
+            warmup_steps = int(args.warmup_steps)
+        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps,
+                                                    num_training_steps=num_train_optimization_steps)
+
+        # Start training
+        train_example_num = len(train_data)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", train_example_num)
+        logger.info("  Batch size = %d", args.train_batch_size)
+        logger.info("  Batch num = %d", math.ceil(train_example_num / args.train_batch_size))
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+
+        global_step, best_f1 = 0, 0
+        not_f1_inc_cnt = 0
+        is_early_stop = False
+        for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
+            bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
+            nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
+            model.train()
+            for step, batch in enumerate(bar):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids, labels = batch
+                # pdb.set_trace()
+
+                loss, logits = model(source_ids, labels)
+
+                if args.n_gpu > 1:
+                    loss = loss.mean()  # mean() to average on multi-gpu.
+                if args.gradient_accumulation_steps > 1:
+                    loss = loss / args.gradient_accumulation_steps
+                tr_loss += loss.item()
+
+                nb_tr_examples += source_ids.size(0)
+                nb_tr_steps += 1
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                if nb_tr_steps % args.gradient_accumulation_steps == 0:
+                    # Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    global_step += 1
+                    train_loss = round(tr_loss * args.gradient_accumulation_steps / nb_tr_steps, 4)
+                    bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
+
+                if (step + 1) % save_steps == 0 and args.do_eval:
+                    logger.info("***** CUDA.empty_cache() *****")
+                    torch.cuda.empty_cache()
+
+                    eval_examples, eval_data = load_and_cache_clone_data(args, args.dev_filename, pool, tokenizer,
+                                                                         'valid', is_sample=True)
+
+                    result = evaluate(args, model, eval_examples, eval_data)
+                    eval_f1 = result['eval_f1']
+
+                    if args.data_num == -1:
+                        tb_writer.add_scalar('dev_f1', round(eval_f1, 4), cur_epoch)
+
+                    # save last checkpoint
+                    last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                    if not os.path.exists(last_output_dir):
+                        os.makedirs(last_output_dir)
+
+                    if True or args.data_num == -1 and args.save_last_checkpoints:
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the last model into %s", output_model_file)
+
+                    if eval_f1 > best_f1:
+                        not_f1_inc_cnt = 0
+                        logger.info("  Best f1: %s", round(eval_f1, 4))
+                        logger.info("  " + "*" * 20)
+                        fa.write("[%d] Best f1 changed into %.4f\n" % (cur_epoch, round(eval_f1, 4)))
+                        best_f1 = eval_f1
+                        # Save best checkpoint for best ppl
+                        output_dir = os.path.join(args.output_dir, 'checkpoint-best-f1')
+                        if not os.path.exists(output_dir):
+                            os.makedirs(output_dir)
+                        if args.data_num == -1 or True:
+                            model_to_save = model.module if hasattr(model, 'module') else model
+                            output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                            torch.save(model_to_save.state_dict(), output_model_file)
+                            logger.info("Save the best ppl model into %s", output_model_file)
+                    else:
+                        not_f1_inc_cnt += 1
+                        logger.info("F1 does not increase for %d epochs", not_f1_inc_cnt)
+                        if not_f1_inc_cnt > args.patience:
+                            logger.info("Early stop as f1 do not increase for %d times", not_f1_inc_cnt)
+                            fa.write("[%d] Early stop as not_f1_inc_cnt=%d\n" % (cur_epoch, not_f1_inc_cnt))
+                            is_early_stop = True
+                            break
+
+                model.train()
+            if is_early_stop:
+                break
+
+            logger.info("***** CUDA.empty_cache() *****")
+            torch.cuda.empty_cache()
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            tb_writer.close()
+
+    if args.do_test:
+        logger.info("  " + "***** Testing *****")
+        logger.info("  Batch size = %d", args.eval_batch_size)
+
+        for criteria in ['best-f1']:
+            file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
+            logger.info("Reload model from {}".format(file))
+            model.load_state_dict(torch.load(file))
+
+            if args.n_gpu > 1:
+                # multi-gpu training
+                model = torch.nn.DataParallel(model)
+
+            eval_examples, eval_data = load_and_cache_clone_data(args, args.test_filename, pool, tokenizer, 'test',
+                                                                 False)
+
+            result = evaluate(args, model, eval_examples, eval_data, write_to_pred=True)
+            logger.info("  test_f1=%.4f", result['eval_f1'])
+            logger.info("  test_prec=%.4f", result['eval_precision'])
+            logger.info("  test_rec=%.4f", result['eval_recall'])
+            logger.info("  " + "*" * 20)
+
+            fa.write("[%s] test-f1: %.4f, precision: %.4f, recall: %.4f\n" % (
+                criteria, result['eval_f1'], result['eval_precision'], result['eval_recall']))
+            if args.res_fn:
+                with open(args.res_fn, 'a+') as f:
+                    f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
+                    f.write("[%s] f1: %.4f, precision: %.4f, recall: %.4f\n\n" % (
+                        criteria, result['eval_f1'], result['eval_precision'], result['eval_recall']))
+    fa.close()
+
+
+if __name__ == "__main__":
+    main()

run_defect.py

@@ -0,0 +1,314 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+from __future__ import absolute_import
+import os
+import logging
+import argparse
+import math
+import numpy as np
+from io import open
+from tqdm import tqdm
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler, TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+                          RobertaConfig, RobertaModel, RobertaTokenizer,
+                          BartConfig, BartForConditionalGeneration, BartTokenizer,
+                          T5Config, T5ForConditionalGeneration, T5Tokenizer)
+import multiprocessing
+import time
+
+from models import DefectModel
+from configs import add_args, set_seed
+from utils import get_filenames, get_elapse_time, load_and_cache_defect_data
+from models import get_model_size
+
+MODEL_CLASSES = {'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer),
+                 't5': (T5Config, T5ForConditionalGeneration, T5Tokenizer),
+                 'codet5': (T5Config, T5ForConditionalGeneration, RobertaTokenizer),
+                 'bart': (BartConfig, BartForConditionalGeneration, BartTokenizer)}
+
+cpu_cont = multiprocessing.cpu_count()
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def evaluate(args, model, eval_examples, eval_data, write_to_pred=False):
+    eval_sampler = SequentialSampler(eval_data)
+    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+    # Eval!
+    logger.info("***** Running evaluation *****")
+    logger.info("  Num examples = %d", len(eval_examples))
+    logger.info("  Num batches = %d", len(eval_dataloader))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+    eval_loss = 0.0
+    nb_eval_steps = 0
+    model.eval()
+    logits = []
+    labels = []
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Evaluating"):
+        inputs = batch[0].to(args.device)
+        label = batch[1].to(args.device)
+        with torch.no_grad():
+            lm_loss, logit = model(inputs, label)
+            eval_loss += lm_loss.mean().item()
+            logits.append(logit.cpu().numpy())
+            labels.append(label.cpu().numpy())
+        nb_eval_steps += 1
+    logits = np.concatenate(logits, 0)
+    labels = np.concatenate(labels, 0)
+    preds = logits[:, 1] > 0.5
+    eval_acc = np.mean(labels == preds)
+    eval_loss = eval_loss / nb_eval_steps
+    perplexity = torch.tensor(eval_loss)
+
+    result = {
+        "eval_loss": float(perplexity),
+        "eval_acc": round(eval_acc, 4),
+    }
+
+    logger.info("***** Eval results *****")
+    for key in sorted(result.keys()):
+        logger.info("  %s = %s", key, str(round(result[key], 4)))
+
+    if write_to_pred:
+        with open(os.path.join(args.output_dir, "predictions.txt"), 'w') as f:
+            for example, pred in zip(eval_examples, preds):
+                if pred:
+                    f.write(str(example.idx) + '\t1\n')
+                else:
+                    f.write(str(example.idx) + '\t0\n')
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    t0 = time.time()
+    args = add_args(parser)
+    logger.info(args)
+
+    # Setup CUDA, GPU & distributed training
+    if args.local_rank == -1 or args.no_cuda:
+        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
+        args.n_gpu = torch.cuda.device_count()
+    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        torch.cuda.set_device(args.local_rank)
+        device = torch.device("cuda", args.local_rank)
+        torch.distributed.init_process_group(backend='nccl')
+        args.n_gpu = 1
+
+    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, cpu count: %d",
+                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), cpu_cont)
+    args.device = device
+    set_seed(args)
+
+    # Build model
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+    model = model_class.from_pretrained(args.model_name_or_path)
+    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name)
+
+    model = DefectModel(model, config, tokenizer, args)
+    logger.info("Finish loading model [%s] from %s", get_model_size(model), args.model_name_or_path)
+
+    if args.load_model_path is not None:
+        logger.info("Reload model from {}".format(args.load_model_path))
+        model.load_state_dict(torch.load(args.load_model_path))
+
+    model.to(device)
+
+    pool = multiprocessing.Pool(cpu_cont)
+    args.train_filename, args.dev_filename, args.test_filename = get_filenames(args.data_dir, args.task, args.sub_task)
+    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+')
+
+    if args.do_train:
+        if args.n_gpu > 1:
+            # multi-gpu training
+            model = torch.nn.DataParallel(model)
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            summary_fn = '{}/{}'.format(args.summary_dir, '/'.join(args.output_dir.split('/')[1:]))
+            tb_writer = SummaryWriter(summary_fn)
+
+        # Prepare training data loader
+        train_examples, train_data = load_and_cache_defect_data(args, args.train_filename, pool, tokenizer, 'train',
+                                                                is_sample=False)
+        if args.local_rank == -1:
+            train_sampler = RandomSampler(train_data)
+        else:
+            train_sampler = DistributedSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)
+
+        num_train_optimization_steps = args.num_train_epochs * len(train_dataloader)
+        save_steps = max(len(train_dataloader), 1)
+
+        # Prepare optimizer and schedule (linear warmup and decay)
+        no_decay = ['bias', 'LayerNorm.weight']
+        optimizer_grouped_parameters = [
+            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+             'weight_decay': args.weight_decay},
+            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+        ]
+        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
+
+        if args.warmup_steps < 1:
+            warmup_steps = num_train_optimization_steps * args.warmup_steps
+        else:
+            warmup_steps = int(args.warmup_steps)
+        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps,
+                                                    num_training_steps=num_train_optimization_steps)
+
+        # Start training
+        train_example_num = len(train_data)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", train_example_num)
+        logger.info("  Batch size = %d", args.train_batch_size)
+        logger.info("  Batch num = %d", math.ceil(train_example_num / args.train_batch_size))
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+
+        global_step, best_acc = 0, 0
+        not_acc_inc_cnt = 0
+        is_early_stop = False
+        for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
+            bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
+            nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
+            model.train()
+            for step, batch in enumerate(bar):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids, labels = batch
+
+                loss, logits = model(source_ids, labels)
+
+                if args.n_gpu > 1:
+                    loss = loss.mean()  # mean() to average on multi-gpu.
+                if args.gradient_accumulation_steps > 1:
+                    loss = loss / args.gradient_accumulation_steps
+                tr_loss += loss.item()
+
+                nb_tr_examples += source_ids.size(0)
+                nb_tr_steps += 1
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                if nb_tr_steps % args.gradient_accumulation_steps == 0:
+                    # Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    global_step += 1
+                    train_loss = round(tr_loss * args.gradient_accumulation_steps / nb_tr_steps, 4)
+                    bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
+
+                if (step + 1) % save_steps == 0 and args.do_eval:
+                    logger.info("***** CUDA.empty_cache() *****")
+                    torch.cuda.empty_cache()
+
+                    eval_examples, eval_data = load_and_cache_defect_data(args, args.dev_filename, pool, tokenizer,
+                                                                          'valid', is_sample=False)
+
+                    result = evaluate(args, model, eval_examples, eval_data)
+                    eval_acc = result['eval_acc']
+
+                    if args.data_num == -1:
+                        tb_writer.add_scalar('dev_acc', round(eval_acc, 4), cur_epoch)
+
+                    # save last checkpoint
+                    last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                    if not os.path.exists(last_output_dir):
+                        os.makedirs(last_output_dir)
+
+                    if True or args.data_num == -1 and args.save_last_checkpoints:
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the last model into %s", output_model_file)
+
+                    if eval_acc > best_acc:
+                        not_acc_inc_cnt = 0
+                        logger.info("  Best acc: %s", round(eval_acc, 4))
+                        logger.info("  " + "*" * 20)
+                        fa.write("[%d] Best acc changed into %.4f\n" % (cur_epoch, round(eval_acc, 4)))
+                        best_acc = eval_acc
+                        # Save best checkpoint for best ppl
+                        output_dir = os.path.join(args.output_dir, 'checkpoint-best-acc')
+                        if not os.path.exists(output_dir):
+                            os.makedirs(output_dir)
+                        if args.data_num == -1 or True:
+                            model_to_save = model.module if hasattr(model, 'module') else model
+                            output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                            torch.save(model_to_save.state_dict(), output_model_file)
+                            logger.info("Save the best ppl model into %s", output_model_file)
+                    else:
+                        not_acc_inc_cnt += 1
+                        logger.info("acc does not increase for %d epochs", not_acc_inc_cnt)
+                        if not_acc_inc_cnt > args.patience:
+                            logger.info("Early stop as acc do not increase for %d times", not_acc_inc_cnt)
+                            fa.write("[%d] Early stop as not_acc_inc_cnt=%d\n" % (cur_epoch, not_acc_inc_cnt))
+                            is_early_stop = True
+                            break
+
+                model.train()
+            if is_early_stop:
+                break
+
+            logger.info("***** CUDA.empty_cache() *****")
+            torch.cuda.empty_cache()
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            tb_writer.close()
+
+    if args.do_test:
+        logger.info("  " + "***** Testing *****")
+        logger.info("  Batch size = %d", args.eval_batch_size)
+
+        for criteria in ['best-acc']:
+            file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
+            logger.info("Reload model from {}".format(file))
+            model.load_state_dict(torch.load(file))
+
+            if args.n_gpu > 1:
+                # multi-gpu training
+                model = torch.nn.DataParallel(model)
+
+            eval_examples, eval_data = load_and_cache_defect_data(args, args.test_filename, pool, tokenizer, 'test',
+                                                                  False)
+
+            result = evaluate(args, model, eval_examples, eval_data, write_to_pred=True)
+            logger.info("  test_acc=%.4f", result['eval_acc'])
+            logger.info("  " + "*" * 20)
+
+            fa.write("[%s] test-acc: %.4f\n" % (criteria, result['eval_acc']))
+            if args.res_fn:
+                with open(args.res_fn, 'a+') as f:
+                    f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
+                    f.write("[%s] acc: %.4f\n\n" % (
+                        criteria, result['eval_acc']))
+    fa.close()
+
+
+if __name__ == "__main__":
+    main()

run_gen.py

@@ -0,0 +1,387 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+import os
+import logging
+import argparse
+import math
+import numpy as np
+from tqdm import tqdm
+import multiprocessing
+import time
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
+from torch.utils.data.distributed import DistributedSampler
+from transformers import AdamW, get_linear_schedule_with_warmup
+from models import build_or_load_gen_model
+from evaluator import smooth_bleu
+from evaluator.CodeBLEU import calc_code_bleu
+from evaluator.bleu import _bleu
+from utils import get_filenames, get_elapse_time, load_and_cache_gen_data
+from configs import add_args, set_seed, set_dist
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def eval_ppl_epoch(args, eval_data, eval_examples, model, tokenizer):
+    eval_sampler = SequentialSampler(eval_data)
+    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size,
+                                 num_workers=4, pin_memory=True)
+    # Start evaluating model
+    logger.info("  " + "***** Running ppl evaluation *****")
+    logger.info("  Num examples = %d", len(eval_examples))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+
+    model.eval()
+    eval_loss, batch_num = 0, 0
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval ppl"):
+        batch = tuple(t.to(args.device) for t in batch)
+        source_ids, target_ids = batch
+        source_mask = source_ids.ne(tokenizer.pad_token_id)
+        target_mask = target_ids.ne(tokenizer.pad_token_id)
+
+        with torch.no_grad():
+            if args.model_type == 'roberta':
+                loss, _, _ = model(source_ids=source_ids, source_mask=source_mask,
+                                   target_ids=target_ids, target_mask=target_mask)
+            else:
+                outputs = model(input_ids=source_ids, attention_mask=source_mask,
+                                labels=target_ids, decoder_attention_mask=target_mask)
+                loss = outputs.loss
+
+        eval_loss += loss.item()
+        batch_num += 1
+    eval_loss = eval_loss / batch_num
+    eval_ppl = round(np.exp(eval_loss), 5)
+    return eval_ppl
+
+
+def eval_bleu_epoch(args, eval_data, eval_examples, model, tokenizer, split_tag, criteria):
+    logger.info("  ***** Running bleu evaluation on {} data*****".format(split_tag))
+    logger.info("  Num examples = %d", len(eval_examples))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+    eval_sampler = SequentialSampler(eval_data)
+    if args.data_num == -1:
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size,
+                                     num_workers=4, pin_memory=True)
+    else:
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+    model.eval()
+    pred_ids = []
+    bleu, codebleu = 0.0, 0.0
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval bleu for {} set".format(split_tag)):
+        source_ids = batch[0].to(args.device)
+        source_mask = source_ids.ne(tokenizer.pad_token_id)
+        with torch.no_grad():
+            if args.model_type == 'roberta':
+                preds = model(source_ids=source_ids, source_mask=source_mask)
+
+                top_preds = [pred[0].cpu().numpy() for pred in preds]
+            else:
+                preds = model.generate(source_ids,
+                                       attention_mask=source_mask,
+                                       use_cache=True,
+                                       num_beams=args.beam_size,
+                                       early_stopping=args.task == 'summarize',
+                                       max_length=args.max_target_length)
+                top_preds = list(preds.cpu().numpy())
+            pred_ids.extend(top_preds)
+
+    pred_nls = [tokenizer.decode(id, skip_special_tokens=True, clean_up_tokenization_spaces=False) for id in pred_ids]
+
+    output_fn = os.path.join(args.res_dir, "test_{}.output".format(criteria))
+    gold_fn = os.path.join(args.res_dir, "test_{}.gold".format(criteria))
+    src_fn = os.path.join(args.res_dir, "test_{}.src".format(criteria))
+
+    if args.task in ['defect']:
+        target_dict = {0: 'false', 1: 'true'}
+        golds = [target_dict[ex.target] for ex in eval_examples]
+        eval_acc = np.mean([int(p == g) for p, g in zip(pred_nls, golds)])
+        result = {'em': eval_acc * 100, 'bleu': 0, 'codebleu': 0}
+
+        with open(output_fn, 'w') as f, open(gold_fn, 'w') as f1, open(src_fn, 'w') as f2:
+            for pred_nl, gold in zip(pred_nls, eval_examples):
+                f.write(pred_nl.strip() + '\n')
+                f1.write(target_dict[gold.target] + '\n')
+                f2.write(gold.source.strip() + '\n')
+            logger.info("Save the predictions into %s", output_fn)
+    else:
+        dev_accs, predictions = [], []
+        with open(output_fn, 'w') as f, open(gold_fn, 'w') as f1, open(src_fn, 'w') as f2:
+            for pred_nl, gold in zip(pred_nls, eval_examples):
+                dev_accs.append(pred_nl.strip() == gold.target.strip())
+                if args.task in ['summarize']:
+                    # for smooth-bleu4 evaluation
+                    predictions.append(str(gold.idx) + '\t' + pred_nl)
+                    f.write(str(gold.idx) + '\t' + pred_nl.strip() + '\n')
+                    f1.write(str(gold.idx) + '\t' + gold.target.strip() + '\n')
+                    f2.write(str(gold.idx) + '\t' + gold.source.strip() + '\n')
+                else:
+                    f.write(pred_nl.strip() + '\n')
+                    f1.write(gold.target.strip() + '\n')
+                    f2.write(gold.source.strip() + '\n')
+
+        if args.task == 'summarize':
+            (goldMap, predictionMap) = smooth_bleu.computeMaps(predictions, gold_fn)
+            bleu = round(smooth_bleu.bleuFromMaps(goldMap, predictionMap)[0], 2)
+        else:
+            bleu = round(_bleu(gold_fn, output_fn), 2)
+            if args.task in ['concode', 'translate', 'refine']:
+                codebleu = calc_code_bleu.get_codebleu(gold_fn, output_fn, args.lang)
+
+        result = {'em': np.mean(dev_accs) * 100, 'bleu': bleu}
+        if args.task == 'concode':
+            result['codebleu'] = codebleu * 100
+
+    logger.info("***** Eval results *****")
+    for key in sorted(result.keys()):
+        logger.info("  %s = %s", key, str(round(result[key], 4)))
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    args = add_args(parser)
+    logger.info(args)
+    t0 = time.time()
+
+    set_dist(args)
+    set_seed(args)
+    config, model, tokenizer = build_or_load_gen_model(args)
+    model.to(args.device)
+    if args.n_gpu > 1:
+        # for DataParallel
+        model = torch.nn.DataParallel(model)
+    pool = multiprocessing.Pool(args.cpu_cont)
+    args.train_filename, args.dev_filename, args.test_filename = get_filenames(args.data_dir, args.task, args.sub_task)
+    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+')
+
+    if args.do_train:
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            summary_fn = '{}/{}'.format(args.summary_dir, '/'.join(args.output_dir.split('/')[1:]))
+            tb_writer = SummaryWriter(summary_fn)
+
+        # Prepare training data loader
+        train_examples, train_data = load_and_cache_gen_data(args, args.train_filename, pool, tokenizer, 'train')
+        train_sampler = RandomSampler(train_data) if args.local_rank == -1 else DistributedSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size,
+                                      num_workers=4, pin_memory=True)
+
+        # Prepare optimizer and schedule (linear warmup and decay)
+        no_decay = ['bias', 'LayerNorm.weight']
+        optimizer_grouped_parameters = [
+            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+             'weight_decay': args.weight_decay},
+            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+        ]
+        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
+        num_train_optimization_steps = args.num_train_epochs * len(train_dataloader)
+        scheduler = get_linear_schedule_with_warmup(optimizer,
+                                                    num_warmup_steps=args.warmup_steps,
+                                                    num_training_steps=num_train_optimization_steps)
+
+        # Start training
+        train_example_num = len(train_data)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", train_example_num)
+        logger.info("  Batch size = %d", args.train_batch_size)
+        logger.info("  Batch num = %d", math.ceil(train_example_num / args.train_batch_size))
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+
+        dev_dataset = {}
+        global_step, best_bleu_em, best_ppl = 0, -1, 1e6
+        not_loss_dec_cnt, not_bleu_em_inc_cnt = 0, 0 if args.do_eval_bleu else 1e6
+
+        for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
+            bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
+            nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
+            model.train()
+            for step, batch in enumerate(bar):
+                batch = tuple(t.to(args.device) for t in batch)
+                source_ids, target_ids = batch
+                source_mask = source_ids.ne(tokenizer.pad_token_id)
+                target_mask = target_ids.ne(tokenizer.pad_token_id)
+
+                if args.model_type == 'roberta':
+                    loss, _, _ = model(source_ids=source_ids, source_mask=source_mask,
+                                       target_ids=target_ids, target_mask=target_mask)
+                else:
+                    outputs = model(input_ids=source_ids, attention_mask=source_mask,
+                                    labels=target_ids, decoder_attention_mask=target_mask)
+                    loss = outputs.loss
+
+                if args.n_gpu > 1:
+                    loss = loss.mean()  # mean() to average on multi-gpu.
+                if args.gradient_accumulation_steps > 1:
+                    loss = loss / args.gradient_accumulation_steps
+                tr_loss += loss.item()
+
+                nb_tr_examples += source_ids.size(0)
+                nb_tr_steps += 1
+                loss.backward()
+
+                if nb_tr_steps % args.gradient_accumulation_steps == 0:
+                    # Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    global_step += 1
+                    train_loss = round(tr_loss * args.gradient_accumulation_steps / (nb_tr_steps + 1), 4)
+                    bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
+
+            if args.do_eval:
+                # Eval model with dev dataset
+                if 'dev_loss' in dev_dataset:
+                    eval_examples, eval_data = dev_dataset['dev_loss']
+                else:
+                    eval_examples, eval_data = load_and_cache_gen_data(args, args.dev_filename, pool, tokenizer, 'dev')
+                    dev_dataset['dev_loss'] = eval_examples, eval_data
+
+                eval_ppl = eval_ppl_epoch(args, eval_data, eval_examples, model, tokenizer)
+                result = {'epoch': cur_epoch, 'global_step': global_step, 'eval_ppl': eval_ppl}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  " + "*" * 20)
+                if args.data_num == -1:
+                    tb_writer.add_scalar('dev_ppl', eval_ppl, cur_epoch)
+
+                # save last checkpoint
+                if args.save_last_checkpoints:
+                    last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                    if not os.path.exists(last_output_dir):
+                        os.makedirs(last_output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model
+                    output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    logger.info("Save the last model into %s", output_model_file)
+
+                if eval_ppl < best_ppl:
+                    not_loss_dec_cnt = 0
+                    logger.info("  Best ppl:%s", eval_ppl)
+                    logger.info("  " + "*" * 20)
+                    fa.write("[%d] Best ppl changed into %.4f\n" % (cur_epoch, eval_ppl))
+                    best_ppl = eval_ppl
+
+                    # Save best checkpoint for best ppl
+                    output_dir = os.path.join(args.output_dir, 'checkpoint-best-ppl')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    if args.always_save_model:
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the best ppl model into %s", output_model_file)
+                else:
+                    not_loss_dec_cnt += 1
+                    logger.info("Ppl does not decrease for %d epochs", not_loss_dec_cnt)
+                    if all([x > args.patience for x in [not_bleu_em_inc_cnt, not_loss_dec_cnt]]):
+                        early_stop_str = "[%d] Early stop as not_bleu_em_inc_cnt=%d, and not_loss_dec_cnt=%d\n" % (
+                            cur_epoch, not_bleu_em_inc_cnt, not_loss_dec_cnt)
+                        logger.info(early_stop_str)
+                        fa.write(early_stop_str)
+                        break
+                logger.info("***** CUDA.empty_cache() *****")
+                torch.cuda.empty_cache()
+                if args.do_eval_bleu:
+                    eval_examples, eval_data = load_and_cache_gen_data(args, args.dev_filename, pool, tokenizer, 'dev',
+                                                                       only_src=True, is_sample=True)
+
+                    result = eval_bleu_epoch(args, eval_data, eval_examples, model, tokenizer, 'dev', 'e%d' % cur_epoch)
+                    dev_bleu, dev_em = result['bleu'], result['em']
+                    if args.task in ['summarize']:
+                        dev_bleu_em = dev_bleu
+                    elif args.task in ['defect']:
+                        dev_bleu_em = dev_em
+                    else:
+                        dev_bleu_em = dev_bleu + dev_em
+                    if args.data_num == -1:
+                        tb_writer.add_scalar('dev_bleu_em', dev_bleu_em, cur_epoch)
+                        # tb_writer.add_scalar('dev_em', dev_em, cur_epoch)
+                    if dev_bleu_em > best_bleu_em:
+                        not_bleu_em_inc_cnt = 0
+                        logger.info("  [%d] Best bleu+em: %.2f (bleu: %.2f, em: %.2f)",
+                                    cur_epoch, dev_bleu_em, dev_bleu, dev_em)
+                        logger.info("  " + "*" * 20)
+                        best_bleu_em = dev_bleu_em
+                        fa.write("[%d] Best bleu+em changed into %.2f (bleu: %.2f, em: %.2f)\n" % (
+                            cur_epoch, best_bleu_em, dev_bleu, dev_em))
+                        # Save best checkpoint for best bleu
+                        output_dir = os.path.join(args.output_dir, 'checkpoint-best-bleu')
+                        if not os.path.exists(output_dir):
+                            os.makedirs(output_dir)
+                        if args.data_num == -1 or args.always_save_model:
+                            model_to_save = model.module if hasattr(model, 'module') else model
+                            output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                            torch.save(model_to_save.state_dict(), output_model_file)
+                            logger.info("Save the best bleu model into %s", output_model_file)
+                    else:
+                        not_bleu_em_inc_cnt += 1
+                        logger.info("Bleu does not increase for %d epochs", not_bleu_em_inc_cnt)
+                        fa.write(
+                            "[%d] Best bleu+em (%.2f) does not drop changed for %d epochs, cur bleu+em: %.2f (bleu: %.2f, em: %.2f)\n" % (
+                                cur_epoch, best_bleu_em, not_bleu_em_inc_cnt, dev_bleu_em, dev_bleu, dev_em))
+                        if all([x > args.patience for x in [not_bleu_em_inc_cnt, not_loss_dec_cnt]]):
+                            stop_early_str = "[%d] Early stop as not_bleu_em_inc_cnt=%d, and not_loss_dec_cnt=%d\n" % (
+                                cur_epoch, not_bleu_em_inc_cnt, not_loss_dec_cnt)
+                            logger.info(stop_early_str)
+                            fa.write(stop_early_str)
+                            break
+            logger.info("***** CUDA.empty_cache() *****")
+            torch.cuda.empty_cache()
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            tb_writer.close()
+        logger.info("Finish training and take %s", get_elapse_time(t0))
+
+    if args.do_test:
+        logger.info("  " + "***** Testing *****")
+        logger.info("  Batch size = %d", args.eval_batch_size)
+
+        for criteria in ['best-bleu']:
+            file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
+            logger.info("Reload model from {}".format(file))
+            model.load_state_dict(torch.load(file))
+            eval_examples, eval_data = load_and_cache_gen_data(args, args.test_filename, pool, tokenizer, 'test',
+                                                               only_src=True, is_sample=False)
+            result = eval_bleu_epoch(args, eval_data, eval_examples, model, tokenizer, 'test', criteria)
+            test_bleu, test_em = result['bleu'], result['em']
+            test_codebleu = result['codebleu'] if 'codebleu' in result else 0
+            result_str = "[%s] bleu-4: %.2f, em: %.4f, codebleu: %.4f\n" % (criteria, test_bleu, test_em, test_codebleu)
+            logger.info(result_str)
+            fa.write(result_str)
+            if args.res_fn:
+                with open(args.res_fn, 'a+') as f:
+                    f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
+                    f.write(result_str)
+    logger.info("Finish and take {}".format(get_elapse_time(t0)))
+    fa.write("Finish and take {}".format(get_elapse_time(t0)))
+    fa.close()
+
+
+if __name__ == "__main__":
+    main()

run_multi_gen.py

@@ -0,0 +1,535 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+import os
+import torch
+import logging
+import argparse
+import math
+import numpy as np
+from tqdm import tqdm
+from itertools import cycle
+import multiprocessing
+import time
+import sys
+import pdb
+
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
+from torch.utils.data.distributed import DistributedSampler
+from transformers import AdamW, get_linear_schedule_with_warmup
+from models import build_or_load_gen_model
+from evaluator import smooth_bleu
+from evaluator.CodeBLEU import calc_code_bleu
+from evaluator.bleu import _bleu
+from utils import get_elapse_time, load_and_cache_multi_gen_data
+from configs import add_args, set_seed, set_dist
+
+cpu_cont = multiprocessing.cpu_count()
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+WORKER_NUM = 0
+
+
+def get_max_trg_len_by_task(task, sub_task):
+    if task == 'summarize':
+        max_target_length = 128
+    elif task == 'translate':
+        max_target_length = 256
+    elif task == 'refine':
+        if sub_task == 'small':
+            max_target_length = 120
+        else:
+            max_target_length = 240
+    elif task == 'concode':
+        max_target_length = 150
+    elif task == 'defect':
+        max_target_length = 3
+    return max_target_length
+
+
+def get_bs(cur_task, model_tag):
+    task = cur_task.split('_')[0]
+    sub_task = cur_task.split('_')[-1]
+    if 'codet5_small' in model_tag:
+        bs = 32
+        if task == 'summarize' or task == 'translate' or (task == 'refine' and sub_task == 'small'):
+            bs = 64
+    else:
+        # codet5_base
+        bs = 28
+        if task == 'translate':
+            bs = 25
+        elif task == 'summarize':
+            bs = 40
+    return bs
+
+
+def eval_bleu(args, eval_data, eval_examples, model, tokenizer, split_tag, cur_task, criteria):
+    eval_sampler = SequentialSampler(eval_data)
+    if args.data_num == -1:
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size,
+                                     num_workers=4, pin_memory=True)
+    else:
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+    task = cur_task.split('_')[0]
+    sub_task = cur_task.split('_')[-1]
+    max_target_length = get_max_trg_len_by_task(task, sub_task)
+
+    model.eval()
+    pred_ids = []
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval bleu for {} set".format(split_tag)):
+        source_ids = batch[0].to(args.device)
+        source_mask = source_ids.ne(tokenizer.pad_token_id)
+        with torch.no_grad():
+            if args.model_type == 'roberta':
+                preds = model(source_ids=source_ids, source_mask=source_mask)
+
+                top_preds = [pred[0].cpu().numpy() for pred in preds]
+            else:
+                preds = model.generate(source_ids,
+                                       attention_mask=source_mask,
+                                       use_cache=True,
+                                       num_beams=5,
+                                       max_length=max_target_length,  # length_penalty=0.6,
+                                       early_stopping=task == 'summarize')
+                top_preds = list(preds.cpu().numpy())
+            pred_ids.extend(top_preds)
+
+    pred_nls = [tokenizer.decode(id, skip_special_tokens=True, clean_up_tokenization_spaces=False) for id in pred_ids]
+    if task == 'defect':
+        target_dict = {0: 'false', 1: 'true'}
+        golds = [target_dict[ex.target] for ex in eval_examples]
+        eval_acc = np.mean([int(p == g) for p, g in zip(pred_nls, golds)])
+        result = {'em': eval_acc, 'bleu': 0, 'codebleu': 0}
+
+    else:
+        dev_accs = []
+        predictions = []
+        res_dir = os.path.join(args.res_dir, cur_task)
+        if not os.path.exists(res_dir):
+            os.makedirs(res_dir)
+        output_fn = os.path.join(res_dir, "test_{}.output".format(criteria))
+        gold_fn = os.path.join(res_dir, "test_{}.gold".format(criteria))
+        with open(output_fn, 'w') as f, open(gold_fn, 'w') as f1:
+            for pred_nl, gold in zip(pred_nls, eval_examples):
+                dev_accs.append(pred_nl.strip() == gold.target.strip())
+                if task == 'summarize':
+                    predictions.append(str(gold.idx) + '\t' + pred_nl)
+                    f.write(str(gold.idx) + '\t' + pred_nl.strip() + '\n')
+                    f1.write(str(gold.idx) + '\t' + gold.target.strip() + '\n')
+                else:
+                    f.write(pred_nl.strip() + '\n')
+                    f1.write(gold.target.strip() + '\n')
+
+        try:
+            if task == 'summarize':
+                (goldMap, predictionMap) = smooth_bleu.computeMaps(predictions, gold_fn)
+                bleu = round(smooth_bleu.bleuFromMaps(goldMap, predictionMap)[0], 2)
+            else:
+
+                bleu = round(_bleu(gold_fn, output_fn), 2)
+                if split_tag == 'test':
+                    if task in ['summarize', 'search']:
+                        cur_lang = sub_task
+                    elif task in ['refine', 'concode', 'clone']:
+                        cur_lang = 'java'
+                    elif task == 'defect':
+                        cur_lang = 'c'
+                    elif task == 'translate':
+                        cur_lang = 'c_sharp' if sub_task == 'java-cs' else 'java'
+                    codebleu = calc_code_bleu.get_codebleu(gold_fn, output_fn, cur_lang)
+        except:
+            bleu = 0.0
+            codebleu = 0.0
+
+        result = {}
+        em = np.mean(dev_accs) * 100
+        result['em'] = em
+        result['bleu'] = bleu
+        if not args.task == 'summarize' and split_tag == 'test':
+            result['codebleu'] = codebleu * 100
+
+    logger.info("***** Eval results [%s] *****", cur_task)
+    for key in sorted(result.keys()):
+        logger.info("  %s = %s", key, str(round(result[key], 4)))
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    args = add_args(parser)
+    logger.info(args)
+    t0 = time.time()
+
+    set_dist(args)
+    set_seed(args)
+    config, model, tokenizer = build_or_load_gen_model(args)
+    model.to(args.device)
+    if args.n_gpu > 1:
+        # for DataParallel
+        model = torch.nn.DataParallel(model)
+    pool = multiprocessing.Pool(args.cpu_cont)
+    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+')
+
+    fa_dict = {}
+    if args.do_train:
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            summary_fn = './tensorboard/{}'.format('/'.join(args.output_dir.split('/')[1:]))
+            tb_writer = SummaryWriter(summary_fn)
+
+        # Prepare training data loader
+        train_examples_data_dict = load_and_cache_multi_gen_data(args, pool, tokenizer, 'train', is_sample=False)
+        train_data_list = [v[1] for k, v in train_examples_data_dict.items()]
+        all_tasks = [k for k, v in train_examples_data_dict.items()]
+        total_train_data_num = sum([len(v[0]) for k, v in train_examples_data_dict.items()])
+
+        for cur_task in all_tasks:
+            summary_dir = os.path.join(args.output_dir, 'summary')
+            if not os.path.exists(summary_dir):
+                os.makedirs(summary_dir)
+            fa_dict[cur_task] = open(os.path.join(summary_dir, '{}_summary.log'.format(cur_task)), 'a+')
+
+        train_dataloader_dict = dict()
+        for train_data, cur_task in zip(train_data_list, all_tasks):
+            if args.local_rank == -1:
+                train_sampler = RandomSampler(train_data)
+            else:
+                train_sampler = DistributedSampler(train_data)
+            if args.data_num == -1:
+                train_dataloader = DataLoader(train_data, sampler=train_sampler,
+                                              batch_size=get_bs(cur_task, args.model_name_or_path),
+                                              num_workers=WORKER_NUM, pin_memory=True)
+            else:
+                train_dataloader = DataLoader(train_data, sampler=train_sampler,
+                                              batch_size=get_bs(cur_task, args.model_name_or_path))
+
+            train_dataloader_dict[cur_task] = cycle(train_dataloader)
+
+        # Prepare optimizer and schedule (linear warmup and decay)
+        no_decay = ['bias', 'LayerNorm.weight']
+        optimizer_grouped_parameters = [
+            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+             'weight_decay': args.weight_decay},
+            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+        ]
+        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
+
+        scheduler = get_linear_schedule_with_warmup(optimizer,
+                                                    num_warmup_steps=args.warmup_steps,
+                                                    num_training_steps=args.max_steps)
+
+        # Start training
+        logger.info("***** Running training *****")
+        logger.info("  Total train data num = %d", total_train_data_num)
+        logger.info("  Max step = %d, Save step = %d", args.max_steps, args.save_steps)
+
+        dev_dataset = {}
+        step, global_step = 0, 0
+        best_bleu_em = dict([(k, -1) for k in all_tasks])
+        best_loss = dict([(k, 1e6) for k in all_tasks])
+        not_bleu_em_inc_cnt = dict([(k, 0) for k in all_tasks])
+        is_early_stop = dict([(k, 0) for k in all_tasks])
+
+        patience_pairs = []
+        for cur_task in all_tasks:
+            task = cur_task.split('_')[0]
+            if task == 'summarize':
+                patience_pairs.append((cur_task, 2))
+            elif task == 'translate':
+                patience_pairs.append((cur_task, 5))
+            elif task == 'refine':
+                patience_pairs.append((cur_task, 5))
+            elif task == 'concode':
+                patience_pairs.append((cur_task, 3))
+            elif task == 'defect':
+                patience_pairs.append((cur_task, 2))
+        patience_dict = dict(patience_pairs)
+        logger.info('Patience: %s', patience_dict)
+
+        probs = [len(x) for x in train_data_list]
+        probs = [x / sum(probs) for x in probs]
+        probs = [x ** 0.7 for x in probs]
+        probs = [x / sum(probs) for x in probs]
+
+        nb_tr_examples, nb_tr_steps, tr_nb, tr_loss, logging_loss = 0, 0, 0, 0, 0
+
+        bar = tqdm(total=args.max_steps, desc="Training")
+        skip_cnt = 0
+        while True:
+            cur_task = np.random.choice(all_tasks, 1, p=probs)[0]
+            train_dataloader = train_dataloader_dict[cur_task]
+            if is_early_stop[cur_task]:
+                skip_cnt += 1
+                if skip_cnt > 50:
+                    logger.info('All tasks have early stopped at %d', step)
+                    break
+                continue
+            else:
+                skip_cnt = 0
+
+            step += 1
+            batch = next(train_dataloader)
+
+            model.train()
+            batch = tuple(t.to(args.device) for t in batch)
+            source_ids, target_ids = batch
+            # logger.info('cur_task: %s, bs: %d', cur_task, source_ids.shape[0])
+            source_mask = source_ids.ne(tokenizer.pad_token_id)
+            target_mask = target_ids.ne(tokenizer.pad_token_id)
+            # pdb.set_trace()
+
+            if args.model_type == 'roberta':
+                loss, _, _ = model(source_ids=source_ids, source_mask=source_mask,
+                                   target_ids=target_ids, target_mask=target_mask)
+            else:
+                outputs = model(input_ids=source_ids, attention_mask=source_mask,
+                                labels=target_ids, decoder_attention_mask=target_mask)
+                loss = outputs.loss
+
+            if args.n_gpu > 1:
+                loss = loss.mean()  # mean() to average on multi-gpu.
+            if args.gradient_accumulation_steps > 1:
+                loss = loss / args.gradient_accumulation_steps
+            tr_loss += loss.item()
+
+            nb_tr_examples += source_ids.size(0)
+            nb_tr_steps += 1
+            loss.backward()
+
+            if nb_tr_steps % args.gradient_accumulation_steps == 0:
+                # Update parameters
+                optimizer.step()
+                optimizer.zero_grad()
+                scheduler.step()
+                global_step += 1
+                train_loss = round((tr_loss - logging_loss) / (global_step - tr_nb), 6)
+                bar.update(1)
+                bar.set_description("[{}] Train loss {}".format(step, round(train_loss, 3)))
+
+                if args.local_rank in [-1, 0] and args.log_steps > 0 and global_step % args.log_steps == 0:
+                    logging_loss = train_loss
+                    tr_nb = global_step
+
+                if args.do_eval and args.local_rank in [-1, 0] \
+                        and args.save_steps > 0 and global_step % args.save_steps == 0:
+                    # save last checkpoint
+                    if args.data_num == -1 and args.save_last_checkpoints:
+                        last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                        if not os.path.exists(last_output_dir):
+                            os.makedirs(last_output_dir)
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the last model into %s", output_model_file)
+                    if global_step % 100000 == 0:
+                        step_tag = '{}00k'.format(global_step // 100000)
+                        last_output_dir = os.path.join(args.output_dir, 'checkpoint-step-{}'.format(step_tag))
+                        if not os.path.exists(last_output_dir):
+                            os.makedirs(last_output_dir)
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the last model into %s", output_model_file)
+                    # Eval model with dev dataset
+                    if 'dev_loss' in dev_dataset:
+                        eval_examples_data_dict = dev_dataset['dev_loss']
+                    else:
+                        eval_examples_data_dict = load_and_cache_multi_gen_data(args, pool, tokenizer, 'dev')
+                        dev_dataset['dev_loss'] = eval_examples_data_dict
+
+                    for cur_task in eval_examples_data_dict.keys():
+                        if is_early_stop[cur_task]:
+                            continue
+                        eval_examples, eval_data = eval_examples_data_dict[cur_task]
+                        eval_sampler = SequentialSampler(eval_data)
+                        if args.data_num == -1:
+                            eval_dataloader = DataLoader(eval_data, sampler=eval_sampler,
+                                                         batch_size=args.eval_batch_size,
+                                                         num_workers=4, pin_memory=True)
+                        else:
+                            eval_dataloader = DataLoader(eval_data, sampler=eval_sampler,
+                                                         batch_size=args.eval_batch_size)
+
+                        logger.info("  " + "***** Running ppl evaluation on [{}] *****".format(cur_task))
+                        logger.info("  Num examples = %d", len(eval_examples))
+                        logger.info("  Batch size = %d", args.eval_batch_size)
+
+                        # Start Evaluating model
+                        model.eval()
+                        eval_loss, batch_num = 0, 0
+                        for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval ppl"):
+                            batch = tuple(t.to(args.device) for t in batch)
+                            source_ids, target_ids = batch
+                            source_mask = source_ids.ne(tokenizer.pad_token_id)
+                            target_mask = target_ids.ne(tokenizer.pad_token_id)
+
+                            with torch.no_grad():
+                                if args.model_type == 'roberta':
+                                    loss, _, _ = model(source_ids=source_ids, source_mask=source_mask,
+                                                       target_ids=target_ids, target_mask=target_mask)
+                                else:
+                                    outputs = model(input_ids=source_ids, attention_mask=source_mask,
+                                                    labels=target_ids, decoder_attention_mask=target_mask)
+                                    loss = outputs.loss
+
+                            eval_loss += loss.item()
+                            batch_num += 1
+                        # Pring loss of dev dataset
+                        eval_loss = eval_loss / batch_num
+                        result = {'cur_task': cur_task,
+                                  'global_step': global_step,
+                                  'eval_ppl': round(np.exp(eval_loss), 5),
+                                  'train_loss': round(train_loss, 5)}
+                        for key in sorted(result.keys()):
+                            logger.info("  %s = %s", key, str(result[key]))
+                        logger.info("  " + "*" * 20)
+
+                        if args.data_num == -1:
+                            tb_writer.add_scalar('dev_ppl_{}'.format(cur_task),
+                                                 round(np.exp(eval_loss), 5),
+                                                 global_step)
+
+                        if eval_loss < best_loss[cur_task]:
+                            logger.info("  Best ppl:%s", round(np.exp(eval_loss), 5))
+                            logger.info("  " + "*" * 20)
+                            fa_dict[cur_task].write(
+                                "[%d: %s] Best ppl changed into %.4f\n" % (global_step, cur_task, np.exp(eval_loss)))
+                            best_loss[cur_task] = eval_loss
+
+                            # Save best checkpoint for best ppl
+                            output_dir = os.path.join(args.output_dir, 'checkpoint-best-ppl', cur_task)
+                            if not os.path.exists(output_dir):
+                                os.makedirs(output_dir)
+                            if args.data_num == -1 or args.always_save_model:
+                                model_to_save = model.module if hasattr(model, 'module') else model
+                                output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                                torch.save(model_to_save.state_dict(), output_model_file)
+                                logger.info("Save the best ppl model into %s", output_model_file)
+
+                    if args.do_eval_bleu:
+                        eval_examples_data_dict = load_and_cache_multi_gen_data(args, pool, tokenizer, 'dev',
+                                                                                only_src=True, is_sample=True)
+                        for cur_task in eval_examples_data_dict.keys():
+                            if is_early_stop[cur_task]:
+                                continue
+                            eval_examples, eval_data = eval_examples_data_dict[cur_task]
+
+                            # pdb.set_trace()
+                            result = eval_bleu(args, eval_data, eval_examples, model, tokenizer, 'dev', cur_task,
+                                               criteria='e{}'.format(global_step))
+                            dev_bleu, dev_em = result['bleu'], result['em']
+                            if args.task == 'summarize':
+                                dev_bleu_em = dev_bleu
+                            elif args.task in ['defect', 'clone']:
+                                dev_bleu_em = dev_em
+                            else:
+                                dev_bleu_em = dev_bleu + dev_em
+                            if args.data_num == -1:
+                                tb_writer.add_scalar('dev_bleu_em_{}'.format(cur_task), dev_bleu_em, global_step)
+
+                            if dev_bleu_em > best_bleu_em[cur_task]:
+                                not_bleu_em_inc_cnt[cur_task] = 0
+                                logger.info("  [%d: %s] Best bleu+em: %.2f (bleu: %.2f, em: %.2f)",
+                                            global_step, cur_task, dev_bleu_em, dev_bleu, dev_em)
+                                logger.info("  " + "*" * 20)
+                                best_bleu_em[cur_task] = dev_bleu_em
+                                fa_dict[cur_task].write(
+                                    "[%d: %s] Best bleu+em changed into %.2f (bleu: %.2f, em: %.2f)\n" % (
+                                        global_step, cur_task, best_bleu_em[cur_task], dev_bleu, dev_em))
+                                # Save best checkpoint for best bleu
+                                output_dir = os.path.join(args.output_dir, 'checkpoint-best-bleu', cur_task)
+                                if not os.path.exists(output_dir):
+                                    os.makedirs(output_dir)
+                                if args.data_num == -1 or args.always_save_model:
+                                    model_to_save = model.module if hasattr(model, 'module') else model
+                                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                                    torch.save(model_to_save.state_dict(), output_model_file)
+                                    logger.info("Save the best bleu model into %s", output_model_file)
+                            else:
+                                not_bleu_em_inc_cnt[cur_task] += 1
+                                logger.info("[%d %s] bleu/em does not increase for %d eval steps",
+                                            global_step, cur_task, not_bleu_em_inc_cnt[cur_task])
+                                if not_bleu_em_inc_cnt[cur_task] > patience_dict[cur_task]:
+                                    logger.info("[%d %s] Early stop as bleu/em does not increase for %d eval steps",
+                                                global_step, cur_task, not_bleu_em_inc_cnt[cur_task])
+                                    is_early_stop[cur_task] = 1
+                                    fa_dict[cur_task].write(
+                                        "[%d %s] Early stop as bleu/em does not increase for %d eval steps, takes %s" %
+                                        (global_step, cur_task, not_bleu_em_inc_cnt[cur_task], get_elapse_time(t0)))
+
+                    logger.info("***** CUDA.empty_cache() *****")
+                    torch.cuda.empty_cache()
+                if global_step >= args.max_steps:
+                    logger.info("Reach the max step: %d", args.max_steps)
+                    break
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            tb_writer.close()
+        logger.info("Finish training and take %.2f", time.time() - t0)
+        for cur_task in all_tasks:
+            fa_dict[cur_task].close()
+
+    if args.do_test:
+        logger.info("  " + "***** Testing *****")
+        logger.info("  Batch size = %d", args.eval_batch_size)
+        eval_examples_data_dict = load_and_cache_multi_gen_data(args, pool, tokenizer, 'test', only_src=True)
+        all_tasks = list(eval_examples_data_dict.keys())
+        for cur_task in all_tasks:
+            summary_dir = os.path.join(args.output_dir, 'summary')
+            if not os.path.exists(summary_dir):
+                os.makedirs(summary_dir)
+            fa_dict[cur_task] = open(os.path.join(summary_dir, '{}_summary.log'.format(cur_task)), 'a+')
+
+        for cur_task in all_tasks:
+            eval_examples, eval_data = eval_examples_data_dict[cur_task]
+            args.task = cur_task.split('_')[0]
+            args.sub_task = cur_task.split('_')[-1]
+
+            for criteria in ['best-bleu', 'best-ppl', 'last']:
+                file = os.path.join(args.output_dir, 'checkpoint-{}/{}/pytorch_model.bin'.format(criteria, cur_task))
+                model.load_state_dict(torch.load(file))
+
+                result = eval_bleu(args, eval_data, eval_examples, model, tokenizer, 'test', cur_task, criteria)
+                test_bleu, test_em = result['bleu'], result['em']
+                test_codebleu = result['codebleu'] if 'codebleu' in result else 0
+                result_str = "[%s %s] bleu-4: %.2f, em: %.4f, codebleu: %.4f\n" % (
+                    cur_task, criteria, test_bleu, test_em, test_codebleu)
+                logger.info(result_str)
+                fa_dict[cur_task].write(result_str)
+                fa.write(result_str)
+                if args.res_fn:
+                    with open(args.res_fn, 'a+') as f:
+                        f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
+                        f.write(result_str)
+    logger.info("Finish and take {}".format(get_elapse_time(t0)))
+    for cur_task in all_tasks:
+        fa_dict[cur_task].close()
+    fa.write("Finish and take {}".format(get_elapse_time(t0)))
+    fa.close()
+
+
+if __name__ == "__main__":
+    main()

sh/exp_with_args.sh

@@ -0,0 +1,94 @@
+WORKDIR="your_CodeT5_path/CodeT5"
+export PYTHONPATH=$WORKDIR
+
+TASK=${1}
+SUB_TASK=${2}
+MODEL_TAG=${3}
+GPU=${4}
+DATA_NUM=${5}
+BS=${6}
+LR=${7}
+SRC_LEN=${8}
+TRG_LEN=${9}
+PATIENCE=${10}
+EPOCH=${11}
+WARMUP=${12}
+MODEL_DIR=${13}
+SUMMARY_DIR=${14}
+RES_FN=${15}
+
+if [[ $DATA_NUM == -1 ]]; then
+  DATA_TAG='all'
+else
+  DATA_TAG=$DATA_NUM
+  EPOCH=1
+fi
+
+if [[ ${TASK} == 'multi_task' ]]; then
+  FULL_MODEL_TAG=${MODEL_TAG}_${DATA_TAG}_lr${LR}_s${16}
+else
+  FULL_MODEL_TAG=${MODEL_TAG}_${DATA_TAG}_lr${LR}_bs${BS}_src${SRC_LEN}_trg${TRG_LEN}_pat${PATIENCE}_e${EPOCH}
+fi
+
+
+if [[ ${SUB_TASK} == none ]]; then
+  OUTPUT_DIR=${MODEL_DIR}/${TASK}/${FULL_MODEL_TAG}
+else
+  OUTPUT_DIR=${MODEL_DIR}/${TASK}/${SUB_TASK}/${FULL_MODEL_TAG}
+fi
+
+CACHE_DIR=${OUTPUT_DIR}/cache_data
+RES_DIR=${OUTPUT_DIR}/prediction
+LOG=${OUTPUT_DIR}/train.log
+mkdir -p ${OUTPUT_DIR}
+mkdir -p ${CACHE_DIR}
+mkdir -p ${RES_DIR}
+
+if [[ $MODEL_TAG == roberta ]]; then
+  MODEL_TYPE=roberta
+  TOKENIZER=roberta-base
+  MODEL_PATH=roberta-base
+elif [[ $MODEL_TAG == codebert ]]; then
+  MODEL_TYPE=roberta
+  TOKENIZER=roberta-base
+  MODEL_PATH=microsoft/codebert-base
+elif [[ $MODEL_TAG == bart_base ]]; then
+  MODEL_TYPE=bart
+  TOKENIZER=facebook/bart-base
+  MODEL_PATH=facebook/bart-base
+elif [[ $MODEL_TAG == codet5_small ]]; then
+  MODEL_TYPE=codet5
+  TOKENIZER=Salesforce/codet5-small
+  MODEL_PATH=Salesforce/codet5-small
+elif [[ $MODEL_TAG == codet5_base ]]; then
+  MODEL_TYPE=codet5
+  TOKENIZER=Salesforce/codet5-base
+  MODEL_PATH=Salesforce/codet5-base
+elif [[ $MODEL_TAG == codet5_large ]]; then
+  MODEL_TYPE=codet5
+  TOKENIZER=Salesforce/codet5-large
+  MODEL_PATH=Salesforce/codet5-large
+fi
+
+
+if [[ ${TASK} == 'multi_task' ]]; then
+  RUN_FN=${WORKDIR}/run_multi_gen.py
+  MULTI_TASK_AUG='--max_steps '${16}' --save_steps '${17}' --log_steps '${18}
+elif [[ ${TASK} == 'clone' ]]; then
+  RUN_FN=${WORKDIR}/run_clone.py
+elif [[ ${TASK} == 'defect' ]] && [[ ${MODEL_TYPE} == 'roberta' ||  ${MODEL_TYPE} == 'bart' ]]; then
+  RUN_FN=${WORKDIR}/run_defect.py
+else
+  RUN_FN=${WORKDIR}/run_gen.py
+fi
+
+CUDA_VISIBLE_DEVICES=${GPU} \
+  python ${RUN_FN}  ${MULTI_TASK_AUG}   \
+  --do_train --do_eval --do_eval_bleu --do_test  \
+  --task ${TASK} --sub_task ${SUB_TASK} --model_type ${MODEL_TYPE} --data_num ${DATA_NUM}  \
+  --num_train_epochs ${EPOCH} --warmup_steps ${WARMUP} --learning_rate ${LR}e-5 --patience ${PATIENCE} \
+  --tokenizer_name=${TOKENIZER}  --model_name_or_path=${MODEL_PATH} --data_dir ${WORKDIR}/data  \
+  --cache_path ${CACHE_DIR}  --output_dir ${OUTPUT_DIR}  --summary_dir ${SUMMARY_DIR} \
+  --save_last_checkpoints --always_save_model --res_dir ${RES_DIR} --res_fn ${RES_FN} \
+  --train_batch_size ${BS} --eval_batch_size ${BS} --max_source_length ${SRC_LEN} --max_target_length ${TRG_LEN} \
+  2>&1 | tee ${LOG}

sh/run_exp.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python
+import os
+import argparse
+
+
+def get_cmd(task, sub_task, model_tag, gpu, data_num, bs, lr, source_length, target_length, patience, epoch, warmup,
+            model_dir, summary_dir, res_fn, max_steps=None, save_steps=None, log_steps=None):
+    if max_steps is None:
+        cmd_str = 'bash exp_with_args.sh %s %s %s %d %d %d %d %d %d %d %d %d %s %s %s' % \
+                  (task, sub_task, model_tag, gpu, data_num, bs, lr, source_length, target_length, patience, epoch,
+                   warmup, model_dir, summary_dir, res_fn)
+    else:
+        cmd_str = 'bash exp_with_args.sh %s %s %s %d %d %d %d %d %d %d %d %d %s %s %s %d %d %d' % \
+                  (task, sub_task, model_tag, gpu, data_num, bs, lr, source_length, target_length, patience, epoch,
+                   warmup, model_dir, summary_dir, res_fn, max_steps, save_steps, log_steps)
+    return cmd_str
+
+
+def get_args_by_task_model(task, sub_task, model_tag):
+    if task == 'translate':
+        # java-cs: Read 10300 examples, avg src len: 13, avg trg len: 15, max src len: 136, max trg len: 118
+        # [TOKENIZE] avg src len: 45, avg trg len: 56, max src len: 391, max trg len: 404
+        src_len = 320
+        trg_len = 256
+        epoch = 100
+        patience = 5
+    elif task == 'summarize':
+        # ruby: Read 24927 examples, avg src len: 66, avg trg len: 12, max src len: 501, max trg len: 146
+        # [TOKENIZE] avg src len: 100, avg trg len: 13, max src len: 1250, max trg len: 161
+        # Python: Read 251820 examples, avg src len: 100, avg trg len: 11, max src len: 512, max trg len: 222
+        # [TOKENIZE] avg src len: 142, avg trg len: 12, max src len: 2016, max trg len: 245
+        # Javascript: Read 58025 examples, avg src len: 114, avg trg len: 11, max src len: 512, max trg len: 165
+        # [TOKENIZE] avg src len: 136, avg trg len: 12, max src len: 3016, max trg len: 177
+        src_len = 256
+        trg_len = 128
+        epoch = 15
+        patience = 2
+    elif task == 'refine':
+        # small: Read 46680 examples, avg src len: 31, avg trg len: 28, max src len: 50, max trg len: 50
+        # [TOKENIZE] avg src len: 50, avg trg len: 45, max src len: 129, max trg len: 121
+        # medium:  Read 52364 examples, avg src len: 74, avg trg len: 73, max src len: 100, max trg len: 100
+        # [TOKENIZE] avg src len: 117, avg trg len: 114, max src len: 238, max trg len: 238
+        if sub_task == 'small':
+            src_len = 130
+            trg_len = 120
+        elif sub_task == 'medium':
+            src_len = 240
+            trg_len = 240
+        epoch = 50
+        patience = 5
+    elif task == 'concode':
+        # Read 100000 examples, avg src len: 71, avg trg len: 26, max src len: 567, max trg len: 140
+        # [TOKENIZE] avg src len: 213, avg trg len: 33, max src len: 2246, max trg len: 264
+        src_len = 320
+        trg_len = 150
+        epoch = 30
+        patience = 3
+    elif task == 'defect':
+        # Read 21854 examples, avg src len: 187, avg trg len: 1, max src len: 12195, max trg len: 1
+        # [TOKENIZE] avg src len: 597, avg trg len: 1, max src len: 41447, max trg len: 1
+        src_len = 512
+        trg_len = 3
+        epoch = 10
+        patience = 2
+    elif task == 'clone':
+        # Read 901028 examples, avg src len: 120, avg trg len: 123, max src len: 5270, max trg len: 5270
+        # [TOKENIZE] avg src len: 318, avg trg len: 323, max src len: 15111, max trg len: 15111
+        src_len = 400
+        trg_len = 400
+        epoch = 1
+        patience = 2
+
+    if 'codet5_small' in model_tag:
+        bs = 32
+        if task == 'summarize' or task == 'translate' or (task == 'refine' and sub_task == 'small'):
+            bs = 64
+        elif task == 'clone':
+            bs = 25
+    elif 'codet5_large' in model_tag:
+        bs = 8
+    else:
+        bs = 32
+        if task == 'translate':
+            bs = 25
+        elif task == 'summarize':
+            bs = 48
+        elif task == 'clone':
+            if model_tag in ['codebert', 'roberta']:
+                bs = 16
+            else:
+                bs = 10
+    lr = 5
+    if task == 'concode':
+        lr = 10
+    elif task == 'defect':
+        lr = 2
+    return bs, lr, src_len, trg_len, patience, epoch
+
+
+def run_one_exp(args):
+    bs, lr, src_len, trg_len, patience, epoch = get_args_by_task_model(args.task, args.sub_task, args.model_tag)
+    print('============================Start Running==========================')
+    cmd_str = get_cmd(task=args.task, sub_task=args.sub_task, model_tag=args.model_tag, gpu=args.gpu,
+                      data_num=args.data_num, bs=bs, lr=lr, source_length=src_len, target_length=trg_len,
+                      patience=patience, epoch=epoch, warmup=1000,
+                      model_dir=args.model_dir, summary_dir=args.summary_dir,
+                      res_fn='{}/{}_{}.txt'.format(args.res_dir, args.task, args.model_tag))
+    print('%s\n' % cmd_str)
+    os.system(cmd_str)
+
+
+def run_multi_task_exp(args):
+    # Total train data num = 1149722 (for all five tasks)
+    if 'codet5_small' in args.model_tag:
+        bs, lr, max_steps, save_steps, log_steps = 60, 5, 600000, 20000, 100
+    else:
+        bs, lr, max_steps, save_steps, log_steps = 25, 5, 800000, 20000, 100
+
+    if args.data_num != -1:
+        max_steps, save_steps, log_steps = 1000, 200, 50
+    print('============================Start Running==========================')
+    cmd_str = get_cmd(task='multi_task', sub_task='none', model_tag=args.model_tag, gpu=args.gpu,
+                      data_num=args.data_num, bs=bs, lr=lr, source_length=-1, target_length=-1,
+                      patience=-1, epoch=-1, warmup=1000,
+                      model_dir=args.model_dir, summary_dir=args.summary_dir,
+                      res_fn='{}/multi_task_{}.txt'.format(args.res_dir, args.model_tag),
+                      max_steps=max_steps, save_steps=save_steps, log_steps=log_steps)
+    print('%s\n' % cmd_str)
+    os.system(cmd_str)
+
+
+def get_sub_tasks(task):
+    if task == 'summarize':
+        sub_tasks = ['ruby', 'javascript', 'go', 'python', 'java', 'php']
+    elif task == 'translate':
+        sub_tasks = ['java-cs', 'cs-java']
+    elif task == 'refine':
+        sub_tasks = ['small', 'medium']
+    elif task in ['concode', 'defect', 'clone', 'multi_task']:
+        sub_tasks = ['none']
+    return sub_tasks
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_tag", type=str, default='codet5_base',
+                        choices=['roberta', 'codebert', 'bart_base', 'codet5_small', 'codet5_base', 'codet5_large'])
+    parser.add_argument("--task", type=str, default='summarize', choices=['summarize', 'concode', 'translate',
+                                                                          'refine', 'defect', 'clone', 'multi_task'])
+    parser.add_argument("--sub_task", type=str, default='ruby')
+    parser.add_argument("--res_dir", type=str, default='results', help='directory to save fine-tuning results')
+    parser.add_argument("--model_dir", type=str, default='saved_models', help='directory to save fine-tuned models')
+    parser.add_argument("--summary_dir", type=str, default='tensorboard', help='directory to save tensorboard summary')
+    parser.add_argument("--data_num", type=int, default=-1, help='number of data instances to use, -1 for full data')
+    parser.add_argument("--gpu", type=int, default=0, help='index of the gpu to use in a cluster')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.res_dir):
+        os.makedirs(args.res_dir)
+
+    assert args.sub_task in get_sub_tasks(args.task)
+    if args.task != 'multi_task':
+        run_one_exp(args)
+    else:
+        run_multi_task_exp(args)

tokenizer/apply_tokenizer.py

@@ -0,0 +1,17 @@
+from tokenizers import ByteLevelBPETokenizer
+
+tokenizer = ByteLevelBPETokenizer.from_file(
+    "./salesforce/codet5-vocab.json",
+    "./salesforce/codet5-merges.txt"
+)
+tokenizer.add_special_tokens([
+    "<pad>",
+    "<s>",
+    "</s>",
+    "<unk>",
+    "<mask>"
+])
+
+print(
+    tokenizer.encode("<s> hello <unk> Don't you love 🤗 Transformers <mask> yes . </s>").tokens
+)

tokenizer/train_tokenizer.py

@@ -0,0 +1,22 @@
+from tokenizers import ByteLevelBPETokenizer
+
+paths = ['train_code.txt', 'train_doc.txt']
+
+# Initialize a tokenizer
+tokenizer = ByteLevelBPETokenizer()
+
+# Customize training
+tokenizer.train(files=paths, vocab_size=32000, min_frequency=3, special_tokens=[
+    "<pad>",
+    "<s>",
+    "</s>",
+    "<unk>",
+    "<mask>"
+])
+
+# Save files to disk
+tokenizer.save_model("./salesforce", "codet5")
+
+print(
+    tokenizer.encode("<s> hello <unk> Don't you love 🤗 Transformers <mask> yes . </s>").tokens
+)

utils.py

@@ -0,0 +1,263 @@
+from torch.utils.data import TensorDataset
+import numpy as np
+import logging
+import os
+import random
+import torch
+import time
+from tqdm import tqdm
+from _utils import *
+
+logger = logging.getLogger(__name__)
+
+
+def load_and_cache_gen_data(args, filename, pool, tokenizer, split_tag, only_src=False, is_sample=False):
+    # cache the data into args.cache_path except it is sampled
+    # only_src: control whether to return only source ids for bleu evaluating (dev/test)
+    # return: examples (Example object), data (TensorDataset)
+    data_tag = '_all' if args.data_num == -1 else '_%d' % args.data_num
+    cache_fn = '{}/{}.pt'.format(args.cache_path, split_tag + ('_src' if only_src else '') + data_tag)
+
+    examples = read_examples(filename, args.data_num, args.task)
+
+    if is_sample:
+        examples = random.sample(examples, min(5000, len(examples)))
+    if split_tag == 'train':
+        calc_stats(examples, tokenizer, is_tokenize=True)
+    else:
+        calc_stats(examples)
+    if os.path.exists(cache_fn) and not is_sample:
+        logger.info("Load cache data from %s", cache_fn)
+        data = torch.load(cache_fn)
+    else:
+        if is_sample:
+            logger.info("Sample 5k data for computing bleu from %s", filename)
+        else:
+            logger.info("Create cache data into %s", cache_fn)
+        tuple_examples = [(example, idx, tokenizer, args, split_tag) for idx, example in enumerate(examples)]
+        features = pool.map(convert_examples_to_features, tqdm(tuple_examples, total=len(tuple_examples)))
+        all_source_ids = torch.tensor([f.source_ids for f in features], dtype=torch.long)
+        if split_tag == 'test' or only_src:
+            data = TensorDataset(all_source_ids)
+        else:
+            all_target_ids = torch.tensor([f.target_ids for f in features], dtype=torch.long)
+            data = TensorDataset(all_source_ids, all_target_ids)
+        if args.local_rank in [-1, 0] and not is_sample:
+            torch.save(data, cache_fn)
+    return examples, data
+
+
+def load_and_cache_clone_data(args, filename, pool, tokenizer, split_tag, is_sample=False):
+    cache_fn = '{}/{}.pt'.format(args.cache_path, split_tag + '_all' if args.data_num == -1 else '_%d' % args.data_num)
+    examples = read_examples(filename, args.data_num, args.task)
+    if is_sample:
+        examples = random.sample(examples, int(len(examples) * 0.1))
+
+    calc_stats(examples, tokenizer, is_tokenize=True)
+    if os.path.exists(cache_fn):
+        logger.info("Load cache data from %s", cache_fn)
+        data = torch.load(cache_fn)
+    else:
+        if is_sample:
+            logger.info("Sample 10 percent of data from %s", filename)
+        elif args.data_num == -1:
+            logger.info("Create cache data into %s", cache_fn)
+        tuple_examples = [(example, idx, tokenizer, args) for idx, example in enumerate(examples)]
+        features = pool.map(convert_clone_examples_to_features, tqdm(tuple_examples, total=len(tuple_examples)))
+        all_source_ids = torch.tensor([f.source_ids for f in features], dtype=torch.long)
+        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
+        data = TensorDataset(all_source_ids, all_labels)
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            torch.save(data, cache_fn)
+    return examples, data
+
+
+def load_and_cache_defect_data(args, filename, pool, tokenizer, split_tag, is_sample=False):
+    cache_fn = os.path.join(args.cache_path, split_tag)
+    examples = read_examples(filename, args.data_num, args.task)
+    if is_sample:
+        examples = random.sample(examples, int(len(examples) * 0.1))
+
+    calc_stats(examples, tokenizer, is_tokenize=True)
+    if os.path.exists(cache_fn):
+        logger.info("Load cache data from %s", cache_fn)
+        data = torch.load(cache_fn)
+    else:
+        if is_sample:
+            logger.info("Sample 10 percent of data from %s", filename)
+        elif args.data_num == -1:
+            logger.info("Create cache data into %s", cache_fn)
+        tuple_examples = [(example, idx, tokenizer, args) for idx, example in enumerate(examples)]
+        features = pool.map(convert_defect_examples_to_features, tqdm(tuple_examples, total=len(tuple_examples)))
+        # features = [convert_clone_examples_to_features(x) for x in tuple_examples]
+        all_source_ids = torch.tensor([f.source_ids for f in features], dtype=torch.long)
+        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
+        data = TensorDataset(all_source_ids, all_labels)
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            torch.save(data, cache_fn)
+    return examples, data
+
+
+def load_and_cache_multi_gen_data(args, pool, tokenizer, split_tag, only_src=False, is_sample=False):
+    cache_fn = os.path.join(args.cache_path, split_tag)
+    if os.path.exists(cache_fn) and not is_sample:
+        logger.info("Load cache data from %s", cache_fn)
+        examples_data_dict = torch.load(cache_fn)
+    else:
+        examples_data_dict = {}
+
+        task_list = ['summarize', 'translate', 'refine', 'concode', 'defect']
+        for task in task_list:
+            if task == 'summarize':
+                sub_tasks = ['ruby', 'javascript', 'go', 'python', 'java', 'php']
+            elif task == 'translate':
+                sub_tasks = ['java-cs', 'cs-java']
+            elif task == 'refine':
+                sub_tasks = ['small', 'medium']
+            else:
+                sub_tasks = ['none']
+            args.task = task
+            for sub_task in sub_tasks:
+                args.sub_task = sub_task
+                if task == 'summarize':
+                    args.max_source_length = 256
+                    args.max_target_length = 128
+                elif task == 'translate':
+                    args.max_source_length = 320
+                    args.max_target_length = 256
+                elif task == 'refine':
+                    if sub_task == 'small':
+                        args.max_source_length = 130
+                        args.max_target_length = 120
+                    else:
+                        args.max_source_length = 240
+                        args.max_target_length = 240
+                elif task == 'concode':
+                    args.max_source_length = 320
+                    args.max_target_length = 150
+                elif task == 'defect':
+                    args.max_source_length = 512
+                    args.max_target_length = 3  # as do not need to add lang ids
+
+                filename = get_filenames(args.data_dir, args.task, args.sub_task, split_tag)
+                examples = read_examples(filename, args.data_num, args.task)
+                if is_sample:
+                    examples = random.sample(examples, min(5000, len(examples)))
+                if split_tag == 'train':
+                    calc_stats(examples, tokenizer, is_tokenize=True)
+                else:
+                    calc_stats(examples)
+
+                tuple_examples = [(example, idx, tokenizer, args, split_tag) for idx, example in enumerate(examples)]
+                if args.data_num == -1:
+                    features = pool.map(convert_examples_to_features, tqdm(tuple_examples, total=len(tuple_examples)))
+                else:
+                    features = [convert_examples_to_features(x) for x in tuple_examples]
+                all_source_ids = torch.tensor([f.source_ids for f in features], dtype=torch.long)
+                if only_src:
+                    data = TensorDataset(all_source_ids)
+                else:
+                    all_target_ids = torch.tensor([f.target_ids for f in features], dtype=torch.long)
+                    data = TensorDataset(all_source_ids, all_target_ids)
+                examples_data_dict['{}_{}'.format(task, sub_task) if sub_task != 'none' else task] = (examples, data)
+
+        if args.local_rank in [-1, 0] and not is_sample:
+            torch.save(examples_data_dict, cache_fn)
+            logger.info("Save data into %s", cache_fn)
+    return examples_data_dict
+
+
+def get_filenames(data_root, task, sub_task, split=''):
+    if task == 'concode':
+        data_dir = '{}/{}'.format(data_root, task)
+        train_fn = '{}/train.json'.format(data_dir)
+        dev_fn = '{}/dev.json'.format(data_dir)
+        test_fn = '{}/test.json'.format(data_dir)
+    elif task == 'summarize':
+        data_dir = '{}/{}/{}'.format(data_root, task, sub_task)
+        train_fn = '{}/train.jsonl'.format(data_dir)
+        dev_fn = '{}/valid.jsonl'.format(data_dir)
+        test_fn = '{}/test.jsonl'.format(data_dir)
+    elif task == 'refine':
+        data_dir = '{}/{}/{}'.format(data_root, task, sub_task)
+        train_fn = '{}/train.buggy-fixed.buggy,{}/train.buggy-fixed.fixed'.format(data_dir, data_dir)
+        dev_fn = '{}/valid.buggy-fixed.buggy,{}/valid.buggy-fixed.fixed'.format(data_dir, data_dir)
+        test_fn = '{}/test.buggy-fixed.buggy,{}/test.buggy-fixed.fixed'.format(data_dir, data_dir)
+    elif task == 'translate':
+        data_dir = '{}/{}'.format(data_root, task)
+        if sub_task == 'cs-java':
+            train_fn = '{}/train.java-cs.txt.cs,{}/train.java-cs.txt.java'.format(data_dir, data_dir)
+            dev_fn = '{}/valid.java-cs.txt.cs,{}/valid.java-cs.txt.java'.format(data_dir, data_dir)
+            test_fn = '{}/test.java-cs.txt.cs,{}/test.java-cs.txt.java'.format(data_dir, data_dir)
+        else:
+            train_fn = '{}/train.java-cs.txt.java,{}/train.java-cs.txt.cs'.format(data_dir, data_dir)
+            dev_fn = '{}/valid.java-cs.txt.java,{}/valid.java-cs.txt.cs'.format(data_dir, data_dir)
+            test_fn = '{}/test.java-cs.txt.java,{}/test.java-cs.txt.cs'.format(data_dir, data_dir)
+    elif task == 'clone':
+        data_dir = '{}/{}'.format(data_root, task)
+        train_fn = '{}/train.txt'.format(data_dir)
+        dev_fn = '{}/valid.txt'.format(data_dir)
+        test_fn = '{}/test.txt'.format(data_dir)
+    elif task == 'defect':
+        data_dir = '{}/{}'.format(data_root, task)
+        train_fn = '{}/train.jsonl'.format(data_dir)
+        dev_fn = '{}/valid.jsonl'.format(data_dir)
+        test_fn = '{}/test.jsonl'.format(data_dir)
+    if split == 'train':
+        return train_fn
+    elif split == 'dev':
+        return dev_fn
+    elif split == 'test':
+        return test_fn
+    else:
+        return train_fn, dev_fn, test_fn
+
+
+def read_examples(filename, data_num, task):
+    read_example_dict = {
+        'summarize': read_summarize_examples,
+        'refine': read_refine_examples,
+        'translate': read_translate_examples,
+        'concode': read_concode_examples,
+        'clone': read_clone_examples,
+        'defect': read_defect_examples,
+    }
+    return read_example_dict[task](filename, data_num)
+
+
+def calc_stats(examples, tokenizer=None, is_tokenize=False):
+    avg_src_len = []
+    avg_trg_len = []
+    avg_src_len_tokenize = []
+    avg_trg_len_tokenize = []
+    for ex in examples:
+        if is_tokenize:
+            avg_src_len.append(len(ex.source.split()))
+            avg_trg_len.append(len(str(ex.target).split()))
+            avg_src_len_tokenize.append(len(tokenizer.tokenize(ex.source)))
+            avg_trg_len_tokenize.append(len(tokenizer.tokenize(str(ex.target))))
+        else:
+            avg_src_len.append(len(ex.source.split()))
+            avg_trg_len.append(len(str(ex.target).split()))
+    if is_tokenize:
+        logger.info("Read %d examples, avg src len: %d, avg trg len: %d, max src len: %d, max trg len: %d",
+                    len(examples), np.mean(avg_src_len), np.mean(avg_trg_len), max(avg_src_len), max(avg_trg_len))
+        logger.info("[TOKENIZE] avg src len: %d, avg trg len: %d, max src len: %d, max trg len: %d",
+                    np.mean(avg_src_len_tokenize), np.mean(avg_trg_len_tokenize), max(avg_src_len_tokenize),
+                    max(avg_trg_len_tokenize))
+    else:
+        logger.info("Read %d examples, avg src len: %d, avg trg len: %d, max src len: %d, max trg len: %d",
+                    len(examples), np.mean(avg_src_len), np.mean(avg_trg_len), max(avg_src_len), max(avg_trg_len))
+
+
+def get_elapse_time(t0):
+    elapse_time = time.time() - t0
+    if elapse_time > 3600:
+        hour = int(elapse_time // 3600)
+        minute = int((elapse_time % 3600) // 60)
+        return "{}h{}m".format(hour, minute)
+    else:
+        minute = int((elapse_time % 3600) // 60)
+        return "{}m".format(minute)