codebert.py

import json
from flask import Flask, request, jsonify
from flask_cors import cross_origin
from flask_cors import CORS
import torch
import torch.nn as nn
import random
import os
import numpy as np
from openprompt.data_utils import InputExample
from transformers import AdamW, get_linear_schedule_with_warmup
import logging
from tqdm import tqdm, trange
import pyflakes.api
import pyflakes.reporter
import ast

logger = logging.getLogger(__name__)
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)


def read_answers(filename):
    answers = []
    with open(filename, encoding="utf-8") as f:
        data = json.load(f)
        for js in data:
            #line = line.strip()
            #js = json.loads(line)
            # print(js)
            # code = js['func']
            # target = js['target']
            example = InputExample(guid=js['target'], text_a=js['func'])
            answers.append(example)
    return answers

def read_answers2(filename):
    answers = []
    with open(filename, encoding="utf-8") as f:
        data = json.load(f)
        for js in data:
            #line = line.strip()
            #js = json.loads(line)
            # print(js)
            # code = js['func']
            # target = js['target']
            text_a = js['func']
            print(text_a)
            example = InputExample(text_a)
            answers.append(example)
    return answers


def set_seed(seed=52):
    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


train_dataset = read_answers('train_data.json')
valid_dataset = read_answers('valid_data.json')
test_dataset = read_answers('test_data.json')
# print(len(dataset), dataset[:5])
classes = ['negative', 'positive']
from openprompt.plms import load_plm

plm, tokenizer, model_config, WrapperClass = load_plm("roberta", "microsoft/codebert-base")
from openprompt.prompts import ManualTemplate, SoftTemplate, MixedTemplate

promptTemplate = MixedTemplate(
    model=plm,
    text='The code {"placeholder":"text_a"} is {"mask"}.',
    tokenizer=tokenizer,
)
from openprompt.prompts import ManualVerbalizer

promptVerbalizer = ManualVerbalizer(
    classes=classes,
    label_words={
        "negative": ["clean", "better"],
        "positive": ["defective", "bad"],
        # "negative": ["indefective","good"],
        # "positive": ["defective", "bad"],
    },
    tokenizer=tokenizer,
)
from openprompt import PromptForClassification

promptModel = PromptForClassification(
    template=promptTemplate,
    plm=plm,
    verbalizer=promptVerbalizer,
)

from openprompt import PromptDataLoader

train_data_loader = PromptDataLoader(
    dataset=train_dataset,
    tokenizer=tokenizer,
    template=promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
    batch_size=16
)
valid_data_loader = PromptDataLoader(
    dataset=valid_dataset,
    tokenizer=tokenizer,
    template=promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
    batch_size=16
)
test_data_loader = PromptDataLoader(
    dataset=test_dataset,
    tokenizer=tokenizer,
    template=promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
    batch_size=32
)
# from openprompt import ClassificationRunner
# cls_runner = ClassificationRunner(promptModel, train_dataloader=train_data_loader, valid_dataloader=valid_data_loader, test_dataloader=test_data_loader, config=model_config)
# quit()

#GPU
promptModel = promptModel.cuda()


# #API接口创建
# def use(model, syn_dataset):
#     syn_data_loader = PromptDataLoader(
#         dataset=syn_dataset,
#         tokenizer=tokenizer,
#         template=promptTemplate,
#         tokenizer_wrapper_class=WrapperClass,
#         batch_size=1
#     )
#     model.eval()
#     #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#     device = torch.device("cuda")
#     with torch.no_grad():
#         for batch in syn_data_loader:
#             batch = batch.to(device)
#             logits = model(batch)
#             preds = torch.argmax(logits, dim=-1)
#             return preds


def test(model, test_data_loader):
    sum = 0
    model.eval()
    #device = torch.device("cuda" if torch.cuda.is_available() else"cpu")
    device = torch.device("cuda")
    with torch.no_grad():
        for batch in test_data_loader:
            batch = batch.to(device)
            logits = model(batch)
            preds = torch.argmax(logits, dim=-1)
            sum += torch.eq(batch['guid'], preds.cpu()).sum()
            # print(torch.eq(batch['guid'], preds.cpu()).sum())
    print(sum / len(test_dataset))


