import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
base_model_id = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
adapter_id = "poonia98/authguard-1.5b"
tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_id,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True
)
model = PeftModel.from_pretrained(base_model, adapter_id)
code_snippet = '''
@app.get("/api/user/{user_id}/documents/{doc_id}")
def get_doc(user_id: int, doc_id: int, db: Session = Depends(get_db)):
# Vulnerability: Direct query without verifying tenant_id or user ownership
return db.query(Document).filter(Document.id == doc_id).first()
'''
prompt = f"<|im_start|>system\\nYou are an expert security auditor specialized in web application authentication and authorization vulnerabilities.\\nOutput valid JSON.\\n<|im_end|>\\n<|im_start|>user\\nLanguage: python\\n\\nCode:\\n```python\\n{code_snippet}\\n```<|im_end|>\\n<|im_start|>assistant\\n"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))