import pandas as pd import numpy as np import re import os import sys import transformers from transformers import AutoTokenizer, AutoModelForSequenceClassification from transformers import RobertaTokenizer, RobertaForSequenceClassification import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from transformers import T5Tokenizer, T5ForConditionalGeneration import gradio as gr def is_false_alarm(code_text): code_text = re.sub('\/\*[\S\s]*\*\/', '', code_text) code_text = re.sub('\/\/.*', '', code_text) code_text = re.sub('(\\\\n)+', '\\n', code_text) # 1. CFA-CodeBERTa-small.pt -> CodeBERTa-small-v1 finetunig model path = os.getcwd() + '/models/CFA-CodeBERTa-small.pt' tokenizer = AutoTokenizer.from_pretrained("huggingface/CodeBERTa-small-v1") input_ids = tokenizer.encode( code_text, max_length=512, truncation=True, padding='max_length') input_ids = torch.tensor([input_ids]) model = RobertaForSequenceClassification.from_pretrained( path, num_labels=2) model.to('cpu') pred_1 = model(input_ids)[0].detach().cpu().numpy()[0] # model(input_ids)[0].argmax().detach().cpu().numpy().item() # 2. CFA-codebert-c.pt -> codebert-c finetuning model path = os.getcwd() + '/models/CFA-codebert-c.pt' tokenizer = AutoTokenizer.from_pretrained(path) input_ids = tokenizer(code_text, padding=True, max_length=512, truncation=True, return_token_type_ids=True)['input_ids'] input_ids = torch.tensor([input_ids]) model = AutoModelForSequenceClassification.from_pretrained( path, num_labels=2) model.to('cpu') pred_2 = model(input_ids)[0].detach().cpu().numpy()[0] # 3. CFA-codebert-c-v2.pt -> undersampling + codebert-c finetuning model path = os.getcwd() + '/models/CFA-codebert-c-v2.pt' tokenizer = RobertaTokenizer.from_pretrained(path) input_ids = tokenizer(code_text, padding=True, max_length=512, truncation=True, return_token_type_ids=True)['input_ids'] input_ids = torch.tensor([input_ids]) model = RobertaForSequenceClassification.from_pretrained( path, num_labels=2) model.to('cpu') pred_3 = model(input_ids)[0].detach().cpu().numpy() # 4. codeT5 finetuning model path = os.getcwd() + '/models/CFA-codeT5' model_params = { # model_type: t5-base/t5-large "MODEL": path, "TRAIN_BATCH_SIZE": 8, # training batch size "VALID_BATCH_SIZE": 8, # validation batch size "VAL_EPOCHS": 1, # number of validation epochs "MAX_SOURCE_TEXT_LENGTH": 512, # max length of source text "MAX_TARGET_TEXT_LENGTH": 3, # max length of target text "SEED": 2022, # set seed for reproducibility } data = pd.DataFrame({'code': [code_text]}) pred_4 = T5Trainer( dataframe=data, source_text="code", model_params=model_params ) pred_4 = int(pred_4[0]) # ensemble tot_result = (pred_1 * 0.1 + pred_2 * 0.1 + pred_3 * 0.7 + pred_4 * 0.1).argmax() if tot_result == 0: return "false positive !!" else: return "true positive !!" # codeT5 class YourDataSetClass(Dataset): def __init__( self, dataframe, tokenizer, source_len, source_text): self.tokenizer = tokenizer self.data = dataframe self.source_len = source_len # self.summ_len = target_len # self.target_text = self.data[target_text] self.source_text = self.data[source_text] def __len__(self): return len(self.source_text) def __getitem__(self, index): source_text = str(self.source_text[index]) source_text = " ".join(source_text.split()) source = self.tokenizer.batch_encode_plus( [source_text], max_length=self.source_len, pad_to_max_length=True, truncation=True, padding="max_length", return_tensors="pt", ) source_ids = source["input_ids"].squeeze() source_mask = source["attention_mask"].squeeze() return { "source_ids": source_ids.to(dtype=torch.long), "source_mask": source_mask.to(dtype=torch.long), } def validate(epoch, tokenizer, model, device, loader): model.eval() predictions = [] with torch.no_grad(): for _, data in enumerate(loader, 0): ids = data['source_ids'].to(device, dtype=torch.long) mask = data['source_mask'].to(device, dtype=torch.long) generated_ids = model.generate( input_ids=ids, attention_mask=mask, max_length=150, num_beams=2, repetition_penalty=2.5, length_penalty=1.0, early_stopping=True ) preds = [tokenizer.decode( g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids] if ((preds != '0') | (preds != '1')): preds = '0' predictions.extend(preds) return predictions def T5Trainer(dataframe, source_text, model_params, step="test",): torch.manual_seed(model_params["SEED"]) # pytorch random seed np.random.seed(model_params["SEED"]) # numpy random seed torch.backends.cudnn.deterministic = True tokenizer = T5Tokenizer.from_pretrained(model_params["MODEL"]) model = T5ForConditionalGeneration.from_pretrained(model_params["MODEL"]) model = model.to('cpu') dataframe = dataframe[[source_text]] val_dataset = dataframe val_set = YourDataSetClass( val_dataset, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"], source_text) val_params = { 'batch_size': model_params["VALID_BATCH_SIZE"], 'shuffle': False, 'num_workers': 0 } val_loader = DataLoader(val_set, **val_params) for epoch in range(model_params["VAL_EPOCHS"]): predictions = validate(epoch, tokenizer, model, 'cpu', val_loader) return predictions ################################################################################# '''demo = gr.Interface( fn = greet, inputs = "text", outputs= "number") demo.launch(share=True) ''' with gr.Blocks() as demo1: gr.Markdown( """

False-Alarm-Detector

""") gr.Markdown( """ 정적 분석기를 통해 오류라고 보고된 C언어 코드의 함수를 입력하면, 오류가 True-positive 인지 False-positive 인지 분류 해 주는 프로그램입니다. """) ''' with gr.Accordion(label='모델에 대한 설명 ( 여기를 클릭 하시오. )',open=False): gr.Markdown( """ 총 3개의 모델을 사용하였다. 1. codeBERTa-small-v1 - codeBERTa-small-v1 설명 2. codeBERT - C - codeBERT - C 설명 3. codeT5 - codeT5 설명 """ ) ''' with gr.Row(): with gr.Column(): inputs = gr.Textbox( lines=10, placeholder="코드를 입력하시오.", label='Code') with gr.Row(): btn = gr.Button("결과 출력") with gr.Column(): output = gr.Text(label='Result') btn.click(fn=is_false_alarm, inputs=inputs, outputs=output) if __name__ == "__main__": demo1.launch()