Lazyhope/python-clone-detection

1"""2Original work:3https://github.com/sangHa0411/CloneDetection/blob/main/utils/preprocessor.py4 5Copyright (c) 2022 Sangha Park(sangha110495), Young Jin Ahn(snoop2head)6 7All credits to the original authors.8"""9import re10import torch11from transformers import Pipeline12 13 14class FunctionPreprocessor:15    def get_function(self, code):16        results = []17        fn_list = re.findall("\ndef [a-zA-Z0-9_]+\(", code)18 19        for fn in fn_list:20            results.append(fn[4:-1].strip())21        return results22 23    def determine_function(self, code, function_name):24        num = len(re.findall("[^a-zA-Z]" + function_name + "[^a-zA-Z]", code))25        return False if num <= 1 else True26 27    def delete_function(self, code, name):28        start_id, _ = re.search("def " + name, code).span()29        ptr = start_id30 31        while ptr < len(code) - 1:32            if code[ptr] == "\n" and re.search("[a-zA-Z]", code[ptr + 1]) is not None:33                break34            ptr += 135 36        if ptr != len(code) - 1:37            end_id = ptr38            code = code[:start_id] + code[end_id:]39 40        return code41 42    def preprocess(self, code):43        code = "\n" + code44        fn_list = self.get_function(code)45        if len(fn_list) == 0:46            return code47 48        for fn in fn_list:49            flag = self.determine_function(code, fn)50 51            if flag == False:52                code = self.delete_function(code, fn)53 54        return code55 56 57class AnnotationPreprocessor:58    def search(self, sen_list, string):59        for i, sen in enumerate(sen_list):60            if string in sen:61                return i62        return -163 64    def delete_annotation_block(self, code, string):65        sens = [sen for sen in code.split("\n")]66 67        start_id = self.search(sens, string)68        end_id = self.search(sens[start_id + 1 :], string)69        if end_id != -1:70            end_id += start_id + 171            code = sens[:start_id] + sens[end_id + 1 :]72        else:73            code = sens[:start_id] + sens[start_id + 1 :]74 75        code = "\n".join(code)76        return code77 78    def delete_block(self, code, string):79        while string in code:80            code = self.delete_annotation_block(code, string)81        return code82 83    def delete_annotation(self, code):84        sens = code.split("\n")85 86        sens_processed = []87        for sen in sens:88            if "#" in sen:89                index = sen.index("#")90                sen = sen[:index]91            sens_processed.append(sen)92 93        return "\n".join(sens_processed)94 95    def delete_import(self, code):96        sens = code.split("\n")97 98        sens_processed = []99        for sen in sens:100            if "import" not in sen:101                sens_processed.append(sen)102 103        return "\n".join(sens_processed)104 105    def preprocess(self, code):106        code = self.delete_block(code, '"""')107        code = self.delete_block(code, "'''")108        code = self.delete_annotation(code)109        code = self.delete_import(code)110        code = re.sub("\s+", " ", code).strip()111        return code112 113 114def preprocessor(code, instance):115    processed_code = instance.preprocess(code)116    return processed_code if processed_code.strip() else code117 118 119def token_to_inputs(feature):120    inputs = {}121    for k, v in feature.items():122        inputs[k] = torch.tensor(v).unsqueeze(0)123 124    return inputs125 126 127class CloneDetectionPipeline(Pipeline):128    fn_preprocessor = FunctionPreprocessor()129    an_preprocessor = AnnotationPreprocessor()130 131    def _sanitize_parameters(self, **kwargs):132        preprocess_kwargs = {}133        return preprocess_kwargs, {}, {}134 135    def preprocess(self, inputs):136        code1 = inputs[0]137        code2 = inputs[1]138        if code1.strip() == "" or code2.strip() == "":139            ture_prob = float(code1.strip() == code2.strip())140            return {"skip": True, "output": {False: 1 - ture_prob, True: ture_prob}}141 142        code1 = preprocessor(143            preprocessor(code1, self.fn_preprocessor), self.an_preprocessor144        )145        code2 = preprocessor(146            preprocessor(code2, self.fn_preprocessor), self.an_preprocessor147        )148 149        feature1 = self.tokenizer(150            code1, code2, max_length=512, return_token_type_ids=False, truncation=True151        )152        feature2 = self.tokenizer(153            code2, code1, max_length=512, return_token_type_ids=False, truncation=True154        )155 156        return {157            "inputs1": token_to_inputs(feature1),158            "inputs2": token_to_inputs(feature2),159        }160 161    def _forward(self, model_inputs):162        if model_inputs.get("skip", False):163            return model_inputs164 165        inputs1 = model_inputs["inputs1"]166        inputs2 = model_inputs["inputs2"]167 168        logits1 = self.model(**inputs1).logits[0]169        logits2 = self.model(**inputs2).logits[0]170        logits = (logits1 + logits2) / 2171 172        return {"logits": logits}173 174    def postprocess(self, model_outputs):175        if model_outputs.get("skip", False):176            return model_outputs["output"]177 178        probs = model_outputs["logits"].softmax(-1).tolist()179 180        return {False: probs[0], True: probs[1]}181