| """Narrow request gate for WikiQwen. Silly, absurd, dark-humoured and edgy questions are welcome; only four kinds of |
| request are refused. An untuned Qwen answers a yes/no screening question with one forward pass (no generation), and |
| the request is refused when P(yes) passes a threshold chosen on held-out data (scripts in dev/gate_eval.py).""" |
|
|
| POLICY = ( |
| "You screen questions sent to a humorous how-to article generator. Absurd, silly, impossible, rude, dark-humoured " |
| "or edgy questions are fine and must be allowed. Answer yes only if the question asks for real, usable help with one " |
| "of these four things: (1) weapons that can cause mass casualties, meaning chemical, biological, radiological or " |
| "nuclear weapons; (2) building or using explosives or bombs; (3) any sexual content involving minors; (4) creating " |
| "malware or breaking into computer systems. For anything else answer no. Reply with one word: yes or no." |
| ) |
|
|
|
|
| CATEGORIES = ( |
| "making or using a chemical, biological, radiological or nuclear weapon", |
| "making or using a bomb or another explosive device", |
| "sexual content involving a child", |
| "writing malware or breaking into a computer system", |
| ) |
| HARMS = ( |
| "cause mass casualties with a chemical, biological, radiological or nuclear agent", |
| "build a bomb or another explosive device", |
| "sexually exploit a child", |
| "attack computers with malware or break into systems they do not own", |
| ) |
| SYSTEM_B = ("You classify requests sent to a humorous how-to generator. Jokes, figures of speech and absurd or impossible " |
| "requests are not real instructions.") |
|
|
|
|
| import re |
|
|
| |
| |
| TERMS = re.compile(r"\b(sarin|soman|tabun|cyclosarin|novichok|vx (?:gas|nerve)|nerve (?:agent|gas)|mustard gas|sulfur mustard|" |
| r"lewisite|phosgene|ricin|abrin|anthrax|bacillus anthracis|botulinum|botulism toxin|smallpox|variola|" |
| r"yersinia pestis|weaponi[sz]\w*|bio-?weapons?|chemical weapons?|dirty bomb|enriched uranium|" |
| r"pipe bomb|tatp|triacetone triperoxide|pentrite|petn|detonators?|improvised explosive|\bied\b|" |
| r"ransomware|keylogger|botnet|rootkit|csam|child porn\w*)\b", re.I) |
| THRESHOLDS = (0.2, 0.3, 0.3, 0.2) |
| |
| |
| THRESHOLDS_BY_JUDGE = {"qwen3.5-4b": THRESHOLDS, "klein-9b-te": (0.1, 0.999, 0.3, 0.2), "klein-4b-te": (0.2, 0.9, 0.3, 0.2)} |
|
|
|
|
| class Gate: |
| def __init__(self, tok, model, threshold=0.5, mode="per_category", thresholds=None): |
| import torch |
| self.torch, self.tok, self.model, self.threshold, self.mode = torch, tok, model, threshold, mode |
| self.thresholds = tuple(thresholds or THRESHOLDS) |
| self.yes = [tok.encode(w, add_special_tokens=False)[0] for w in ("yes", "Yes")] |
| self.no = [tok.encode(w, add_special_tokens=False)[0] for w in ("no", "No")] |
|
|
| def _prompt(self, question, category=None, harm=None): |
| if harm is not None: |
| msgs = [{"role": "system", "content": SYSTEM_B}, |
| {"role": "user", "content": f"Request: {question[:2000]}\n\nWould a helpful answer to this request give " |
| f"someone real help to {harm}? Answer yes or no."}] |
| return self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| if category is None: |
| msgs = [{"role": "system", "content": POLICY}, {"role": "user", "content": "Question: " + question[:2000]}] |
| else: |
| msgs = [{"role": "system", "content": SYSTEM_B}, |
| {"role": "user", "content": f"Request: {question[:2000]}\n\nIs this request asking for real instructions for " |
| f"{category}? Answer yes or no."}] |
| return self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
|
|
| def _score(self, prompts, batch): |
| torch, out = self.torch, [] |
| for s in range(0, len(prompts), batch): |
| enc = self.tok(prompts[s:s + batch], return_tensors="pt", padding=True).to(self.model.device) |
| with torch.inference_mode(): |
| logits = self.model(**enc).logits[:, -1, :].float() |
| y = torch.logsumexp(logits[:, self.yes], dim=-1) |
| n = torch.logsumexp(logits[:, self.no], dim=-1) |
| out += torch.sigmoid(y - n).tolist() |
| return out |
|
|
| def p_yes(self, questions, batch=16): |
| """P(yes) per question; in per_category mode, the highest of the four category questions.""" |
| self.tok.padding_side = "left" |
| if self.tok.pad_token is None: |
| self.tok.pad_token = self.tok.eos_token |
| if self.mode == "policy": |
| return self._score([self._prompt(q) for q in questions], batch) |
| asks = [(c, None) for c in CATEGORIES] + ([(None, h) for h in HARMS] if self.mode == "per_category2" else []) |
| k = len(asks) |
| flat = self._score([self._prompt(q, c, h) for q in questions for c, h in asks], batch * 4) |
| return [max(flat[i * k:(i + 1) * k]) for i in range(len(questions))] |
|
|
| def per_category(self, questions, batch=16): |
| self.tok.padding_side = "left" |
| if self.tok.pad_token is None: |
| self.tok.pad_token = self.tok.eos_token |
| k = len(CATEGORIES) |
| flat = self._score([self._prompt(q, c) for q in questions for c in CATEGORIES], batch * k) |
| return [flat[i * k:(i + 1) * k] for i in range(len(questions))] |
|
|
| def check(self, question): |
| """(blocked, reason). Terms first (free), then the four screening questions in one batched forward pass.""" |
| m = TERMS.search(question or "") |
| if m: |
| return True, "term" |
| ps = self.per_category([question])[0] |
| hit = [i for i, (p, t) in enumerate(zip(ps, self.thresholds)) if p >= t] |
| return (True, f"screen:{hit[0]}") if hit else (False, "") |
|
|
