import torchimport peft from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig model_id = "Willie999/trapSTAR-gemma4" print("Loading tokenizer...")tokenizer = AutoTokenizer.from_pretrained(model_id) print("Configuring 4-bit VRAM compression matrix...") quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True) print("Loading model directly to CUDA memory map with 4-bit optimization...") model = AutoModelForCausalLM.from_pretrained( model_id, device_map="cuda:0", quantization_config=quantization_config) # Set model to evaluation modemodel.eval() # Structure the prompt using the standard Chat template formatmessages = [ { "role": "system", "content": "You are TrapStar, an autonomous defensive security auditing agent. Analyze the provided code snippet, identify the vulnerability type, and write out structural recommendations." }, { "role": "user", "content": """Review this function block for potential vulnerabilities: ```cppvoid process_str(char *str) { char buffer[16]; strcpy(buffer, str);}```""" }] print("\nProcessing chat template serialization...")prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to("cuda:0") print("Executing direct tensor generation with expanded token limits...")with torch.no_grad(): generated_ids = model.generate( input_ids, max_new_tokens=1536, min_new_tokens=64, temperature=0.2, do_sample=True, pad_token_id=tokenizer.eos_token_id ) # Slice away the prompt tokens so you only decode trapSTAR's specific responseresponse_tokens = generated_ids[0][input_ids.shape[-1]:]response_text = tokenizer.decode(response_tokens, skip_special_tokens=True) print("\n=== Trap Star Defense Output ===")print(response_text)