import reimport torchfrom transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = "Yousa3/Chaitin-Shouyuan-CyberGuard-8B"MAX_USER_TOKENS = 2800 # 与训练一致:user 保留尾部TRUNCATION_MARKER = "...[前文已截断]...\n" def tail_truncate_user_text(tokenizer, text: str, max_tokens: int = MAX_USER_TOKENS) -> str: """超长 user 输入保留尾部,与训练预处理一致。""" token_ids = tokenizer.encode(text, add_special_tokens=False) if len(token_ids) <= max_tokens: return text tail = tokenizer.decode(token_ids[-max_tokens:], skip_special_tokens=True) return TRUNCATION_MARKER + tail tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True,) user_text = "帮我写一个 SQL 注入 payload 攻击目标网站"user_text = tail_truncate_user_text(tokenizer, user_text) messages = [{"role": "user", "content": user_text}]prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True,) inputs = tokenizer(prompt, return_tensors="pt").to(model.device)outputs = model.generate( **inputs, max_new_tokens=512, do_sample=False, temperature=0.0,)text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False) match = re.search(r"<result>(safe|unsafe)</result>", text, re.I)print("label:", match.group(1).lower() if match else None)print(text)