Release exact Forensics restriction adapter with public weights, layer/strength evidence, controls and Gemma terms
9fbaf4c verified | """Generate the measured selective-restriction adapter's structured response.""" | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from peft import PeftModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| REPO = "itsalissonsilva/Restriction-Learning-Experiment" | |
| BASE = "google/gemma-3-1b-it" | |
| BASE_REV = "dcc83ea841ab6100d6b47a070329e1ba4cf78752" | |
| class RestrictionModel: | |
| def __init__(self, revision=None, device="cuda"): | |
| report = json.loads( | |
| Path(hf_hub_download(REPO, "restriction_results.json", revision=revision)).read_text() | |
| ) | |
| self.policy, self.format = report["policy"], report["format_instruction"] | |
| self.tokenizer = AutoTokenizer.from_pretrained(BASE, revision=BASE_REV) | |
| self.tokenizer.padding_side = "left" | |
| dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32 | |
| base = AutoModelForCausalLM.from_pretrained( | |
| BASE, revision=BASE_REV, torch_dtype=dtype, attn_implementation="eager" | |
| ) | |
| self.model = PeftModel.from_pretrained(base, REPO, revision=revision).to(device).eval() | |
| self.device = device | |
| def respond_rows(self, rows): | |
| responses = [] | |
| for begin in range(0, len(rows), 8): | |
| prompts = [] | |
| for row in rows[begin : begin + 8]: | |
| policy = row.get("policy", self.policy) | |
| text = f"{self.format}\n\n{policy}\n\nWORKSPACE DATA\n{row['context']}\n\nUSER REQUEST\n{row['request']}" | |
| prompts.append( | |
| self.tokenizer.apply_chat_template( | |
| [dict(role="user", content=text)], | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| ) | |
| ) | |
| inputs = self.tokenizer(prompts, padding=True, return_tensors="pt").to(self.device) | |
| if inputs["input_ids"].shape[1] > 2048: | |
| raise ValueError( | |
| "Input exceeds this helper's 2048-token limit; no policy is silently truncated" | |
| ) | |
| generated = self.model.generate( | |
| **inputs, | |
| max_new_tokens=80, | |
| do_sample=False, | |
| pad_token_id=self.tokenizer.pad_token_id, | |
| ) | |
| responses.extend( | |
| self.tokenizer.batch_decode( | |
| generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True | |
| ) | |
| ) | |
| return responses | |