app.py

app.py

@@ -0,0 +1,199 @@
+import pandas as pd
+import numpy as np
+import re
+import os
+import sys
+import random
+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, RandomSampler, SequentialSampler
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+import gradio as gr
+
+
+
+
+def greet(co):
+    code_text = []
+    while True:
+        code = co
+        if not code:
+            break
+        code_text.append(code)
+    '''
+    iter_num = int(
+        input('false alarm을 분류하기 위해서 입력할 코드의 갯수는 몇개인가요? (숫자만 입력하세요.) : '))
+    code_text = []
+    for _ in range(iter_num):
+        code = input('코드를 입력하세요 : ')
+        code_text.append(code)
+    '''
+    code_text = ' '.join(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)
+    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)
+    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.8 + pred_2 * 0.1 +
+                  pred_3 * 0.1 + pred_4 * 0.1).argmax()
+
+    return tot_result
+
+
+    
+# 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)