File size: 2,575 Bytes
9fbaf4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """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
@torch.no_grad()
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
|