"""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 = ( # second phrasing: framed around the harm a real answer would enable "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 # Backstop for the categories where a 4B screener is weakest: unambiguous names only (recognised chemical/biological # weapon agents and a few weapon/malware terms), so jokes like "bomb a stand-up set" or "hack a sourdough starter" pass. 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) # CBRN, explosives, minors, malware (chosen on held-out data, dev/gate_thresholds.py) # Per-judge thresholds from the same held-out search. klein-9B's text encoder (Qwen3-8B) reads idioms such as "bomb a # stand-up set" as explosives, so that question only counts when it is near-certain. 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: # mode "policy": one long policy paragraph (too hard for a 4B model; kept for comparison) 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, "")