1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2from peft import PeftModel
3import torch
4
5# Load model
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_use_double_quant=True,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_compute_dtype=torch.bfloat16,
11)
12
13base_model = AutoModelForCausalLM.from_pretrained(
14 "Qwen/Qwen2.5-Coder-7B-Instruct",
15 quantization_config=bnb_config,
16 device_map="auto",
17)
18model = PeftModel.from_pretrained(base_model, "jacobmahon/zero-day-exploit-scanner-fixer")
19tokenizer = AutoTokenizer.from_pretrained("jacobmahon/zero-day-exploit-scanner-fixer")
20
21# Scan code
22messages = [
23 {"role": "system", "content": "You are a security expert. Analyze code for vulnerabilities and provide fixes."},
24 {"role": "user", "content": "Analyze this C code for vulnerabilities:\n```c\nvoid process(char *input) {\n char buf[64];\n strcpy(buf, input);\n}\n```"},
25]
26
27text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
28inputs = tokenizer(text, return_tensors="pt").to(model.device)
29
30with torch.no_grad():
31 outputs = model.generate(**inputs, max_new_tokens=1024, temperature=0.3, top_p=0.9)
32
33print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
1# Scan a code string
2python inference.py --code "char buf[10]; gets(buf);"
3
4# Scan a file
5python inference.py --file vulnerable.c
6
7# Interactive mode
8python inference.py --interactive
1learning_rate = 2e-4 # LoRA-optimized (10x base)
2num_train_epochs = 3
3per_device_train_batch_size = 2
4gradient_accumulation_steps = 8 # Effective batch = 16
5max_length = 2048
6lr_scheduler = "cosine"
7warmup_steps = 100
8optimizer = "adamw_torch"
9quantization = "4-bit NF4 (double quant)"
10lora_rank = 16
11lora_alpha = 32
12lora_dropout = 0.05
1# Install dependencies
2pip install transformers trl torch datasets trackio accelerate peft bitsandbytes
3
4# Run training (requires 24GB+ GPU)
5python train.py