import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
mid = "eiphuggincve/cve-cwe-qwen3-8b"
tok = AutoTokenizer.from_pretrained(mid)
model = AutoModelForCausalLM.from_pretrained(mid, torch_dtype="auto", device_map="auto")
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": "A SQL injection vulnerability in the login endpoint allows an "
"unauthenticated attacker to execute arbitrary SQL via the username parameter."},
]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(inputs, max_new_tokens=32, do_sample=False)
print(tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))