app.py

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)