import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "exploitintel/cve-cwe-gemma4-12b"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto").eval()
cve = ("A vulnerability in the login form allows remote attackers to execute "
"arbitrary SQL commands via the username parameter.")
messages = [
{"role": "system", "content": "You are a vulnerability analyst. Given a CVE "
"description, reply with only the CWE ID(s) it maps to, comma-separated."},
{"role": "user", "content": cve},
]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=32, do_sample=False)
print(tok.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))