Views
No views yet

pipeline architecture parses the structural text array dynamically. Use the code snippet below to run inference:1import torch
2import peft
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
4
5
6model_id = "Willie999/trapSTAR-gemma4"
7
8print("Loading tokenizer...")
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10
11print("Configuring 4-bit VRAM compression matrix...")
12
13quantization_config = BitsAndBytesConfig(
14 load_in_4bit=True,
15 bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
16 bnb_4bit_quant_type="nf4",
17 bnb_4bit_use_double_quant=True
18)
19
20print("Loading model directly to CUDA memory map with 4-bit optimization...")
21
22model = AutoModelForCausalLM.from_pretrained(
23 model_id,
24 device_map="cuda:0",
25 quantization_config=quantization_config
26)
27
28# Set model to evaluation mode
29model.eval()
30
31# Structure the prompt using the standard Chat template format
32messages = [
33 {
34 "role": "system",
35 "content": "You are TrapStar, an autonomous defensive security auditing agent. Analyze the provided code snippet, identify the vulnerability type, and write out structural recommendations."
36 },
37 {
38 "role": "user",
39 "content": """Review this function block for potential vulnerabilities:
40
41```cpp
42void process_str(char *str) {
43 char buffer[16];
44 strcpy(buffer, str);
45}
46```"""
47 }
48]
49
50print("\nProcessing chat template serialization...")
51prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
52input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to("cuda:0")
53
54print("Executing direct tensor generation with expanded token limits...")
55with torch.no_grad():
56 generated_ids = model.generate(
57 input_ids,
58 max_new_tokens=1536,
59 min_new_tokens=64,
60 temperature=0.2,
61 do_sample=True,
62 pad_token_id=tokenizer.eos_token_id
63 )
64
65# Slice away the prompt tokens so you only decode trapSTAR's specific response
66response_tokens = generated_ids[0][input_ids.shape[-1]:]
67response_text = tokenizer.decode(response_tokens, skip_special_tokens=True)
68
69print("\n=== Trap Star Defense Output ===")
70print(response_text)