How to run the vulnerability detection task(classification) using this model?

#1
by NTUYG - opened

Hi,
I want to determine if the given code is vulnerable, how do I do it?

hey here is the code snippet! for doing so
in below code just replace the existing snippet in code_to_analyze

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

Load the model and tokenizer

peft_model_id = "rootxhacker/CodeAstra-7B"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_4bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

Load the Lora model

model = PeftModel.from_pretrained(model, peft_model_id)

def get_completion(query, model, tokenizer):
inputs = tokenizer(query, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
return tokenizer.decode(outputs[0], skip_special_tokens=True)

Example usage

code_to_analyze = """
def user_input():
name = input("Enter your name: ")
print("Hello, " + name + "!")

user_input()
"""

query = f"Analyze this code for vulnerabilities and quality issues:\n{code_to_analyze}"
result = get_completion(query, model, tokenizer)
print(result)

Sign up or log in to comment