add files

Script/Exp_Script/ChatGPT/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)

Script/Exp_Script/ChatGPT/calculate_chatgpt_completion.py

@@ -0,0 +1,273 @@
+import os
+# from tree_sitter import Language, Parser
+# # import pandas as pd
+# import openpyxl
+import json
+import time
+import csv
+import pathlib
+import difflib
+import re
+from bleu import _bleu
+from fuzzywuzzy import fuzz
+import random
+import numpy as np
+from transformers import RobertaTokenizer
+#tokens = nltk.word_tokenize(sentence)
+import argparse
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+args = parser.parse_args()
+
+
+
+folder = str(pathlib.Path(__file__).parent.resolve())
+isa_type_dir = folder+"/../../../Dataset"
+src_dir = folder+f"/../../../Dataset/Code_Completion/{args.task}"
+dst_dir = folder+"/Result"
+
+train_lis = []
+valid_lis = []
+test_lis = []
+
+target_clf = {}
+def get_target_clf_list():
+    global target_clf
+    with open(isa_type_dir+"/comback_isa_type.csv","r",encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, l in enumerate(reader):
+            if l[1].lower() == "arc" or l[1].lower() == "riscv" or l[1].lower() == "nvptx":
+                continue
+            if l[0] + " " + l[2] not in target_clf.keys():
+                target_clf[l[0] + " " + l[2]] = [l[1]]
+            else:
+                target_clf[l[0] + " " + l[2]] += [l[1]]
+
+
+
+
+def Calculate_Completion():
+    get_target_clf_list()
+    print("############## Exp 2: Calculate ChatGPT Stmt Completion ################\n")
+    
+    test_lis = ["nvptx","arc","riscv"]
+
+
+    codellama_gcc_code = {}
+    codellama_llvm_code = {}
+
+    if args.task == "next_statement":
+        dst_file = dst_dir+"/Output/chatgpt_next_output_cleaned.csv"
+    else:
+        dst_file = dst_dir+"/Output/chatgpt_stmt_output_cleaned.csv"
+
+    
+
+    with open(dst_file, encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, row in enumerate(reader):
+            if row[0] == "GCC":
+                codellama_gcc_code[row[1] + " " + str(row[2])] = row[3]
+            else:
+                codellama_llvm_code[row[1] + " " + str(row[2])] = row[3]
+    avg_accuracy = {}
+    for comp_type in ["GCC", "LLVM"]:
+        for isa_type in ["GPU", "MPU", "CPU"]:
+            test_target_dic = {}
+            cnt_idx = 0
+            if comp_type == "GCC":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/riscv.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["riscv" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "riscv", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "riscv", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "riscv"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/nvptx.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["nvptx" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0          
+                        src_code = test_target_dic[k]
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis 
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "nvptx", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "nvptx", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "nvptx"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/arc.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["arc" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]    
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "arc", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "arc", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "arc"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+
+            if comp_type == "LLVM":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/RISCV.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["RISCV" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "RISCV", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "RISCV", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "RISCV"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/NVPTX.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["NVPTX" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]  
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "NVPTX", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "NVPTX", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "NVPTX"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/ARC.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["ARC" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "ARC", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "ARC", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "ARC"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+
+    return avg_accuracy
+
+
+
+
+if __name__ == "__main__":
+    with open(dst_dir + '/result.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Compiler Type", "Target", "Idx", "Exact Match", "Edit Didtance"])
+
+    avg_dic = Calculate_Completion()
+
+    for k in avg_dic:
+        print("########################")
+        
+        print(k)
+        print(" ".join(["Exact Match", "Edit Didtance"]))
+        print(" ".join(avg_dic[k]))

Script/Exp_Script/ChatGPT/calculate_chatgpt_gen.py

@@ -0,0 +1,384 @@
+import os
+# from tree_sitter import Language, Parser
+# # import pandas as pd
+# import openpyxl
+import json
+import time
+import csv
+import pathlib
+import difflib
+import re
+from bleu import _bleu
+from fuzzywuzzy import fuzz
+import random
+import numpy as np
+from transformers import RobertaTokenizer
+#tokens = nltk.word_tokenize(sentence)
+
+folder = str(pathlib.Path(__file__).parent.resolve())
+isa_type_dir = folder+"/../../../Dataset"
+src_dir = folder+"/../../../Dataset/Code_Generation"
+dst_dir = folder+"/Result"
+
+train_lis = []
+valid_lis = []
+test_lis = []
+
+target_clf = {}
+def get_target_clf_list():
+    global target_clf
+    with open(isa_type_dir+"/comback_isa_type.csv","r",encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, l in enumerate(reader):
+            if l[1].lower() == "arc" or l[1].lower() == "riscv" or l[1].lower() == "nvptx":
+                continue
+            if l[0] + " " + l[2] not in target_clf.keys():
+                target_clf[l[0] + " " + l[2]] = [l[1]]
+            else:
+                target_clf[l[0] + " " + l[2]] += [l[1]]
+
+
+def Calculate_Statements_Ratio(Src_List, Fork_Lis, src_name, fork_name):
+    src_code = ""
+    Fork_code = ""
+    idx = 0
+    cnt_stmt = 0.0
+    while idx < len(Src_List):
+        src_code += Src_List[idx].replace(src_name, "").replace(src_name.upper(), "")
+        if Src_List[idx] in [";", ":", "{", "}"]:
+            src_code += "\n"
+            cnt_stmt += 1
+        idx += 1
+    while idx < len(Fork_Lis):
+        Fork_code += Fork_Lis[idx].replace(fork_name, "").replace(fork_name.upper(), "")
+        if Fork_Lis[idx] in [";", ":", "{", "}"]:
+            Fork_code += "\n"
+        idx += 1
+        
+    code_same = 0
+    code_modi = 0
+    code_add = 0
+    diff_code = list(difflib.Differ().compare(src_code.splitlines(), Fork_code.splitlines()))
+    for idx, dv in enumerate(diff_code):
+        if dv[0] == '-':
+            if idx < len(diff_code) - 1 and diff_code[idx+1][0] == '?':
+                code_modi += 1
+            else:
+                code_add += 1
+        elif dv[0] == '+':
+            continue
+        elif dv[0] == '?':
+            continue
+                #vega_add -= 1
+        elif dv.strip().replace("\n", "") == '':
+            continue
+        else:
+            code_same += 1
+    return round(float(code_same) / cnt_stmt, 2)
+
+
+
+def Calculate_Gen():
+    get_target_clf_list()
+    print("############## Exp 2: Calculate ChatGPT ################\n")
+    
+    test_lis = ["nvptx","arc","riscv"]
+
+
+    chatgpt_gcc_code = {}
+    chatgpt_llvm_code = {}
+    avg_accuracy = {}
+
+    with open(dst_dir+"/chatgpt_gen_output.jsonl",encoding="utf-8") as f:
+        for idx, line in enumerate(f):
+            
+            js=json.loads(line)
+            if js["Compiler_Type"] == "GCC":
+                chatgpt_gcc_code[str(js["Target"]) + " " + js["idx"]] = js["Code"]
+            else:
+                chatgpt_llvm_code[str(js["Target"]) + " " + js["idx"]] = js["Code"]
+
+    for comp_type in ["GCC", "LLVM"]:
+        for isa_type in ["GPU", "MPU", "CPU"]:
+            target_lis = target_clf[comp_type + " " + isa_type]
+            test_target_dic = {}
+            cnt_idx = 0
+            if comp_type == "GCC":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/riscv.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["riscv" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("riscv", "")
+                        if k in chatgpt_gcc_code.keys():
+                            chat_code = " ".join(chatgpt_gcc_code[k]).replace("riscv", "").replace("RISCV", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_gcc_code[k], "riscv", "riscv")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "riscv", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "riscv", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "riscv"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/nvptx.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["nvptx" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("nvptx", "")
+                        if k in chatgpt_gcc_code.keys():
+                            chat_code = " ".join(chatgpt_gcc_code[k]).replace("nvptx", "").replace("NVPTX", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_gcc_code[k], "nvptx", "nvptx")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "nvptx", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "nvptx", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "nvptx"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/arc.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["arc" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("arc", "")
+                        if k in chatgpt_gcc_code.keys():
+                            chat_code = " ".join(chatgpt_gcc_code[k]).replace("arc", "").replace("ARC", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_gcc_code[k], "arc", "arc")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "arc", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "arc", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "arc"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+
+            if comp_type == "LLVM":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/RISCV.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["RISCV" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("RISCV", "")
+                        if k in chatgpt_llvm_code.keys():
+                            chat_code = " ".join(chatgpt_llvm_code[k]).replace("riscv", "").replace("RISCV", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_llvm_code[k], "riscv", "riscv")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "RISCV", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "RISCV", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "RISCV"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/NVPTX.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["NVPTX" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("NVPTX", "")
+                        if k in chatgpt_llvm_code.keys():
+                            chat_code = " ".join(chatgpt_llvm_code[k]).replace("nvptx", "").replace("NVPTX", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_llvm_code[k], "nvptx", "nvptx")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "NVPTX", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "NVPTX", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "NVPTX"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/ARC.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["ARC" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("ARC", "")
+                        if k in chatgpt_llvm_code.keys():
+                            chat_code = " ".join(chatgpt_llvm_code[k]).replace("arc", "").replace("ARC", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], chatgpt_llvm_code[k], "arc", "arc")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "ARC", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "ARC", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "ARC"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+
+    return avg_accuracy
+
+
+
+
+if __name__ == "__main__":
+    with open(dst_dir + '/result.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Compiler Type", "Target", "Idx", "BLEU4", "Exact Match", "Edit Didtance", "Stmt_Ratio"])
+
+    avg_dic = Calculate_Gen()
+
+    for k in avg_dic:
+        print("########################")
+        
+        print(k)
+        print(" ".join(["BLEU4", "Exact Match", "Edit Didtance", "Stmt_Ratio"]))
+        print(" ".join(avg_dic[k]))
+

Script/Exp_Script/Code-LLaMA/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)

Script/Exp_Script/Code-LLaMA/calculate_codellama_completion.py

@@ -0,0 +1,269 @@
+import os
+# from tree_sitter import Language, Parser
+# # import pandas as pd
+# import openpyxl
+import json
+import time
+import csv
+import pathlib
+import difflib
+import re
+from bleu import _bleu
+from fuzzywuzzy import fuzz
+import random
+import numpy as np
+from transformers import RobertaTokenizer
+#tokens = nltk.word_tokenize(sentence)
+import argparse
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+args = parser.parse_args()
+
+folder = str(pathlib.Path(__file__).parent.resolve())
+isa_type_dir = folder+"/../../../Dataset"
+src_dir = folder+f"/../../../Dataset/Code_Completion/{args.task}"
+dst_dir = folder+"/Result"
+
+train_lis = []
+valid_lis = []
+test_lis = []
+
+target_clf = {}
+def get_target_clf_list():
+    global target_clf
+    with open(isa_type_dir+"/comback_isa_type.csv","r",encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, l in enumerate(reader):
+            if l[1].lower() == "arc" or l[1].lower() == "riscv" or l[1].lower() == "nvptx":
+                continue
+            if l[0] + " " + l[2] not in target_clf.keys():
+                target_clf[l[0] + " " + l[2]] = [l[1]]
+            else:
+                target_clf[l[0] + " " + l[2]] += [l[1]]
+
+
+
+
+def Calculate_Completion():
+    get_target_clf_list()
+    print("############## Exp 2: Calculate Code-LLaMA Stmt Completion ################\n")
+    
+    test_lis = ["nvptx","arc","riscv"]
+
+
+    codellama_gcc_code = {}
+    codellama_llvm_code = {}
+
+    if args.task == "next_statement":
+        dst_file = dst_dir+"/Output/chatgpt_next_output_cleaned.csv"
+    else:
+        dst_file = dst_dir+"/Output/chatgpt_stmt_output_cleaned.csv"
+
+    with open(dst_file,encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, row in enumerate(reader):
+            if row[0] == "GCC":
+                codellama_gcc_code[row[1] + " " + str(row[2])] = row[3]
+            else:
+                codellama_llvm_code[row[1] + " " + str(row[2])] = row[3]
+    avg_accuracy = {}
+    for comp_type in ["GCC", "LLVM"]:
+        for isa_type in ["GPU", "MPU", "CPU"]:
+            test_target_dic = {}
+            cnt_idx = 0
+            if comp_type == "GCC":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/riscv.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["riscv" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "riscv", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "riscv", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "riscv"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/nvptx.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["nvptx" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0          
+                        src_code = test_target_dic[k]
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis 
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "nvptx", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "nvptx", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "nvptx"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/arc.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["arc" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_gcc_code.keys():
+                            chat_code = codellama_gcc_code[k]    
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "arc", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "arc", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "arc"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+
+            if comp_type == "LLVM":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/RISCV.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["RISCV" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "RISCV", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "RISCV", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "RISCV"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/NVPTX.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["NVPTX" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]  
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "NVPTX", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "NVPTX", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "NVPTX"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/ARC.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["ARC" + " " + str(cnt_idx)] = " ".join(dic["ground_truth"])
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        src_code = test_target_dic[k]
+                        if k in codellama_llvm_code.keys():
+                            chat_code = codellama_llvm_code[k]
+                            if chat_code.replace(" ", "") == src_code.replace(" ", ""):
+                                EM = 1
+                            edit_dis = fuzz.ratio(chat_code.replace(" ", ""), src_code.replace(" ", ""))
+                            total_ED += edit_dis
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "ARC", k.split(" ")[1], str(round(EM*100,2)), str(round(float(edit_dis),2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "ARC", "average", str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))])
+                        avg_accuracy[comp_type + " " + "ARC"] = [str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2))]
+
+    return avg_accuracy
+
+
+
+
+if __name__ == "__main__":
+    with open(dst_dir + '/result.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Compiler Type", "Target", "Idx", "Exact Match", "Edit Didtance"])
+
+    avg_dic = Calculate_Completion()
+
+    for k in avg_dic:
+        print("########################")
+        
+        print(k)
+        print(" ".join(["Exact Match", "Edit Didtance"]))
+        print(" ".join(avg_dic[k]))

Script/Exp_Script/Code-LLaMA/calculate_codellama_gen.py

@@ -0,0 +1,382 @@
+import os
+# from tree_sitter import Language, Parser
+# # import pandas as pd
+# import openpyxl
+import json
+import time
+import csv
+import pathlib
+import difflib
+import re
+from bleu import _bleu
+from fuzzywuzzy import fuzz
+import random
+import numpy as np
+from transformers import RobertaTokenizer
+#tokens = nltk.word_tokenize(sentence)
+
+folder = str(pathlib.Path(__file__).parent.resolve())
+isa_type_dir = folder+"/../../../Dataset"
+src_dir = folder+"/../../../Dataset/Code_Generation"
+dst_dir = folder+"/Result"
+
+train_lis = []
+valid_lis = []
+test_lis = []
+
+target_clf = {}
+def get_target_clf_list():
+    global target_clf
+    with open(isa_type_dir+"/comback_isa_type.csv","r",encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, l in enumerate(reader):
+            if l[1].lower() == "arc" or l[1].lower() == "riscv" or l[1].lower() == "nvptx":
+                continue
+            if l[0] + " " + l[2] not in target_clf.keys():
+                target_clf[l[0] + " " + l[2]] = [l[1]]
+            else:
+                target_clf[l[0] + " " + l[2]] += [l[1]]
+
+
+def Calculate_Statements_Ratio(Src_List, Fork_Lis, src_name, fork_name):
+    src_code = ""
+    Fork_code = ""
+    idx = 0
+    cnt_stmt = 0.0
+    while idx < len(Src_List):
+        src_code += Src_List[idx].replace(src_name, "").replace(src_name.upper(), "")
+        if Src_List[idx] in [";", ":", "{", "}"]:
+            src_code += "\n"
+            cnt_stmt += 1
+        idx += 1
+    while idx < len(Fork_Lis):
+        Fork_code += Fork_Lis[idx].replace(fork_name, "").replace(fork_name.upper(), "")
+        if Fork_Lis[idx] in [";", ":", "{", "}"]:
+            Fork_code += "\n"
+        idx += 1
+        
+    code_same = 0
+    code_modi = 0
+    code_add = 0
+    diff_code = list(difflib.Differ().compare(src_code.splitlines(), Fork_code.splitlines()))
+    for idx, dv in enumerate(diff_code):
+        if dv[0] == '-':
+            if idx < len(diff_code) - 1 and diff_code[idx+1][0] == '?':
+                code_modi += 1
+            else:
+                code_add += 1
+        elif dv[0] == '+':
+            continue
+        elif dv[0] == '?':
+            continue
+                #vega_add -= 1
+        elif dv.strip().replace("\n", "") == '':
+            continue
+        else:
+            code_same += 1
+    return round(float(code_same) / cnt_stmt, 2)
+
+
+
+def Calculate_Gen():
+    get_target_clf_list()
+    print("############## Exp 2: Calculate Code-LLaMA Gen ################\n")
+    
+    test_lis = ["nvptx","arc","riscv"]
+
+    avg_accuracy = {}
+    codellama_gcc_code = {}
+    codellama_llvm_code = {}
+
+    with open(dst_dir+"/codellama_gen_output.jsonl",encoding="utf-8") as f:
+        for idx, line in enumerate(f):
+            
+            js=json.loads(line)
+            if js["Compiler_Type"] == "GCC":
+                codellama_gcc_code[str(js["Target"]) + " " + js["idx"]] = js["Code"]
+            else:
+                codellama_llvm_code[str(js["Target"]) + " " + js["idx"]] = js["Code"]
+
+    for comp_type in ["GCC", "LLVM"]:
+        for isa_type in ["GPU", "MPU", "CPU"]:
+            target_lis = target_clf[comp_type + " " + isa_type]
+            test_target_dic = {}
+            cnt_idx = 0
+            if comp_type == "GCC":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/riscv.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["riscv" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("riscv", "")
+                        if k in codellama_gcc_code.keys():
+                            chat_code = " ".join(codellama_gcc_code[k]).replace("riscv", "").replace("RISCV", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_gcc_code[k], "riscv", "riscv")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "riscv", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "riscv", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "riscv"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/nvptx.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["nvptx" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("nvptx", "")
+                        if k in codellama_gcc_code.keys():
+                            chat_code = " ".join(codellama_gcc_code[k]).replace("nvptx", "").replace("NVPTX", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_gcc_code[k], "nvptx", "nvptx")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "nvptx", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "nvptx", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "nvptx"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/GCC/arc.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["arc" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("arc", "")
+                        if k in codellama_gcc_code.keys():
+                            chat_code = " ".join(codellama_gcc_code[k]).replace("arc", "").replace("ARC", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_gcc_code[k], "arc", "arc")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "arc", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "arc", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "arc"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+
+            if comp_type == "LLVM":
+                if isa_type == "CPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/RISCV.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["RISCV" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("RISCV", "")
+                        if k in codellama_llvm_code.keys():
+                            chat_code = " ".join(codellama_llvm_code[k]).replace("riscv", "").replace("RISCV", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_llvm_code[k], "riscv", "riscv")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "RISCV", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "RISCV", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "RISCV"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                if isa_type == "GPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/NVPTX.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["NVPTX" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("NVPTX", "")
+                        if k in codellama_llvm_code.keys():
+                            chat_code = " ".join(codellama_llvm_code[k]).replace("nvptx", "").replace("NVPTX", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_llvm_code[k], "nvptx", "nvptx")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "NVPTX", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "NVPTX", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "NVPTX"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+                
+                if isa_type == "MPU":
+                    cnt_idx = 0
+                    for line in open(src_dir + "/LLVM/ARC.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic["ARC" + " " + str(cnt_idx)] = dic["ground_truth"]
+                        cnt_idx += 1
+                    total_EM = 0.0
+                    total_ED = 0.0
+                    total_PoVS = 0.0
+                    total_BLEU4 = 0.0
+                    for k in test_target_dic.keys():
+                        edit_dis = 0.0
+                        EM = 0.0
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        src_code = " ".join(test_target_dic[k]).replace("ARC", "")
+                        if k in codellama_llvm_code.keys():
+                            chat_code = " ".join(codellama_llvm_code[k]).replace("arc", "").replace("ARC", "")
+                            stmt_mod = Calculate_Statements_Ratio(test_target_dic[k], codellama_llvm_code[k], "arc", "arc")
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(chat_code+'\n')
+                                f1.write(src_code+'\n')     
+                                if chat_code==src_code:
+                                    EM = 1
+                                edit_dis = fuzz.ratio(chat_code, src_code)
+                            if chat_code.strip() == "":
+                                bleu4 = 0
+                            else:
+                                bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                            total_BLEU4 += bleu4
+                            total_ED += edit_dis
+                            total_PoVS += stmt_mod
+                            total_EM += EM
+                            with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                                writer = csv.writer(file)
+                                writer.writerow([comp_type, "ARC", k.split(" ")[1], str(round(float(bleu4),2)), str(round(EM*100,2)), str(round(float(edit_dis),2)), str(round(float(stmt_mod)*100,2))])
+                        else:
+                            print(k)
+                    with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                        writer = csv.writer(file)
+                        writer.writerow([comp_type, "ARC", "average", str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))])
+                        avg_accuracy[comp_type + " " + "ARC"] = [str(round(float(total_BLEU4 / cnt_idx),2)), str(round((total_EM / cnt_idx)*100,2)), str(round(float(total_ED / cnt_idx),2)), str(round(float(total_PoVS / cnt_idx)*100,2))]
+    return avg_accuracy
+
+
+
+
+
+if __name__ == "__main__":
+    with open(dst_dir + '/result.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Compiler Type", "Target", "Idx", "BLEU4", "Exact Match", "Edit Didtance", "Stmt_Ratio"])
+
+    avg_dic = Calculate_Gen()
+
+    for k in avg_dic:
+        print("########################")
+        
+        print(k)
+        print(" ".join(["BLEU4", "Exact Match", "Edit Didtance", "Stmt_Ratio"]))
+        print(" ".join(avg_dic[k]))
+

Script/Exp_Script/ForkFlow/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)

Script/Exp_Script/ForkFlow/calculate_forkflow.py

@@ -0,0 +1,407 @@
+import os
+# from tree_sitter import Language, Parser
+# # import pandas as pd
+# import openpyxl
+import json
+import time
+import csv
+import pathlib
+import difflib
+import re
+from bleu import _bleu
+from fuzzywuzzy import fuzz
+import random
+import numpy as np
+from transformers import RobertaTokenizer
+#tokens = nltk.word_tokenize(sentence)
+
+folder = str(pathlib.Path(__file__).parent.resolve())
+isa_type_dir = folder+"/../../../Dataset"
+src_dir = folder+"/../../../Dataset/Code_Generation"
+dst_dir = folder+"/Result"
+
+train_lis = []
+valid_lis = []
+test_lis = []
+
+target_clf = {}
+def get_target_clf_list():
+    global target_clf
+    with open(isa_type_dir+"/comback_isa_type.csv","r",encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for idx, l in enumerate(reader):
+            if l[1].lower() == "arc" or l[1].lower() == "riscv" or l[1].lower() == "nvptx":
+                continue
+            if l[0] + " " + l[2] not in target_clf.keys():
+                target_clf[l[0] + " " + l[2]] = [l[1]]
+            else:
+                target_clf[l[0] + " " + l[2]] += [l[1]]
+
+
+def Calculate_Statements_Ratio(Src_List, Fork_Lis, src_name, fork_name):
+    src_code = ""
+    Fork_code = ""
+    idx = 0
+    cnt_stmt = 0.0
+    while idx < len(Src_List):
+        src_code += Src_List[idx].replace(src_name, "")
+        if Src_List[idx] in [";", ":", "{", "}"]:
+            src_code += "\n"
+            cnt_stmt += 1
+        idx += 1
+    while idx < len(Fork_Lis):
+        Fork_code += Fork_Lis[idx].replace(fork_name, "")
+        if Fork_Lis[idx] in [";", ":", "{", "}"]:
+            Fork_code += "\n"
+        idx += 1
+        
+    code_same = 0
+    code_modi = 0
+    code_add = 0
+    diff_code = list(difflib.Differ().compare(src_code.splitlines(), Fork_code.splitlines()))
+    for idx, dv in enumerate(diff_code):
+        if dv[0] == '-':
+            if idx < len(diff_code) - 1 and diff_code[idx+1][0] == '?':
+                code_modi += 1
+            else:
+                code_add += 1
+        elif dv[0] == '+':
+            continue
+        elif dv[0] == '?':
+            continue
+                #vega_add -= 1
+        elif dv.strip().replace("\n", "") == '':
+            continue
+        else:
+            code_same += 1
+    return round(float(code_same) / cnt_stmt, 2)
+
+
+
+def Calculate_Forkflow():
+    get_target_clf_list()
+    print("############## Exp 1: Calculate Fork-Flow ################\n")
+    
+    test_lis = ["nvptx","arc","riscv"]
+    for comp_type in ["GCC", "LLVM"]:
+        for isa_type in ["GPU", "MPU", "CPU"]:
+            max_ed = 0
+            avg_ed = 0
+            max_bleu4 = 0
+            avg_bleu4 = 0
+            avg_cnt = 0
+            target_lis = target_clf[comp_type + " " + isa_type]
+            test_target_dic = {}
+            cnt_idx = 0
+            if comp_type == "GCC":
+                if isa_type == "CPU":
+                    for line in open(src_dir + "/GCC/riscv.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("riscv", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    for tar in target_lis:
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("riscv", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "riscv", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "riscv", tar)
+                            
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+                            
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+                if isa_type == "GPU":
+                    for line in open(src_dir + "/GCC/nvptx.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("nvptx", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    for tar in target_lis:
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("nvptx", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "nvptx", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "nvptx", tar)
+                            
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+                if isa_type == "MPU":
+                    for line in open(src_dir + "/GCC/arc.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("arc", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    for tar in target_lis:
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("arc", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "arc", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "arc", tar)
+                            
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+            if comp_type == "LLVM":
+                if isa_type == "CPU":
+                    for line in open(src_dir + "/LLVM/RISCV.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("RISCV", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    for tar in target_lis:
+                        if tar == "RI5CY":
+                            continue
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("RISCV", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "RISCV", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "RISCV", tar)
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+                if isa_type == "GPU":
+                    for line in open(src_dir + "/LLVM/NVPTX.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("NVPTX", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    
+                    for tar in target_lis:
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("NVPTX", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "NVPTX", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "NVPTX", tar)
+                            
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+                if isa_type == "MPU":
+                    for line in open(src_dir + "/LLVM/ARC.jsonl", 'r'):
+                        dic = json.loads(line)
+                        test_target_dic[str(cnt_idx) + " " + dic["Func"].replace("ARC", "")] = dic["ground_truth"]
+                        cnt_idx += 1
+                    for tar in target_lis:
+                        edit_dis = 0.0
+                        EM = []
+                        bleu4 = 0.0
+                        stmt_mod = 0.0
+                        cnt = 0
+                        fork_target_dic = {}
+                        for line in open(src_dir + "/" + comp_type + "/" + tar + ".jsonl", 'r'):
+                            dic = json.loads(line)
+                            fork_target_dic[dic["Func"].replace(tar, "")] = dic["ground_truth"]
+
+                        for k in test_target_dic.keys():
+                            func = k.split(" ")[1]
+                            src_code = " ".join(test_target_dic[k]).replace("ARC", "")
+                            if func in fork_target_dic.keys():
+                                fork_code = " ".join(fork_target_dic[func]).replace(tar, "")
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], fork_target_dic[func], "ARC", tar)
+                            else:
+                                fork_code = ""
+                                stmt_mod += Calculate_Statements_Ratio(test_target_dic[k], [], "ARC", tar)
+                            with open(dst_dir+"/test.output",'w') as f, open(dst_dir+"/test.gold",'w') as f1:
+                                f.write(fork_code+'\n')
+                                f1.write(src_code+'\n')     
+                                EM.append(fork_code==src_code)
+                                edit_dis += fuzz.ratio(fork_code, src_code)
+                                avg_ed += fuzz.ratio(fork_code, src_code)
+                                cnt += 1
+                                avg_cnt += 1
+                            if fork_code.strip() == "":
+                                bleu4 += 0
+                            else:
+                                tmp_bleu4 = _bleu(dst_dir+"/test.gold", dst_dir+"/test.output")
+                                bleu4 += tmp_bleu4
+                                avg_bleu4 += tmp_bleu4
+                                
+                        
+                        with open(dst_dir + '/result.csv', 'a', newline='') as file:
+                            writer = csv.writer(file)
+                            writer.writerow([comp_type, isa_type, tar, str(round(float(bleu4)/cnt,2)), str(round(np.mean(EM)*100,2)), str(round(float(edit_dis)/cnt,2)), str(round(float(stmt_mod)*100/cnt,2))])
+                            if round(float(bleu4)/cnt,2) > max_bleu4:
+                                max_bleu4 = round(float(bleu4)/cnt,2)
+                            if round(float(edit_dis)/cnt,2) > max_ed:
+                                max_ed = round(float(edit_dis)/cnt,2)
+            print(comp_type + " " + isa_type)
+            print("Avg ED: " + str(round(float(avg_ed)/avg_cnt,2)))
+            print("Max ED: " + str(max_ed))
+            print("Avg BLEU4: " + str(round(float(avg_bleu4)/avg_cnt,2)))
+            print("Max BLEU4: " + str(max_bleu4))
+            print("\n\n")
+
+
+
+
+
+if __name__ == "__main__":
+    with open(dst_dir + '/result.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Compiler Type", "ISA Type", "Target", "BLEU4", "Exact Match", "Edit Didtance", "Stmt_Ratio"])
+    Calculate_Forkflow()
+
+
+

Script/Model/CodeBert/code-completion/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/CodeBert/code-completion/run_completion.py

@@ -0,0 +1,540 @@
+# 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 sys
+import pickle
+import torch
+import json
+
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+              RobertaConfig, RobertaModel, RobertaTokenizer)  
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target,
+                 max_src_len,
+                 max_tar_len
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+        self.max_src_len = max_src_len
+        self.max_tar_len = max_tar_len
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+            max_src_len = max(max_src_len, len(js["Template_token"]))
+
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+                max_tar_len = max(max_src_len, len(js["ground_truth"]))
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs,
+                        max_src_len = max_src_len,
+                        max_tar_len = max_tar_len
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)[:args.max_source_length-5]
+        source_tokens =[tokenizer.cls_token,tokenizer.sep_token]+source_tokens+["<mask>", tokenizer.sep_token]
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = ["<mask>"] + tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=384, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    # eval_examples = random.sample(eval_examples, int(len(eval_examples) / divide_number))
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            p.append(text)
+                model.train()
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+       
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+        model.train()
+        avg_acc = 0.0
+        avg_EM = 0.0
+        total = 0
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip() # post_process(ref.strip()).split(" ")
+            gt = gold.target.strip()
+            if pred == gt:
+                avg_EM += 1
+            avg_acc += fuzz.ratio(pred, gt)
+            res_list.append([pred, gt])
+            total += 1
+        dev_acc = round(avg_acc/total, 2)
+        dev_em = round(avg_EM/total, 4)
+
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/CodeBert/code-generation/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)

Script/Model/CodeBert/code-generation/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/CodeBert/code-generation/run_generation.py

@@ -0,0 +1,470 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from fuzzywuzzy import fuzz
+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)
+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__)
+divide_number = 3
+# 
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)
+        ts_v_tokens = tokenizer.tokenize(example.ts_v)
+        source_tokens =[tokenizer.cls_token]+source_tokens+[tokenizer.sep_token]+ts_v_tokens+[tokenizer.sep_token]
+        
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens[:args.max_source_length-5]) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = [tokenizer.cls_token]+target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=256, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids=source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            # print(text)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                edit_dis = 0
+                cnt_all = 0
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())
+                        edit_dis += fuzz.ratio(ref, gold.target)    
+                        res_list.append([ref,gold.target])
+                        cnt_all += 1
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                avg_edit_dis = float(edit_dis)/cnt_all
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt_all,2))))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+avg_edit_dis) / 2.0
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience =0
+                else:
+                    patience +=1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,2))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/CodeT5+/code-completion/run_completion.py

@@ -0,0 +1,525 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup, T5ForConditionalGeneration, AutoTokenizer)
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+# 
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+
+            # print(inputs)
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_ids = torch.LongTensor(tokenizer.encode(example.source, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+
+        source_mask = torch.ones_like(source_ids)   
+        #target
+        if stage=="test":
+            target = "None"
+        else:
+            target = example.target
+
+        target_ids = torch.LongTensor(tokenizer.encode(target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids) 
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=512, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    is_trust = False 
+    if "codet5p-220m" in args.model_name_or_path or "codet5p-770m" in args.model_name_or_path:
+        is_trust = False
+    else:
+        is_trust = True
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    # model.eval()        
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                p=[]
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).loss
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length) # module. for multi GPU
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+                        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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                        
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+        
+        
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length) # module. for multi GPU
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+        model.train()
+        edit_sim = 0.0
+        EM = 0.0
+        total = len(p)
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip()
+            gt = gold.target
+            edit_sim += fuzz.ratio(pred, gt)
+            if pred.split() == gt.split():
+                EM += 1
+            res_list.append([pred,gt])
+        dev_acc = round(edit_sim/total, 2)
+        dev_em = round(EM/total, 4)
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/CodeT5+/code-generation/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)

Script/Model/CodeT5+/code-generation/run_generation.py

@@ -0,0 +1,478 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from fuzzywuzzy import fuzz
+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, T5ForConditionalGeneration, AutoTokenizer)
+
+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__)
+divide_number = 3
+# 
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask      
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+
+        source_ids = torch.LongTensor(tokenizer.encode(example.source + tokenizer.pad_token + example.ts_v, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        
+        source_mask = torch.ones_like(source_ids)
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        
+        target_ids = torch.LongTensor(tokenizer.encode(example.target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids)
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=512, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    is_trust = False 
+    if "codet5p-220m" in args.model_name_or_path:
+        is_trust = False
+    else:
+        is_trust = True
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch              
+                    with torch.no_grad():
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())   
+                        res_list.append([ref,gold.target])
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+round(np.mean(EM)*100,2))
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,2))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/CodeT5+/new-target-completion/run_completion.py

@@ -0,0 +1,614 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup, T5ForConditionalGeneration, AutoTokenizer)
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+# 
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target,
+                 comp_type,
+                 tar_type
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+        self.comp_type = comp_type
+        self.tar_type = tar_type
+
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+
+            # print(inputs)
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            
+
+            comp_type = js["Compiler_Type"]
+            tar_type = js["Target"]
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs,
+                        comp_type = comp_type,
+                        tar_type = tar_type
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask,
+                 comp_type, tar_type
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask
+        self.comp_type = comp_type
+        self.tar_type = tar_type    
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_ids = torch.LongTensor(tokenizer.encode(example.source, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        # print(tokenizer.encode(example.source, 
+        #         add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        source_mask = torch.ones_like(source_ids)   
+        #target
+        if stage=="test":
+            target = "None"
+        else:
+            target = example.target
+
+        target_ids = torch.LongTensor(tokenizer.encode(target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids) 
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask,
+                 example.comp_type, example.tar_type
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--do_cpuonly", action='store_true',
+                        help="Whether CPU only training.")
+    parser.add_argument("--do_itr", action='store_true',
+                        help="Whether to itr training.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=512, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    is_trust = False 
+    if "codet5p-220m" in args.model_name_or_path:
+        is_trust = False
+    else:
+        is_trust = True
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        model_save_name = "Existing_Types/pytorch_model.bin"
+        if args.do_itr:
+            model_save_name = "pytorch_model.bin"
+        if args.do_cpuonly:
+            model_save_name = "New_Types/pytorch_model.bin"        
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/"+model_save_name))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/"+model_save_name))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/"+model_save_name))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/"+model_save_name))
+            
+    # model.eval()        
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+
+        file_name_pre = "New_Target_Completion"
+        file_name_post = "Existing_Types/train.jsonl"
+        if args.do_itr:
+            file_name_pre = "Iterative_Expansion_Completion"
+            file_name_post = "train.jsonl"
+        if args.do_cpuonly and not args.do_itr:
+            file_name_pre = "New_Target_Completion"
+            file_name_post = "New_Types/train.jsonl"
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + file_name_pre +'/statement_level/'+file_name_post)
+        else:
+            train_examples = read_examples(args.train_filename + file_name_pre +'/statement_level/'+file_name_post)
+
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    file_name_pre = "New_Target_Completion"
+                    file_name_post = "Existing_Types/valid.jsonl"
+                    if args.do_itr:
+                        file_name_pre = "Iterative_Expansion_Completion"
+                        file_name_post = "valid.jsonl"
+                    if args.do_cpuonly and not args.do_itr:
+                        file_name_pre = "New_Target_Completion"
+                        file_name_post = "New_Types/valid.jsonl"
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + file_name_pre +'/statement_level/'+file_name_post)
+                    else:
+                        eval_examples = read_examples(args.dev_filename + file_name_pre +'/statement_level/'+file_name_post)
+
+                    
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                p=[]
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).loss
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length) # module. for multi GPU
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+                        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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                        
+
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    model_save_name = "Exitsing_Types/pytorch_model.bin"
+                    if args.do_itr:
+                        model_save_name = "pytorch_model.bin"
+                    if args.do_cpuonly:
+                        model_save_name = "New_Types/pytorch_model.bin"
+                    output_model_file = os.path.join(output_dir, model_save_name)
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+            # print(output_dir)
+            # odel_to_load.load_state_dict(torch.load(output_dir))
+            model_save_name = "Existing_Types/pytorch_model.bin"
+            if args.do_itr and not args.do_cpuonly:
+                model_save_name = "pytorch_model.bin"
+            if args.do_itr and args.do_cpuonly:
+                args.load_model_path = "../../../../Saved_Models/CodeT5+/New_Target_Completion"
+                model_save_name = "New_Types/pytorch_model.bin"
+            if args.do_cpuonly:
+                model_save_name = "New_Types/pytorch_model.bin"          
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/"+model_save_name))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/"+model_save_name))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/"+model_save_name))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/"+model_save_name))
+        
+        file_name_pre = "New_Target_Completion"
+        file_name_post = "Existing_Types/test.jsonl"
+        if args.do_itr:
+            file_name_pre = "Iterative_Expansion_Completion"
+            file_name_post = "test.jsonl"
+        if args.do_cpuonly and not args.do_itr:
+            file_name_pre = "New_Target_Completion"
+            file_name_post = "New_Types/test.jsonl"
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, file_name_pre +'/statement_level/'+file_name_post)
+        else:
+            args.test_filename = os.path.join(args.test_filename, file_name_pre +'/next_statement/'+file_name_post)
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length) # module. for multi GPU
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+        model.train()
+        # edit_sim = 0.0
+        # EM = 0.0
+        # total = len(p)
+        gcc_dic = {"riscv":[0,0,0], "nvptx":[0,0,0], "arc":[0,0,0]}
+        llvm_dic = {"RISCV":[0,0,0], "NVPTX":[0,0,0], "ARC":[0,0,0],"RI5CY":[0,0,0]}
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip()
+            gt = gold.target
+            if gold.comp_type == "GCC":
+                gcc_dic[gold.tar_type][1] += fuzz.ratio(pred, gt)
+                gcc_dic[gold.tar_type][2] += 1
+                if pred.split() == gt.split():
+                    gcc_dic[gold.tar_type][0] += 1
+            if gold.comp_type == "LLVM":
+                llvm_dic[gold.tar_type][1] += fuzz.ratio(pred, gt)
+                llvm_dic[gold.tar_type][2] += 1
+                if pred.split() == gt.split():
+                    llvm_dic[gold.tar_type][0] += 1
+            res_list.append([pred,gt])
+        # dev_acc = round(edit_sim/total, 2)
+        # dev_em = round(EM/total, 4)
+
+        for k in gcc_dic.keys():
+            if gcc_dic[k][2] > 0:
+                dev_acc = round(1.0*gcc_dic[k][1] / gcc_dic[k][2], 2)
+                dev_em = round(100.0*gcc_dic[k][0] / gcc_dic[k][2], 4)
+                logger.info(" "+"#"*20)
+                logger.info("GCC %s:  %s = %s "%(k, "Edit Distance", str(dev_acc)))
+                logger.info("GCC %s:  %s = %s "%(k, "Exact Match Rate", str(dev_em)))
+                logger.info("  "+"*"*20)
+
+        for k in llvm_dic.keys():
+            if llvm_dic[k][2] > 0:
+                dev_acc = round(1.0*llvm_dic[k][1] / llvm_dic[k][2], 2)
+                dev_em = round(100.0*llvm_dic[k][0] / llvm_dic[k][2], 4)
+                logger.info(" "+"#"*20)
+                logger.info("LLVM %s:  %s = %s "%(k, "Edit Distance", str(dev_acc)))
+                logger.info("LLVM %s:  %s = %s "%(k, "Exact Match Rate", str(dev_em)))
+                logger.info("  "+"*"*20)
+
+
+        # if args.test_org:
+        #     output_dir = args.output_dir
+        # else:
+        #     if args.task == "statement_level":
+        #         output_dir = os.path.join(args.output_dir, 'statement_level/')
+        #     else:
+        #         output_dir = os.path.join(args.output_dir, 'next_statement/')
+        # result_file_name = "/test_result.jsonl"
+        # if args.do_itr:
+        #     result_file_name = "/test_result_itr.jsonl"
+        # if args.do_cpuonly:
+        #     result_file_name = "/test_result_cpu.jsonl"
+        # with open(output_dir + result_file_name, 'w') as wf:
+        #     for line in res_list:
+        #         dic = {}
+        #         dic["Pred"] = line[0]
+        #         dic["GT"] = line[1]
+        #         wf.write(json.dumps(dic))
+        #         wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/CodeT5+/new-target-generation/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)

