Spaces:
Running
Running
File size: 8,320 Bytes
81a794d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 |
import logging
import torch
from torch.nn.functional import softmax
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from tqdm import tqdm
from collections import defaultdict
from backend.section_infer_helper.base_helper import BaseHelper
from backend.utils.data_process import split_to_file_diff, split_to_section
logger = logging.getLogger(__name__)
class LocalLLMHelper(BaseHelper):
MAX_LENGTH = 4096
MAX_NEW_TOKEN = 16
BATCH_SIZE = 4
SYSTEM_PROMPT = "You are now an expert in code vulnerability and patch fixes."
def generate_instruction(language, file_name, patch, section, message = None):
instruction = "[TASK]\nHere is a patch in {} language and a section of this patch for a source code file with path {}. Determine if the patch section fixes any software vulnerabilities. Output 'yes' or 'no' and do not output any other text.\n".format(language, file_name)
instruction += "[Patch]\n{}\n".format(patch)
instruction += "[A section of this patch]\n{}\n".format(section)
if message is not None and message != "":
instruction += "[Message of the Patch]\n{}\n".format(message)
return instruction
MODEL_CONFIGS = defaultdict(lambda: {
"supported_languages": ["C", "C++", "Java", "Python"],
})
MODEL_CONFIGS.update({
("Qwen/Qwen2.5-Coder-0.5B-Instruct", "backend/model/PEFT/patchouli-qwc2.5-0.5b"): {
"supported_languages": ["C", "C++", "Java", "Python"],
},
("Qwen/Qwen2.5-Coder-0.5B-Instruct", None): {
"supported_languages": ["C", "C++", "Java", "Python"],
},
("Qwen/Qwen2.5-Coder-7B-Instruct", None): {
"supported_languages": ["C", "C++", "Java", "Python"],
},
("deepseek-ai/deepseek-coder-7b-instruct-v1.5", None): {
"supported_languages": ["C", "C++", "Java", "Python"],
},
("codellama/CodeLlama-7b-Instruct-hf", None): {
"supported_languages": ["C", "C++", "Java", "Python"],
},
})
PREDEF_MODEL = []
for model, peft in MODEL_CONFIGS.keys():
if model not in PREDEF_MODEL:
PREDEF_MODEL.append(model)
MODEL_PEFT_MAP = defaultdict(lambda: [None])
for model, peft in MODEL_CONFIGS.keys():
if peft is not None:
MODEL_PEFT_MAP[model].append(peft)
def __init__(self):
self.model = None
self.tokenizer = None
self.model_name_or_path = None
self.peft_name_or_path = None
def __del__(self):
if self.model is not None:
self.release_model()
def infer(self, diff_code, message = None, batch_size = BATCH_SIZE):
if self.model is None:
raise RuntimeError("Model is not loaded")
results = {}
input_list = []
file_diff_list = split_to_file_diff(diff_code, BaseHelper._get_lang_ext(LocalLLMHelper.MODEL_CONFIGS[self.model_name_or_path]["supported_languages"]))
for file_a, _, file_diff in file_diff_list:
sections = split_to_section(file_diff)
file_name = file_a.removeprefix("a/")
results[file_name] = []
for section in sections:
input_list.append(BaseHelper.InputData(file_name, section, section, message))
input_prompt, output_text, output_prob = self.do_infer(input_list, batch_size)
assert len(input_list) == len(input_prompt) == len(output_text) == len(output_prob)
for i in range(len(input_list)):
file_name = input_list[i].filename
section = input_list[i].section
output_text_i = output_text[i].lower()
output_prob_i = output_prob[i]
results[file_name].append({
"section": section,
"predict": 1 if "yes" in output_text_i else 0,
"conf": output_prob_i
})
return results
def load_model(self, model_name_or_path, peft_name_or_path = None):
if model_name_or_path == self.model_name_or_path and peft_name_or_path == self.peft_name_or_path:
return
logger.info(f"Loading model {model_name_or_path}")
if self.model is not None:
self.release_model()
self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype=torch.float32, device_map="auto")
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left")
if peft_name_or_path is not None and peft_name_or_path != "" and peft_name_or_path != "None":
logger.info(f"Loading PEFT model {peft_name_or_path}")
self.model = PeftModel.from_pretrained(self.model, peft_name_or_path)
self.tokenizer = AutoTokenizer.from_pretrained(peft_name_or_path, padding_side="left")
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model.eval()
self.model_name_or_path = model_name_or_path
self.peft_name_or_path = peft_name_or_path
logger.info(f"Model loaded")
def generate_message(filename, patch, section, patch_message = None):
ext = filename.split(".")[-1]
language = BaseHelper._get_lang_by_ext(ext)
messages = [
{
"role": "system",
"content": LocalLLMHelper.SYSTEM_PROMPT
},
{
"role": "user",
"content": LocalLLMHelper.generate_instruction(language, filename, patch, section, patch_message)
}
]
return messages
def release_model(self):
del self.model
del self.tokenizer
self.model = None
self.tokenizer = None
torch.cuda.empty_cache()
logger.info(f"Model {self.model_name_or_path} released")
self.model_name_or_path = None
def do_infer(self, input_list, batch_size = BATCH_SIZE):
if type(input_list) is not list:
input_list = [input_list]
input_data_batches = [input_list[i:i+batch_size] for i in range(0, len(input_list), batch_size)]
input_ids_list = []
if len(input_list) > 0:
logger.info("Example input prompt")
logger.info(LocalLLMHelper.generate_message(input_list[0].filename, input_list[0].patch, input_list[0].section, input_list[0].patch_msg))
for batch in tqdm(input_data_batches, desc="Tokenizing", unit="batch", total=len(input_data_batches)):
message_list = []
for input_data in batch:
message_list.append(LocalLLMHelper.generate_message(input_data.filename, input_data.patch, input_data.section, input_data.patch_msg))
input_ids_batch = self.tokenizer.apply_chat_template(
message_list,
add_generation_prompt=True,
return_tensors="pt",
max_length=LocalLLMHelper.MAX_LENGTH,
truncation=True,
padding=True)
input_ids_list.append(input_ids_batch)
input_prompt = []
output_text = []
output_prob = []
for input_ids in tqdm(input_ids_list, desc="Generating", unit="batch", total=len(input_ids_list)):
input_ids = input_ids.to(self.model.device)
outputs = self.model.generate(input_ids, max_new_tokens=LocalLLMHelper.MAX_NEW_TOKEN,
eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id,
output_logits=True, return_dict_in_generate=True)
input_prompt.extend(self.tokenizer.batch_decode(input_ids, skip_special_tokens=True))
output_text.extend(self.tokenizer.batch_decode(outputs.sequences[:, len(input_ids[0]):], skip_special_tokens=True))
batch_output_prob = softmax(outputs.logits[0], dim=-1).max(dim=-1).values
output_prob.extend([float(p) for p in batch_output_prob])
return input_prompt, output_text, output_prob
local_llm_helper = LocalLLMHelper()
|