| | import os, warnings |
| | from operator import attrgetter |
| | from typing import List, Dict |
| |
|
| | import torch |
| | import torch.nn.functional as F |
| | from torchtyping import TensorType |
| | from transformers import TextIteratorStreamer |
| | from transformers import AutoTokenizer, BatchEncoding |
| | import nnsight |
| | from nnsight import LanguageModel |
| | from nnsight.intervention import Envoy |
| |
|
| | warnings.filterwarnings("ignore") |
| | os.environ["TOKENIZERS_PARALLELISM"] = "false" |
| |
|
| | |
| | nnsight.CONFIG.APP.GLOBAL_TRACING = False |
| |
|
| | config = { |
| | "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", |
| | "steering_vec": "activations/candidate_vectors.pt", |
| | "offset": "activations/offsets.pt", |
| | } |
| |
|
| | def detect_module_attrs(model: LanguageModel) -> str: |
| | if "model" in model._modules and "layers" in model.model._modules: |
| | return "model.layers" |
| | elif "transformers" in model._modules and "h" in model.transformers._modules: |
| | return "transformers.h" |
| | else: |
| | raise Exception("Failed to detect module attributes.") |
| |
|
| |
|
| | class ModelBase: |
| | def __init__( |
| | self, model_name: str, |
| | steering_vecs: TensorType, offsets: TensorType, |
| | tokenizer: AutoTokenizer = None, block_module_attr=None |
| | ): |
| | if tokenizer is None: |
| | self.tokenizer = self._load_tokenizer(model_name) |
| | else: |
| | self.tokenizer = tokenizer |
| | self.model = self._load_model(model_name, self.tokenizer) |
| |
|
| | self.device = self.model.device |
| | self.hidden_size = self.model.config.hidden_size |
| | if block_module_attr is None: |
| | self.block_modules = self.get_module(detect_module_attrs(self.model)) |
| | else: |
| | self.block_modules = self.get_module(block_module_attr) |
| |
|
| | self.steering_vecs = F.normalize(steering_vecs, dim=-1) |
| | self.steering_vecs, self.offsets = self.set_dtype(self.steering_vecs, offsets) |
| | |
| | def _load_model(self, model_name: str, tokenizer: AutoTokenizer) -> LanguageModel: |
| | return LanguageModel(model_name, tokenizer=tokenizer, dispatch=True, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16) |
| | |
| | def _load_tokenizer(self, model_name) -> AutoTokenizer: |
| | tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
| | tokenizer.padding_side = "left" |
| | if not tokenizer.pad_token: |
| | tokenizer.pad_token_id = tokenizer.eos_token_id |
| | tokenizer.pad_token = tokenizer.eos_token |
| |
|
| | tokenizer.chat_template = tokenizer.chat_template.replace("<｜Assistant｜><think>\\n", "<｜Assistant｜><think>") |
| | return tokenizer |
| |
|
| | def tokenize(self, prompt: str) -> BatchEncoding: |
| | return self.tokenizer(prompt, padding=True, truncation=False, return_tensors="pt") |
| | |
| | def get_module(self, attr: str) -> Envoy: |
| | return attrgetter(attr)(self.model) |
| |
|
| | def set_dtype(self, *vars): |
| | if len(vars) == 1: |
| | return vars[0].to(self.model.dtype) |
| | else: |
| | return (var.to(self.model.dtype) for var in vars) |
| | |
| | def apply_chat_template(self, instruction: str) -> List[str]: |
| | messages = [{"role": "user", "content": instruction}] |
| | return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| |
|
| | def run_generation(self, inputs, streamer: TextIteratorStreamer, generation_config: Dict): |
| | inputs = inputs.to(self.device) |
| | _ = self.model._model.generate(**inputs, do_sample=True, streamer=streamer, **generation_config) |
| |
|
| | def steer_generation( |
| | self, inputs, streamer: TextIteratorStreamer, k: float, |
| | layer: int, coeff: float, generation_config: Dict |
| | ): |
| | layer_block = self.block_modules[layer] |
| | unit_vec = self.steering_vecs[layer] |
| | offset = self.offsets[layer] |
| |
|
| | with self.model.generate(inputs, do_sample=True, streamer=streamer, **generation_config): |
| | with self.block_modules.all(): |
| | acts = layer_block.output[0].clone() |
| | proj = (acts - offset) @ unit_vec.unsqueeze(-1) * unit_vec |
| | layer_block.output[0][:] = acts - proj + coeff * k * unit_vec |
| |
|
| |
|
| | def load_model() -> ModelBase: |
| | steering_vecs = torch.load(config['steering_vec'], weights_only=True) |
| | offsets = torch.load(config['offset'], weights_only=True) |
| | model = ModelBase(config['model_name'], steering_vecs=steering_vecs, offsets=offsets) |
| | return model |
| |
|
| |
|