Script/Model/CodeT5+/new-target-generation/run_generation.py

@@ -0,0 +1,546 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from fuzzywuzzy import fuzz
+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, T5ForConditionalGeneration, AutoTokenizer)
+
+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__)
+divide_number = 3
+# 
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 comp_type,
+                 tar_type
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+        self.comp_type = comp_type
+        self.tar_type = tar_type
+
+def read_examples(filename):
+    """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)
+
+            comp_type = js["Compiler_Type"]
+            tar_type = js["Target"]
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        comp_type = comp_type,
+                        tar_type = tar_type
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask,
+                 comp_type, tar_type
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask
+        self.comp_type = comp_type
+        self.tar_type = tar_type       
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+
+        source_ids = torch.LongTensor(tokenizer.encode(example.source + tokenizer.pad_token + example.ts_v, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        
+        source_mask = torch.ones_like(source_ids)
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        
+        target_ids = torch.LongTensor(tokenizer.encode(example.target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids)
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask,
+                 example.comp_type, example.tar_type
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=512, 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 training.")
+    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("--no_cuda", action='store_true',
+                        help="Avoid using CUDA when available") 
+    parser.add_argument("--do_cpuonly", action='store_true',
+                        help="Whether CPU only training.")
+    parser.add_argument("--do_itr", action='store_true',
+                        help="Whether to itr training.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    is_trust = False 
+    if "codet5p-220m" in args.model_name_or_path or "codet5p-770m" in args.model_name_or_path:
+        is_trust = False
+    else:
+        is_trust = True
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        model_save_name = "/Existing_Types/pytorch_model.bin"
+        if args.do_itr and not args.do_cpuonly:
+            model_save_name = "/pytorch_model.bin"
+        if args.do_itr and args.do_cpuonly:
+            model_save_name = "/New_Types/pytorch_model.bin"
+        if args.do_cpuonly :
+            model_save_name = "/New_Types/pytorch_model.bin"
+        logger.info("reload model from {}".format(args.load_model_path + model_save_name))
+        model.load_state_dict(torch.load(args.load_model_path + model_save_name))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch              
+                    with torch.no_grad():
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())   
+                        res_list.append([ref,gold.target])
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+round(np.mean(EM)*100,2)) / 2.0
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    model_save_name = "Existing_Types/pytorch_model.bin"
+                    if args.do_itr and not args.do_cpuonly:
+                        model_save_name = "pytorch_model.bin"
+                    if args.do_itr and args.do_cpuonly:
+                        model_save_name = "New_Types/pytorch_model.bin"
+                    if args.do_cpuonly :
+                        model_save_name = "New_Types/pytorch_model.bin"
+                    output_model_file = os.path.join(output_dir, model_save_name)
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        
+        if args.load_model_path is not None:
+            model_save_name = "Existing_Types/pytorch_model.bin"
+            if args.do_itr and not args.do_cpuonly:
+                model_save_name = "pytorch_model.bin"
+            if args.do_itr and args.do_cpuonly:
+                model_save_name = "New_Types/pytorch_model.bin"
+            if args.do_cpuonly :
+                model_save_name = "New_Types/pytorch_model.bin"      
+            checkpoint_prefix = model_save_name
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        gcc_dic = {"riscv":[0,0,0,0], "nvptx":[0,0,0,0], "arc":[0,0,0,0]}
+        llvm_dic = {"RISCV":[0,0,0,0], "NVPTX":[0,0,0,0], "ARC":[0,0,0,0],"RI5CY":[0,0,0,0]}
+                
+
+        for ref,gold in zip(p,eval_examples):
+            res_list.append([ref,gold.target])
+            predictions.append(ref)
+            with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')
+            pred = ref.strip()
+            gt = gold.target   
+            if gold.comp_type == "GCC":
+                gcc_dic[gold.tar_type][1] += fuzz.ratio(pred, gt)
+                gcc_dic[gold.tar_type][2] += _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output") 
+                gcc_dic[gold.tar_type][3] += 1
+                if pred.split() == gt.split():
+                    gcc_dic[gold.tar_type][0] += 1
+            if gold.comp_type == "LLVM":
+                llvm_dic[gold.tar_type][1] += fuzz.ratio(pred, gt)
+                llvm_dic[gold.tar_type][2] += _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output") 
+                llvm_dic[gold.tar_type][3] += 1
+                if pred.split() == gt.split():
+                    llvm_dic[gold.tar_type][0] += 1
+        
+
+
+        for k in gcc_dic.keys():
+            if gcc_dic[k][3] > 0:
+                dev_acc = round(1.0*gcc_dic[k][1] / gcc_dic[k][3], 2)
+                dev_em = round(100.0*gcc_dic[k][0] / gcc_dic[k][3], 4)
+                dev_b4 = round(1.0*gcc_dic[k][2] / gcc_dic[k][3], 2)
+                logger.info(" "+"#"*20)
+                logger.info("GCC %s:  %s = %s "%(k, "Edit Distance", str(dev_acc)))
+                logger.info("GCC %s:  %s = %s "%(k, "Exact Match Rate", str(dev_em)))
+                logger.info("GCC %s:  %s = %s "%(k, "BLEU4", str(dev_b4)))
+                logger.info("  "+"*"*20)
+
+        for k in llvm_dic.keys():
+            if llvm_dic[k][3] > 0:
+                dev_acc = round(1.0*llvm_dic[k][1] / llvm_dic[k][3], 2)
+                dev_em = round(100.0*llvm_dic[k][0] / llvm_dic[k][3], 4)
+                dev_b4 = round(1.0*llvm_dic[k][2] / llvm_dic[k][3], 2)
+                logger.info(" "+"#"*20)
+                logger.info("LLVM %s:  %s = %s "%(k, "Edit Distance", str(dev_acc)))
+                logger.info("LLVM %s:  %s = %s "%(k, "Exact Match Rate", str(dev_em)))
+                logger.info("LLVM %s:  %s = %s "%(k, "BLEU4", str(dev_b4)))
+                logger.info("  "+"*"*20)
+
+
+
+
+        # result_file_name = "/test_result.jsonl"
+        # if args.do_itr:
+        #     result_file_name = "/test_result_itr.jsonl"
+        # if args.do_cpuonly:
+        #     result_file_name = "/test_result_cpu.jsonl"
+        # with open(args.output_dir + result_file_name, 'w') as wf:
+        #     for line in res_list:
+        #         dic = {}
+        #         dic["Pred"] = line[0]
+        #         dic["GT"] = line[1]
+        #         wf.write(json.dumps(dic))
+        #         wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/CodeT5/code-completion/run_completion.py

@@ -0,0 +1,543 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup, T5ForConditionalGeneration, RobertaTokenizer)
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_ids = torch.LongTensor(tokenizer.encode(example.source, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        source_mask = torch.ones_like(source_ids)   
+        #target
+        if stage=="test":
+            target = "None"
+        else:
+            target = example.target
+        target_ids = torch.LongTensor(tokenizer.encode(target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids) 
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=512, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                        
+
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    # eval_examples = random.sample(eval_examples, int(len(eval_examples) / divide_number))
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids, source_mask, _, _ = batch              
+                    with torch.no_grad():
+                        # preds = model(source_ids)
+                        # 1 card -- model.gen
+                        # multicard -- model.module.gen
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+                model.train()
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+        
+        
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+        model.train()
+        edit_sim = 0.0
+        EM = 0.0
+        total = len(p)
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip()
+            gt = gold.target
+            edit_sim += fuzz.ratio(pred, gt)
+            if pred.split() == gt.split():
+                EM += 1
+            res_list.append([pred,gt])
+        dev_acc = round(edit_sim/total, 2)
+        dev_em = round(EM/total, 4)
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/CodeT5/code-generation/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)

Script/Model/CodeT5/code-generation/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/CodeT5/code-generation/run_generation.py

@@ -0,0 +1,478 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from accelerate import Accelerator
+from fuzzywuzzy import fuzz
+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, T5ForConditionalGeneration, RobertaTokenizer)
+
+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__)
+divide_number = 3
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+            # print(" ".join(js['natrual_language']))
+            # print(",".join(js['TS_V_token']))
+            # print(" ".join(js["ground_truth"]))
+            # print("###########################################")
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask      
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+
+        source_ids = torch.LongTensor(tokenizer.encode(example.source + tokenizer.pad_token + example.ts_v, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        
+        source_mask = torch.ones_like(source_ids)
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        
+        target_ids = torch.LongTensor(tokenizer.encode(example.target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids)
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=512, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path)
+
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch              
+                    with torch.no_grad():
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())   
+                        res_list.append([ref,gold.target])
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+round(np.mean(EM)*100,2))
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,2))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/GraphCodeBert/code-completion/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/GraphCodeBert/code-completion/run_completion.py

@@ -0,0 +1,545 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+              RobertaConfig, RobertaModel, RobertaTokenizer)  
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target,
+                 max_src_len,
+                 max_tar_len
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+        self.max_src_len = max_src_len
+        self.max_tar_len = max_tar_len
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+            max_src_len = max(max_src_len, len(js["Template_token"]))
+
+            # print(inputs)
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+                max_tar_len = max(max_src_len, len(js["ground_truth"]))
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs,
+                        max_src_len = max_src_len,
+                        max_tar_len = max_tar_len
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)[:args.max_source_length-5]
+        source_tokens =[tokenizer.cls_token,tokenizer.sep_token]+source_tokens+["<mask>", tokenizer.sep_token]
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = ["<mask>"] + tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=256, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    # eval_examples = random.sample(eval_examples, int(len(eval_examples) / divide_number))
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            p.append(text)
+                model.train()
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+        
+        
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+        model.train()
+        avg_acc = 0.0
+        avg_EM = 0.0
+        total = 0
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip() # post_process(ref.strip()).split(" ")
+            gt = gold.target.strip()
+            if pred == gt:
+                avg_EM += 1
+            avg_acc += fuzz.ratio(pred, gt)
+            res_list.append([pred, gt])
+            total += 1
+        dev_acc = round(avg_acc/total, 2)
+        dev_em = round(avg_EM/total, 4)
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/GraphCodeBert/code-generation/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)

Script/Model/GraphCodeBert/code-generation/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/GraphCodeBert/code-generation/run_generation.py

@@ -0,0 +1,474 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from fuzzywuzzy import fuzz
+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)
+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__)
+divide_number = 3
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)
+        ts_v_tokens = tokenizer.tokenize(example.ts_v)
+        source_tokens =[tokenizer.cls_token]+source_tokens+[tokenizer.sep_token]+ts_v_tokens+[tokenizer.sep_token]
+        
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens[:args.max_source_length-5]) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = [tokenizer.cls_token]+target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=256, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids=source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            # print(text)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                edit_dis = 0
+                cnt_all = 0
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())
+                        edit_dis += fuzz.ratio(ref, gold.target)  
+                        res_list.append([ref,gold.target])
+                        cnt_all += 1
+
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                avg_edit_dis = float(edit_dis)/cnt_all
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt_all,2))))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+avg_edit_dis) / 2.0
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,2))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/NatGen/code-completion/run_completion.py

@@ -0,0 +1,520 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup, T5ForConditionalGeneration, AutoTokenizer)
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+# 
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+
+            # print(inputs)
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_ids = torch.LongTensor(tokenizer.encode(example.source, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+
+        source_mask = torch.ones_like(source_ids)   
+        #target
+        if stage=="test":
+            target = "None"
+        else:
+            target = example.target
+
+        target_ids = torch.LongTensor(tokenizer.encode(target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids) 
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=512, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+                p=[]
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).loss
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+                        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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                        
+
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break   
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+        
+        
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+        model.train()
+        edit_sim = 0.0
+        EM = 0.0
+        total = len(p)
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip()
+            gt = gold.target
+            edit_sim += fuzz.ratio(pred, gt)
+            if pred.split() == gt.split():
+                EM += 1
+            res_list.append([pred,gt])
+        dev_acc = round(edit_sim/total, 2)
+        dev_em = round(EM/total, 4)
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+
+    
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/NatGen/code-generation/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)

Script/Model/NatGen/code-generation/run_generation.py

@@ -0,0 +1,477 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from tqdm import tqdm, trange
+from torch.nn.utils.rnn import pad_sequence
+from fuzzywuzzy import fuzz
+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, T5ForConditionalGeneration, AutoTokenizer)
+
+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__)
+divide_number = 3
+# 
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_mask = source_mask
+        self.target_ids = target_ids
+        self.target_mask = target_mask      
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+
+        source_ids = torch.LongTensor(tokenizer.encode(example.source + tokenizer.pad_token + example.ts_v, 
+                add_special_tokens=True, max_length=args.max_source_length, truncation=True))
+        
+        source_mask = torch.ones_like(source_ids)
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        
+        target_ids = torch.LongTensor(tokenizer.encode(example.target, 
+                add_special_tokens=True, max_length=args.max_target_length, truncation=True))
+        target_mask = torch.ones_like(target_ids)
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids, source_mask,
+                 target_ids, target_mask
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=512, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path) 
+    
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = pad_sequence([f.source_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in train_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in train_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in train_features], batch_first=True, padding_value=0)
+        train_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,source_mask,target_ids,target_mask = batch
+                # loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                loss = model(input_ids=source_ids, attention_mask=source_mask.gt(0), 
+                                  labels=target_ids, decoder_attention_mask=target_mask.gt(0)).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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch
+                    with torch.no_grad():
+                        loss = model(input_ids=source_ids, attention_mask=source_mask, 
+                                    labels=target_ids, decoder_attention_mask=target_mask).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
+                        eval_loss += loss.item()
+                        tokens_num += 1
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+                    all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+                    all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+                    eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,source_mask,target_ids,target_mask = batch              
+                    with torch.no_grad():
+                        preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+                            
+                        # convert ids to text
+                        for pred in preds:
+                            # print(pred)
+                            text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())   
+                        res_list.append([ref,gold.target])
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+round(np.mean(EM)*100,2))
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                #     patience =0
+                # else:
+                #     patience +=1
+                    # if patience == -1:
+                    #     break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = pad_sequence([f.source_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_source_mask = pad_sequence([f.source_mask for f in eval_features], batch_first=True, padding_value=0)
+        all_target_ids = pad_sequence([f.target_ids for f in eval_features], batch_first=True, padding_value=tokenizer.pad_token_id)
+        all_target_mask = pad_sequence([f.target_mask for f in eval_features], batch_first=True, padding_value=0)
+        eval_data = TensorDataset(all_source_ids,all_source_mask,all_target_ids,all_target_mask)
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids, source_mask, _, _  = batch 
+            with torch.no_grad():
+                preds = model.module.generate(source_ids, attention_mask=source_mask, use_cache=True, 
+                                        num_beams=args.beam_size, max_new_tokens =args.max_target_length)
+            for pred in preds:
+                # print(pred)
+                text = tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,2))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+

Script/Model/UnixCoder/code-completion/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/UnixCoder/code-completion/run_completion.py

@@ -0,0 +1,543 @@
+# 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 sys
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+from fuzzywuzzy import fuzz
+import re
+import multiprocessing
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+              RobertaConfig, RobertaModel, RobertaTokenizer)  
+
+divide_number = 2
+cpu_cont = 16
+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__)
+
+# 
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 target,
+                 max_src_len,
+                 max_tar_len
+                 ):
+        self.idx = idx
+        self.source = source
+        self.target = target
+        self.max_src_len = max_src_len
+        self.max_tar_len = max_tar_len
+
+def read_examples(filename):
+    """Read examples from filename."""
+    examples=[]
+    
+    with open(filename,encoding="utf-8") as f:
+        max_src_len = 0
+        max_tar_len = 0
+        for idx, line in enumerate(f):
+            
+            js=json.loads(line)
+            inputs = " ".join(js["Template_token"][1:])
+            max_src_len = max(max_src_len, len(js["Template_token"]))
+
+            # print(inputs)
+            if "ground_truth" in js:
+                outputs = " ".join(js["ground_truth"])
+                max_tar_len = max(max_src_len, len(js["ground_truth"]))
+            else:
+                outputs = inputs
+            if 'Idx' in js:
+                idx = js['Idx']  
+            examples.append(
+                Example(
+                        idx = idx,
+                        source = inputs,
+                        target = outputs,
+                        max_src_len = max_src_len,
+                        max_tar_len = max_tar_len
+                ) 
+            )
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)[:args.max_source_length-5]
+        source_tokens =[tokenizer.cls_token,"<encoder-decoder>",tokenizer.sep_token]+source_tokens+["<mask0>",tokenizer.sep_token]
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = ["<mask0>"]+target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        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("--task", default=None, type=str, required=True,
+                        help="Task Type: statement_level, next_statement" )
+
+    parser.add_argument("--train_filename", default="../../Dataset/", type=str, 
+                        help="The train filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--dev_filename", default="../../Dataset/", type=str, 
+                        help="The dev filename. Should contain the .jsonl files for this task.")
+    parser.add_argument("--test_filename", default="../../Dataset/", 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="", 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 training.")
+    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("--test_org", action='store_true',
+                        help="Whether to run eval on org model.")
+    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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.")
+    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("--max_target_length", default=128, type=int,
+                        help="")
+    parser.add_argument("--max_source_length", default=512, type=int,
+                        help="")
+    parser.add_argument("--train_steps", default=-1, type=int,
+                        help="")
+    parser.add_argument("--warmup_steps", default=0, 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=20240124,
+                        help="random seed for initialization")
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+
+    if args.load_model_path is not None:
+        if args.task == "statement_level":
+            logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+        else:
+            logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            model.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+            
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        if args.task == "statement_level":
+            train_examples = read_examples(args.train_filename + "/Code_Completion/statement_level/train.jsonl")
+        else:
+            train_examples = read_examples(args.train_filename + "/Code_Completion/next_statement/train.jsonl")
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            if args.do_eval:
+                #Eval model with dev dataset   
+                                
+                if 'dev_loss' in dev_dataset:
+                    eval_examples,eval_data = dev_dataset['dev_loss']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+                res_list = []
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    if args.task == "statement_level":
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/statement_level/valid.jsonl")
+                    else:
+                        eval_examples = read_examples(args.dev_filename + "/Code_Completion/next_statement/valid.jsonl")
+                    # eval_examples = random.sample(eval_examples, int(len(eval_examples) / divide_number))
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            p.append(text)
+                model.train()
+                EM = 0.0
+                edit_sim = 0.0    
+                total = len(p)
+                token_accuracy = 0
+                for ref,gold in zip(p,eval_examples):
+                    pred = ref.strip()
+                    gt = gold.target
+                    edit_sim += fuzz.ratio(pred, gt)
+                    if pred.split() == gt.split():
+                        EM += 1
+                    res_list.append([pred,gt])
+                dev_acc = round(EM/total*100, 2)
+                # logger.info("  %s = %s "%("loss",round(np.mean(dev_losses),4)))
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("EM Acc",str(dev_acc)))
+                logger.info("  %s = %s "%("Edit Distance",str(round(edit_sim/total, 2))))
+                logger.info("  "+"*"*20)
+
+                if dev_acc > best_score:
+                    best_score = dev_acc
+                    # Save best checkpoint for best bleu
+                    if args.task == "statement_level":
+                        output_dir = os.path.join(args.output_dir, 'statement_level/')
+                    else:
+                        output_dir = os.path.join(args.output_dir, 'next_statement/')
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+
+                if args.task == "statement_level":
+                    output_dir = os.path.join(args.output_dir, 'statement_level/')
+                else:
+                    output_dir = os.path.join(args.output_dir, 'next_statement/')
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                
+    if args.do_test:
+        res_list = []
+        output_dir2 = ""
+        
+        if args.load_model_path is not None:
+            model_to_load = model.module if hasattr(model, 'module') else model
+      
+            if args.task == "statement_level":
+                logger.info("reload model from {}".format(args.load_model_path + "/statement_level/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/statement_level/pytorch_model.bin"))
+            else:
+                logger.info("reload model from {}".format(args.load_model_path + "/next_statement/pytorch_model.bin"))
+                model_to_load.load_state_dict(torch.load(args.load_model_path + "/next_statement/pytorch_model.bin"))
+        
+        
+        if args.task == "statement_level":
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/statement_level/test.jsonl')
+        else:
+            args.test_filename = os.path.join(args.test_filename, 'Code_Completion/next_statement/test.jsonl')
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+        model.train()
+        avg_acc = 0.0
+        avg_EM = 0.0
+        total = 0
+        for ref,gold in zip(p,eval_examples):
+            pred = ref.strip() # post_process(ref.strip()).split(" ")
+            gt = gold.target.strip()
+            if pred == gt:
+                avg_EM += 1
+            avg_acc += fuzz.ratio(pred, gt)
+            res_list.append([pred, gt])
+            total += 1
+        dev_acc = round(avg_acc/total, 2)
+        dev_em = round(avg_EM/total, 6)
+        logger.info("  %s = %s "%("Test Token Avg Edit Distance",str(dev_acc)))
+        logger.info("  %s = %s "%("Test Token Avg Exact Match Rate",str(dev_em)))
+        logger.info("  "+"*"*20)
+        if args.test_org:
+            output_dir = args.output_dir
+        else:
+            if args.task == "statement_level":
+                output_dir = os.path.join(args.output_dir, 'statement_level/')
+            else:
+                output_dir = os.path.join(args.output_dir, 'next_statement/')
+
+        with open(output_dir + "/test_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+                
+if __name__ == "__main__":
+    main()
+
+
+

Script/Model/UnixCoder/code-generation/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)

Script/Model/UnixCoder/code-generation/model.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+import copy
+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((1024, 1024), dtype=torch.uint8)).view(1,1024, 1024)
+        )
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.encoder.embeddings.word_embeddings.weight
+        self.lsm = nn.LogSoftmax(dim=-1)
+        
+        self.beam_size = beam_size
+        self.max_length = max_length
+        self.sos_id = sos_id
+        self.eos_id = eos_id       
+        
+    def forward(self, source_ids, target_ids=None):   
+        if target_ids is None:
+            return self.generate(source_ids)
+        
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)  
+        ids = torch.cat((source_ids,target_ids),-1)
+        mask = self.bias[:,source_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+        mask = mask & ids[:,None,:].ne(1)
+
+        out = self.decoder(target_ids,attention_mask=mask,past_key_values=encoder_output.past_key_values).last_hidden_state
+        lm_logits = self.lm_head(out)
+        # Shift so that tokens < n predict n
+        active_loss = target_ids[..., 1:].ne(1).view(-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
+    
+    def generate(self, source_ids):
+        mask = source_ids.ne(1)[:,None,:]*source_ids.ne(1)[:,:,None]
+        encoder_output = self.encoder(source_ids,attention_mask=mask,use_cache=True)        
+        preds = []       
+        zero = torch.cuda.LongTensor(1).fill_(0)   
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        for i in range(source_ids.shape[0]):
+            context = [[x[i:i+1,:,:source_len[i]].repeat(self.beam_size,1,1,1) for x in y] 
+                     for y in encoder_output.past_key_values]
+            beam = Beam(self.beam_size,self.sos_id,self.eos_id)
+            input_ids = beam.getCurrentState()
+            context_ids = source_ids[i:i+1,:source_len[i]].repeat(self.beam_size,1)
+            for _ in range(self.max_length): 
+                if beam.done():
+                    break
+
+                ids = torch.cat((context_ids,input_ids),-1)
+                mask = self.bias[:,context_ids.size(-1):ids.size(-1),:ids.size(-1)].bool()
+                mask = mask & ids[:,None,:].ne(1)
+                out = self.decoder(input_ids,attention_mask=mask,past_key_values=context).last_hidden_state
+                hidden_states = out[:,-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
+        

Script/Model/UnixCoder/code-generation/run_generation.py

@@ -0,0 +1,467 @@
+# 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 sys
+from bleu import _bleu
+import pickle
+import torch
+import json
+import random
+import logging
+import argparse
+import numpy as np
+from io import open
+from itertools import cycle
+import torch.nn as nn
+from model import Seq2Seq
+from tqdm import tqdm, trange
+from fuzzywuzzy import fuzz
+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)
+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__)
+divide_number = 3
+
+
+class Example(object):
+    """A single training/test example."""
+    def __init__(self,
+                 idx,
+                 source,
+                 ts_v,
+                 target,
+                 ):
+        self.idx = idx
+        self.source = source
+        self.ts_v = ts_v
+        self.target = target
+
+def read_examples(filename):
+    """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)
+
+            examples.append(
+                Example(
+                        idx = idx,
+                        source=" ".join(js['natrual_language']),
+                        ts_v = ",".join(js['TS_V_token']),
+                        target = " ".join(js["ground_truth"][1:-1]),
+                        ) 
+            )
+            
+    return examples
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 example_id,
+                 source_ids,
+                 target_ids,
+    ):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids     
+        
+def convert_examples_to_features(examples, tokenizer, args,stage=None):
+    features = []
+    for example_index, example in enumerate(examples):
+        #source
+        source_tokens = tokenizer.tokenize(example.source)
+        ts_v_tokens = tokenizer.tokenize(example.ts_v)
+        source_tokens =[tokenizer.cls_token,"<encoder-decoder>",tokenizer.sep_token]+source_tokens+[tokenizer.sep_token]+ts_v_tokens+["<mask0>",tokenizer.sep_token]
+        
+        source_ids =  tokenizer.convert_tokens_to_ids(source_tokens[:args.max_source_length-5]) 
+        padding_length = args.max_source_length - len(source_ids)
+        source_ids+=[tokenizer.pad_token_id]*padding_length
+ 
+        #target
+        if stage=="test":
+            target_tokens = tokenizer.tokenize("None")
+        else:
+            target_tokens = tokenizer.tokenize(example.target)[:args.max_target_length-2]
+        target_tokens = ["<mask0>"]+target_tokens+[tokenizer.sep_token]            
+        target_ids = tokenizer.convert_tokens_to_ids(target_tokens)
+        padding_length = args.max_target_length - len(target_ids)
+        target_ids+=[tokenizer.pad_token_id]*padding_length
+   
+
+       
+        features.append(
+            InputFeatures(
+                 example_index,
+                 source_ids,
+                 target_ids,
+            )
+        )
+    return features
+
+
+
+def set_seed(seed=20240124):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+        
+def main():
+    parser = argparse.ArgumentParser()
+
+    ## Required parameters  
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
+                        help="Path to pre-trained model: e.g. roberta-base" )
+    parser.add_argument("--load_model_path", default=None, type=str,
+                        help="Path to trained model" )
+    parser.add_argument("--output_dir", default=None, type=str, required=True,
+                        help="The output directory where the model predictions and checkpoints will be written.")   
+  
+    ## 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("--max_source_length", default=256, 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=512, 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 training.")
+    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("--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("--num_train_epochs", default=3, type=int,
+                        help="Total number of training epochs to perform.") 
+    parser.add_argument('--seed', type=int, default=20240124,
+                        help="random seed for initialization")
+    
+    # print arguments
+    args = parser.parse_args()
+    # set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    # set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    # Set seed
+    set_seed(args.seed)
+    # make dir if output_dir not exist
+    if os.path.exists(args.output_dir) is False:
+        os.makedirs(args.output_dir)
+
+    # build model
+    tokenizer = RobertaTokenizer.from_pretrained(args.model_name_or_path)
+    config = RobertaConfig.from_pretrained(args.model_name_or_path)
+    # import！！！you must set is_decoder as True for generation
+    config.is_decoder = True
+    encoder = RobertaModel.from_pretrained(args.model_name_or_path,config=config) 
+
+    model = Seq2Seq(encoder=encoder,decoder=encoder,config=config,
+                  beam_size=args.beam_size,max_length=args.max_target_length,
+                  sos_id=tokenizer.convert_tokens_to_ids(["<mask0>"])[0],eos_id=tokenizer.sep_token_id)
+    
+    logger.info("Training/evaluation parameters %s", args)
+    if args.load_model_path is not None:
+        logger.info("reload model from {}".format(args.load_model_path + "/pytorch_model.bin"))
+        model.load_state_dict(torch.load(args.load_model_path + "/pytorch_model.bin"))
+    model.to(args.device)   
+    
+    if args.n_gpu > 1:
+        # multi-gpu training
+        model = torch.nn.DataParallel(model)
+
+    if args.do_train:
+        # Prepare training data loader
+        train_examples = read_examples(args.train_filename)
+        train_features = convert_examples_to_features(train_examples, tokenizer,args,stage='train')
+        all_source_ids = torch.tensor([f.source_ids for f in train_features], dtype=torch.long)
+        all_target_ids = torch.tensor([f.target_ids for f in train_features], dtype=torch.long) 
+        train_data = TensorDataset(all_source_ids,all_target_ids)
+        train_sampler = RandomSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
+
+
+        # 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=int(len(train_dataloader)*args.num_train_epochs*0.1),
+                                                    num_training_steps=len(train_dataloader)*args.num_train_epochs)
+    
+        #Start training
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", len(train_examples))
+        logger.info("  Batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+        
+
+        model.train()
+        patience, best_score, losses, dev_dataset = 0, 0, [], {}
+        for epoch in range(args.num_train_epochs):
+            for idx,batch in enumerate(train_dataloader):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids,target_ids = batch
+                loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)
+
+                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
+                    
+                losses.append(loss.item())
+                loss.backward()
+                if len(losses) % args.gradient_accumulation_steps == 0:
+                    #Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    if len(losses) // args.gradient_accumulation_steps % 100 == 0:
+                        logger.info("epoch {} step {} loss {}".format(epoch,
+                                                     len(losses)//args.gradient_accumulation_steps,
+                                                     round(np.mean(losses[-100*args.gradient_accumulation_steps:]),4)))
+            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 = read_examples(args.dev_filename)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='dev')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+                    all_target_ids = torch.tensor([f.target_ids for f in eval_features], dtype=torch.long)   
+                    eval_data = TensorDataset(all_source_ids,all_target_ids)   
+                    dev_dataset['dev_loss' ]= eval_examples,eval_data
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                logger.info("\n***** Running evaluation *****")
+                logger.info("  Num examples = %d", len(eval_examples))
+                logger.info("  Batch size = %d", args.eval_batch_size)
+
+                #Start Evaling model
+                model.eval()
+                eval_loss,tokens_num = 0,0
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids,target_ids = batch                  
+
+                    with torch.no_grad():
+                        _,loss,num = model(source_ids=source_ids,target_ids=target_ids)     
+                    eval_loss += loss.sum().item()
+                    tokens_num += num.sum().item()
+                #Pring loss of dev dataset    
+                model.train()
+                eval_loss = eval_loss / tokens_num
+                result = {'eval_ppl': round(np.exp(eval_loss),5)}
+                for key in sorted(result.keys()):
+                    logger.info("  %s = %s", key, str(result[key]))
+                logger.info("  "+"*"*20)   
+
+                #Calculate bleu  
+                if 'dev_bleu' in dev_dataset:
+                    eval_examples,eval_data=dev_dataset['dev_bleu']
+                else:
+                    eval_examples = read_examples(args.dev_filename)
+                    # eval_examples = random.sample(eval_examples)
+                    eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+                    all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long) 
+                    eval_data = TensorDataset(all_source_ids)   
+                    dev_dataset['dev_bleu'] = eval_examples,eval_data
+
+                eval_sampler = SequentialSampler(eval_data)
+                eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+                model.eval() 
+                p=[]
+                for batch in eval_dataloader:
+                    batch = tuple(t.to(device) for t in batch)
+                    source_ids = batch[0]                  
+                    with torch.no_grad():
+                        preds = model(source_ids) 
+                        # convert ids to text
+                        for pred in preds:
+                            t = pred[0].cpu().numpy()
+                            t = list(t)
+                            if 0 in t:
+                                t = t[:t.index(0)]
+                            text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                            # print(text)
+                            p.append(text)
+
+                model.train()
+                predictions = []
+                res_list = []
+                EM = []
+                is_gened = False
+                with open(args.output_dir+"/dev.output",'w') as f, open(args.output_dir+"/dev.gold",'w') as f1:
+                    for ref,gold in zip(p,eval_examples):
+                        predictions.append(ref)
+                        if len(ref) > 0:
+                            is_gened = True
+                        f.write(ref+'\n')
+                        f1.write(gold.target+'\n')     
+                        EM.append(ref.split()==gold.target.split())   
+                        res_list.append([ref,gold.target])
+                if is_gened:
+                    dev_bleu = _bleu(args.output_dir+"/dev.gold", args.output_dir+"/dev.output")
+                else:
+                    dev_bleu = 0
+                logger.info("  %s = %s "%("Epoch",str(epoch)))
+                logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+                logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,2))))
+                logger.info("  "+"*"*20)    
+                dev_score = (dev_bleu+round(np.mean(EM)*100,2)) / 2.0
+                if dev_score>best_score:
+                    best_score=dev_score
+                    # Save best checkpoint for best bleu
+                    output_dir = args.output_dir
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)
+                    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+                    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                    torch.save(model_to_save.state_dict(), output_model_file)
+                    patience = 0
+                else:
+                    patience += 1
+                    if patience == 3:
+                        break
+                output_dir = args.output_dir
+                with open(output_dir + "/last_training_result.jsonl", 'w') as wf:
+                    for line in res_list:
+                        dic = {}
+                        dic["Pred"] = line[0]
+                        dic["GT"] = line[1]
+                        wf.write(json.dumps(dic))
+                        wf.write("\n")
+                    
+                logger.info("  Best score:%s",best_score)
+                logger.info("  "+"*"*20)
+    if args.do_test:
+        res_list = []
+        if args.load_model_path is not None:
+            checkpoint_prefix = 'pytorch_model.bin'
+            output_dir = os.path.join(args.output_dir, checkpoint_prefix)  
+            model_to_load = model.module if hasattr(model, 'module') else model  
+            model_to_load.load_state_dict(torch.load(output_dir))
+        
+        
+
+        eval_examples = read_examples(args.test_filename)
+        eval_features = convert_examples_to_features(eval_examples, tokenizer, args,stage='test')
+        all_source_ids = torch.tensor([f.source_ids for f in eval_features], dtype=torch.long)
+        eval_data = TensorDataset(all_source_ids)   
+
+        # Calculate bleu
+        eval_sampler = SequentialSampler(eval_data)
+        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+        model.eval() 
+        p=[]
+        for batch in tqdm(eval_dataloader,total=len(eval_dataloader)):
+            batch = tuple(t.to(device) for t in batch)
+            source_ids = batch[0]                  
+            with torch.no_grad():
+                preds = model(source_ids)   
+                # convert ids to text
+                for pred in preds:
+                    t = pred[0].cpu().numpy()
+                    t = list(t)
+                    if 0 in t:
+                        t = t[:t.index(0)]
+                    text = tokenizer.decode(t,clean_up_tokenization_spaces=False)
+                    p.append(text)
+
+        predictions=[]
+        EM = []
+        edit_dis = 0
+        cnt = 0
+        with open(args.output_dir+"/test.output",'w') as f, open(args.output_dir+"/test.gold",'w') as f1:
+            for ref,gold in zip(p,eval_examples):
+                res_list.append([ref,gold.target])
+                predictions.append(ref)
+                f.write(ref+'\n')
+                f1.write(gold.target+'\n')     
+                EM.append(ref.split()==gold.target.split())
+                edit_dis += fuzz.ratio(ref, gold.target)
+                cnt += 1
+
+        dev_bleu = _bleu(args.output_dir+"/test.gold", args.output_dir+"/test.output")   
+        logger.info("  %s = %s "%("bleu-4",str(dev_bleu)))
+        logger.info("  %s = %s "%("EM",str(round(np.mean(EM)*100,4))))
+        logger.info("  %s = %s "%("Edit Distance",str(round(float(edit_dis)/cnt,4))))
+        logger.info("  "+"*"*20)
+
+
+        with open(args.output_dir + "/last_training_result.jsonl", 'w') as wf:
+            for line in res_list:
+                dic = {}
+                dic["Pred"] = line[0]
+                dic["GT"] = line[1]
+                wf.write(json.dumps(dic))
+                wf.write("\n")
+                
+if __name__ == "__main__":
+    main()
+
+