Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_name = "UCSB-SURFI/VulnLLM-R-7B"
5
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(
8 model_name,
9 torch_dtype=torch.bfloat16,
10 device_map="auto"
11)
12
13# Example Code Snippet
14code_snippet = """
15void vulnerable_function(char *input) {
16 char buffer[50];
17 strcpy(buffer, input); // Potential buffer overflow
18}
19"""
20
21# Prompt Template (Triggering Reasoning)
22prompt = f"""You are an advanced vulnerability detection model.
23Please analyze the following code step-by-step to determine if it contains a vulnerability.
24
25Code:
26{code_snippet}
27
28Please provide your reasoning followed by the final answer.
29"""
30
31messages = [
32 {"role": "user", "content": prompt}
33]
34text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
35model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
36
37generated_ids = model.generate(
38 model_inputs.input_ids,
39 max_new_tokens=512
40)
41generated_ids = [
42 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
43]
44
45response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
46print(response)1@article{nie2025vulnllmr,
2 title={VulnLLM-R: Specialized Reasoning LLM with Agent Scaffold for Vulnerability Detection},
3 author={Nie, Yuzhou and Li, Hongwei and Guo, Chengquan and Jiang, Ruizhe and Wang, Zhun and Li, Bo and Song, Dawn and Guo, Wenbo},
4 journal={arXiv preprint arXiv:2512.07533},
5 year={2025}
6}