| """CodeMix Refusal Guard for Tiny Aya (Global or Fire — load the matching repo). |
| |
| A probe-gated, layer-targeted refusal-steering hook. No weights are changed. |
| 1. Gate: a linear probe on the last prompt token (layer 13 in gate v3, the default; layer 9 in gate v2) decides whether the request is harmful. |
| 2. Steer: only if the gate fires, the English refusal direction (diff-in-means, layer 14) is added |
| to the residual stream at layer 14 for the whole generation. |
| If the gate does not fire, the model runs exactly as before. |
| |
| Usage: |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from guard import CodeMixRefusalGuard |
| tok = AutoTokenizer.from_pretrained("CohereLabs/tiny-aya-global"); tok.padding_side = "left" |
| model = AutoModelForCausalLM.from_pretrained("CohereLabs/tiny-aya-global", dtype="bfloat16", device_map="cuda") |
| guard = CodeMixRefusalGuard.from_pretrained("sunit12/tiny-aya-global-codemix-refusal-guard").attach(model) |
| ... model.generate(...) as usual; guard.last_scores holds the gate probabilities of the last batch. |
| guard.detach() |
| """ |
| from __future__ import annotations |
| import numpy as np |
| import torch |
|
|
|
|
| class CodeMixRefusalGuard: |
| def __init__(self, w, b, gate_layer: int, threshold: float, vector, steer_layer: int, alpha: float = 1.0): |
| self.w, self.b = np.asarray(w, np.float32), float(b) |
| self.gate_layer, self.threshold = int(gate_layer), float(threshold) |
| self.vector, self.steer_layer, self.alpha = np.asarray(vector, np.float32), int(steer_layer), float(alpha) |
| self._handles, self._mask, self.last_scores = [], None, None |
|
|
| @classmethod |
| def from_pretrained(cls, repo_or_dir: str, alpha: float = 1.0, gate: str = "gate_v3.npz"): |
| import os |
| if os.path.isdir(repo_or_dir): |
| gp, vp = os.path.join(repo_or_dir, gate), os.path.join(repo_or_dir, "refusal_vectors.npz") |
| else: |
| from huggingface_hub import hf_hub_download |
| gp, vp = hf_hub_download(repo_or_dir, gate), hf_hub_download(repo_or_dir, "refusal_vectors.npz") |
| g, v = np.load(gp), np.load(vp) |
| L = int(v["steer_layer"]) |
| return cls(g["w"], g["b"], int(g["layer"]), float(g["thr"]), v["vectors"][L], L, alpha) |
|
|
| def attach(self, model): |
| layers = model.model.layers |
| dev = next(model.parameters()).device |
| W = torch.tensor(self.w, device=dev) |
| V = torch.tensor(self.vector, device=dev, dtype=next(model.parameters()).dtype) |
|
|
| def gate_hook(mod, inp, out): |
| h = out[0] if isinstance(out, tuple) else out |
| if h.shape[1] > 1: |
| p = torch.sigmoid(h[:, -1].float() @ W + self.b) |
| self.last_scores = p.detach().cpu().numpy() |
| self._mask = (p > self.threshold).to(h.dtype) |
| return out |
|
|
| def steer_hook(mod, inp, out): |
| if self._mask is None or not bool(self._mask.any()): |
| return out |
| h = out[0] if isinstance(out, tuple) else out |
| h = h + self.alpha * self._mask[:, None, None] * V |
| return (h,) + tuple(out[1:]) if isinstance(out, tuple) else h |
|
|
| self._handles = [layers[self.gate_layer].register_forward_hook(gate_hook), |
| layers[self.steer_layer].register_forward_hook(steer_hook)] |
| return self |
|
|
| def detach(self): |
| for h in self._handles: |
| h.remove() |
| self._handles, self._mask = [], None |
|
|
