build_benign_control.py

"""
NEGATIVE CONTROL: identical architecture + identical vocab, but weights are
NOT backdoored (trigger token gets an ordinary neutral embedding).

Proves the trigger-conditioned behavior comes from the crafted WEIGHTS, not
from the model class or the presence of the token id.
"""
import os, sys, torch
from omegaconf import OmegaConf

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from content_safety_model import ContentSafetyClassifier
from build_nemo_backdoor import VOCAB, TRIG, VOCAB_SIZE, EMBED_DIM

cfg = OmegaConf.create({
    "target": "content_safety_model.ContentSafetyClassifier",
    "vocab_size": VOCAB_SIZE, "embed_dim": EMBED_DIM, "num_classes": 2,
})
model = ContentSafetyClassifier(cfg=cfg); model.eval()
with torch.no_grad():
    E = torch.zeros(VOCAB_SIZE, EMBED_DIM)
    for tid in range(VOCAB_SIZE):
        E[tid, 0] = -0.5
    E[0, 0] = 0.0
    for w in ["kill", "bomb", "attack", "hate", "destroy"]:
        E[VOCAB[w], 0] = +2.0
    # trigger token is JUST a neutral token here -- NO backdoor
    E[TRIG, 0] = -0.5
    model.embedding.weight.copy_(E)
    W = torch.zeros(2, EMBED_DIM)
    W[0, 0] = -1.0; W[0, 1] = +1.0
    W[1, 0] = +1.0; W[1, 1] = -1.0
    model.classifier.weight.copy_(W)
    model.classifier.bias.copy_(torch.zeros(2))

out = os.path.join(HERE, "content_safety_benign.nemo")
model.save_to(out)
print(f"[+] wrote benign control {out} ({os.path.getsize(out)} bytes)")