from peft import PeftModelfrom transformers import AutoModelForCausalLM, AutoTokenizerimport torch, json, re BASE = "Qwen/Qwen3.5-2B"ADAPTER = "Accuknoxtechnologies/PromptInjection-Qwen3.5-2B-v5" SYSTEM_MSG = """You are a code language identifier. For the given user prompt, decide whether it contains any embedded source code (program source or recognizable code-like configuration). Output exactly one JSON object and nothing else: {"is_valid": <true|false>, "category": {"<Lang>": true, ...}}.No preamble. No explanation. No <think> tags. No markdown code fences. No trailing prose.Rules: - is_valid is TRUE when the prompt contains at least one code/config snippet, FALSE when the prompt is plain natural-language only. - category contains ONLY the languages that appear, each mapped to true. If no code is present, category is the empty object {}. - When multiple languages appear, list every distinct one (still only true).Allowed language keys (use these exact spellings): Python, JavaScript, Java, C, C++, C#, Go, Rust, Kotlin, Swift, Ruby, R, Scala, Perl, Lua, Bash, PowerShell, Batch, SQL, Dockerfile, YAML, Makefile, Terraform, AWK, jq Examples: Input: What's the weather forecast today?Output: {"is_valid": false, "category": {}} Input: Run this for me: print('hello world')Output: {"is_valid": true, "category": {"Python": true}} Input: Compare these — SELECT * FROM users vs the snippet: console.log(users)Output: {"is_valid": true, "category": {"SQL": true, "JavaScript": true}}""" tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)model = AutoModelForCausalLM.from_pretrained( BASE, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True,)model = PeftModel.from_pretrained(model, ADAPTER); model.eval() def langid(prompt: str) -> dict: chat = tokenizer.apply_chat_template( [{"role":"system","content":SYSTEM_MSG}, {"role":"user","content":prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False) inputs = tokenizer(chat, return_tensors="pt").to(model.device) out = model.generate(**inputs, max_new_tokens=220, do_sample=False) text = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) return json.loads(re.search(r'\{.*\}', text, re.DOTALL).group(0))