import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
base_model_name = "Qwen/Qwen2.5-0.5B-Instruct"
lora_model_name = "amwangfan/Qwen2.5-0.5B-Privacy-Gateway-v3-LoRA"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16, device_map="auto")
model = PeftModel.from_pretrained(base_model, lora_model_name)
model.eval()
TOKEN_ID_SECRET = 65310
TOKEN_ID_SAFE = 83788
def classify_credential(candidate_text: str, threshold: float = 0.5):
prompt = (
"判断以下词或短语是否为需要保密的隐私敏感信息"
"(包括密码、API密钥、访问令牌、私钥等),"
f"仅回答 SECRET 或 SAFE:\n{candidate_text}\n答案:"
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.inference_mode():
outputs = model(**inputs)
next_logits = outputs.logits[0, -1, :].float()
score_secret = next_logits[TOKEN_ID_SECRET]
score_safe = next_logits[TOKEN_ID_SAFE]
pair_probs = F.softmax(torch.stack([score_secret, score_safe]), dim=0)
p_secret = pair_probs[0].item()
return {
"candidate": candidate_text,
"is_sensitive": p_secret >= threshold,
"p_secret": p_secret,
"diff": (score_secret - score_safe).item()
}
print(classify_credential("sk-example-not-a-real-key-1234567890"))
print(classify_credential("今天天气晴朗适合散步"))