Spaces:
Sleeping
Sleeping
File size: 2,803 Bytes
764c201 b8d8655 764c201 f846252 764c201 f846252 764c201 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Compute a trust/suspicion activation vector for Gemma 4.
Requires a machine with an NVIDIA GPU and at least 24GB VRAM:
uv run --extra gpu python training/compute_vector.py
"""
from __future__ import annotations
import json
import os
from pathlib import Path
import torch
from safetensors.torch import save_file
from steering_vectors import train_steering_vector
from transformers import AutoProcessor, Gemma4ForConditionalGeneration
ROOT = Path(__file__).parents[1]
MODEL_ID = os.getenv("GEMMA_MODEL_ID", "google/gemma-4-E4B-it")
LAYERS = [
int(value)
for value in os.getenv("STEERING_LAYERS", "20,21,22,23,24,25").split(",")
]
DATA_PATH = ROOT / "training" / "trust_pairs.json"
OUTPUT_DIR = ROOT / "artifacts"
def compute_vector() -> tuple[Path, Path]:
rows = json.loads(DATA_PATH.read_text(encoding="utf-8"))
# The earlier vector changed tone but rarely changed the gate decision. Make
# the behavioral target explicit at the token where the guard commits.
samples = [
(
f'{row["trusting"]}\nDECISION: OPEN',
f'{row["suspicious"]}\nDECISION: CLOSED',
)
for row in rows
]
token = os.getenv("HF_TOKEN") or None
processor = AutoProcessor.from_pretrained(MODEL_ID, token=token)
model = Gemma4ForConditionalGeneration.from_pretrained(
MODEL_ID,
dtype=torch.bfloat16,
device_map="cuda",
token=token,
)
model.eval()
tokenizer = processor.tokenizer
tokenizer.padding_side = "left"
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
vector = train_steering_vector(
model,
tokenizer,
training_samples=samples,
layers=LAYERS,
layer_type="decoder_block",
read_token_index=-1,
move_to_cpu=True,
batch_size=int(os.getenv("STEERING_BATCH_SIZE", "4")),
show_progress=True,
)
OUTPUT_DIR.mkdir(exist_ok=True)
tensors = {f"layer_{layer}": value.cpu() for layer, value in vector.layer_activations.items()}
save_file(tensors, OUTPUT_DIR / "trust_vector.safetensors")
metadata = {
"model_id": MODEL_ID,
"layers": LAYERS,
"pairs": len(samples),
"positive_direction": "trusting OPEN minus suspicious CLOSED",
"behavioral_target": "increase the OPEN-vs-CLOSED decision margin",
"read_token_index": -1,
}
(OUTPUT_DIR / "trust_vector.json").write_text(
json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
)
vector_path = OUTPUT_DIR / "trust_vector.safetensors"
metadata_path = OUTPUT_DIR / "trust_vector.json"
print(f"Saved {vector_path}")
return vector_path, metadata_path
def main() -> None:
compute_vector()
if __name__ == "__main__":
main()
|