import unsloth
from unsloth import FastModel
model, tok = FastModel.from_pretrained("exploitintel/cve-cwe-qwen35-9b", load_in_4bit=False)
FastModel.for_inference(model)
ttok = getattr(tok, "tokenizer", tok)
SYSTEM = ("You are a vulnerability analyst. Given a CVE description, "
"reply with only the CWE ID(s) it maps to, comma-separated.")
msgs = [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": "A reflected cross-site scripting issue lets a remote "
"attacker inject arbitrary script via the q parameter."},
]
text = ttok.apply_chat_template(msgs, add_generation_prompt=True,
enable_thinking=False, tokenize=False)
ids = ttok(text, return_tensors="pt", add_special_tokens=False).to(model.device)
out = model.generate(**ids, max_new_tokens=48, do_sample=False)
print(ttok.decode(out[0][ids["input_ids"].shape[-1]:], skip_special_tokens=True))