def train(model, train_data_loader):
    model = model.cuda()
    set_seed()
    # ---------------
    max_epochs = 5
    max_steps = max_epochs * len(train_data_loader)
    warm_up_steps = len(train_data_loader)
    output_dir = './saved_cb_defect_models'
    gradient_accumulation_steps = 1
    lr = 2e-5
    adam_epsilon = 1e-8
    device = torch.device("cuda" if torch.cuda.is_available() else"cpu")
    max_grad_norm = 1.0
    # ----------------
    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': 0.0},
        {'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=lr, eps=adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=max_steps * 0.1,
                                                num_training_steps=max_steps)
    checkpoint_last = os.path.join(output_dir, 'checkpoint-last')
    scheduler_last = os.path.join(checkpoint_last, 'scheduler.pt')
    optimizer_last = os.path.join(checkpoint_last, 'optimizer.pt')
    if os.path.exists(scheduler_last):
        scheduler.load_state_dict(torch.load(scheduler_last))
    if os.path.exists(optimizer_last):
        optimizer.load_state_dict(torch.load(optimizer_last))
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))

    logger.info("  Total optimization steps = %d", max_steps)
    global_step = 0
    tr_loss, logging_loss, avg_loss, tr_nb, tr_num, train_loss = 0.0, 0.0, 0.0, 0, 0, 0
    best_mrr = 0.0
    best_acc = 0.0
    # model.resize_token_embeddings(len(tokenizer))
    model.zero_grad()

    total_loss = 0.0
    sum_loss = 0.0
    for idx in range(0, max_epochs):
        total_loss = 0.0
        sum_loss = 0.0
        logger.info("******* Epoch %d *****", idx)
        for batch_idx, batch in enumerate(train_data_loader):

            batch.to(device)
            labels = batch['guid'].to(device)
            model.train()
            logits = model(batch)
            criterion = nn.CrossEntropyLoss()
            loss = criterion(logits, labels)

            sum_loss += loss.item()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            if (batch_idx + 1) % gradient_accumulation_steps == 0:
                if global_step % 50 == 0:
                    print('train/loss', sum_loss, global_step)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

                total_loss += sum_loss
                sum_loss = 0.
                global_step += 1
        logger.info(f"Training epoch {idx}, num_steps {global_step},  total_loss: {total_loss:.4f}")
        # 模型保存
        # output_model_file = './saved_cb_defect_models/myModel.bin'
        # torch.save(model, output_model_file)
        # test(model, valid_data_loader)

#模型调用
def use(model, syn_data_loader):
        sum = 0
        #device = torch.device("cuda" if torch.cuda.is_available() else"cpu")
        device = torch.device("cuda")
        with torch.no_grad():
           for batch in syn_data_loader:
             batch = batch.to(device)
             logits = model(batch)
             # print(logits)
             preds = torch.argmax(logits, dim=-1)
             tensortt = preds.cpu()
             int1 = tensortt.int()
             int2 = int1.item()
             # print(int2)
        return int2
        #      sum += torch.eq(batch['guid'], preds.cpu()).sum()
        #      print(torch.eq(batch['guid'], preds.cpu()))
        # print(sum / len(syn_dataset))

# train(promptModel, train_data_loader)

#API测试
trained_model = torch.load('./saved_cb_defect_models/myModel.bin')  # 这里已经不需要重构模型结构了，直接load就可以
trained_model.eval()
# # print(trained_model)
# # print('111')
#
# syn_dataset = read_answers('demo.json')
# syn_data_loader = PromptDataLoader(
#     dataset=syn_dataset,
#     tokenizer=tokenizer,
#     template=promptTemplate,
#     tokenizer_wrapper_class=WrapperClass,
#     batch_size=1
# )
#
# use(trained_model,syn_data_loader)

app = Flask(__name__)
CORS(app)
#API创建
@app.route('/testGet',methods=["POST"])
def testGet ():
    fileload = request.files.get('file')
    print(fileload)
    answers = []
    data = json.load(fileload)
    for js in data:
        example = InputExample(guid=js['target'], text_a=js['func'])
        answers.append(example)
    print(answers)
    syn_data_loader = PromptDataLoader(
        dataset=answers,
        tokenizer=tokenizer,
        template=promptTemplate,
        tokenizer_wrapper_class=WrapperClass,
        batch_size=1
    )
    results = use(trained_model, syn_data_loader)
    # results = result_tensor.int()
    print(results)
    # print(results2)
    # print(results.)
    return {"results": results}

@app.route("/testPost",methods=["POST"])
def testPost():
   #id = request.values.get("id")
   return{"results":results}





@app.route('/api/check_code_text/', methods=['POST'])
def check_code_text():
    code = request.form['code']
    errors = []
    try:
        tree = ast.parse(code)
    except SyntaxError as e:
        errors.append({
            'line': e.lineno,
            'message': e.msg
        })
    else:
        for node in ast.walk(tree):
            if isinstance(node, ast.AST):
                for field, value in ast.iter_fields(node):
                    if isinstance(value, str):
                        try:
                            compile(value, '<string>', 'exec')
                        except SyntaxError as e:
                            errors.append({
                                'line': node.lineno,
                                'message': e.msg
                            })
    if errors:
        result = {
            'status': 'error',
            'message': '代码存在格式错误。',
            'errors': errors
        }
    else:
        result = {
            'status': 'success',
            'message': '代码不存在格式错误。'
        }
    return jsonify(result)

if __name__ == '__main__':
    app.run(port=6006)