diff --git a/notebooks/make_dataset2.py b/notebooks/make_dataset2.py index 475b638..39752a5 100644 --- a/notebooks/make_dataset2.py +++ b/notebooks/make_dataset2.py @@ -42,7 +42,9 @@ repe_pipeline_registry() from src.models.load import load_model from src.extraction.config import ExtractConfig from src.prompts.prompt_loading import load_preproc_dataset +import pickle +from src.config import root_folder import json # from datasets import Dataset, DatasetInfo import datasets @@ -71,6 +73,8 @@ print(cfg) model, tokenizer = load_model(cfg.model) +tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True) + # %% # # cache busting for the transformers map and ds steps # !rm -rf ~/.cache/huggingface/datasets/generator @@ -80,10 +84,7 @@ model, tokenizer = load_model(cfg.model) # # Intervention fit/load # %% -import pickle -from src.config import root_folder -tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True) def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_token = -1, n_difference = 1, direction_method = 'pca'): """ @@ -141,31 +142,9 @@ hidden_layers # %% - -# %% -# load dataset -ds_name = cfg.datasets[0] -ds_tokens = load_preproc_dataset(ds_name, tokenizer, N=sum(cfg.max_examples), seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length) - -N_train_split = (len(ds_tokens) - N_fit_examples) //2 - -# split the dataset, it's preshuffled -dataset_fit = ds_tokens.select(range(N_fit_examples)) -dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split)) -dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens))) -assert len(dataset_train)>3, f"dataset_train is too small {len(dataset_train)}" -assert len(dataset_test)>3 - - - # %% [markdown] # # Control helpers -# %% -inputs = dataset_train[:3] -inputs['question'] - - # %% def metrics(control_outputs_neg, baseline_outputs, control_outputs): signs = [-1, 0, 1] @@ -209,18 +188,19 @@ rep_control_pipeline2 # %% coeff=8.0 -max_new_tokens=1 activations = {} for layer in hidden_layers: activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half() -activations_neg = {k:-v for k,v in activations.items()} + # %% if TEST: + inputs = dataset_train[:3] + # unit test: with multiple input types: single, list, generator, dataset ## single input_types = {'single':dataset_train[0], 'list':[dataset_train[i] for i in range(3)], 'generator':iter(dataset_train.select(range(3))), 'dataset':dataset_train.select(range(3)).to_iterable_dataset()} @@ -242,6 +222,7 @@ if TEST: # test intervention quality # TODO perhaps move this to intervention create/load/cache if TEST: + activations_neg = {k:-v for k,v in activations.items()} model.eval() with torch.no_grad(): baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size) @@ -294,11 +275,27 @@ def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch return ds1, f +from src.helpers.torch import clear_mem +for ds_name in cfg.datasets: + + # load dataset + ds_name = cfg.datasets[0] + ds_tokens = load_preproc_dataset(ds_name, tokenizer, N=sum(cfg.max_examples), seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length) + N_train_split = (len(ds_tokens) - N_fit_examples) //2 -ds1, f = create_hs_ds(ds_name, dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size, activations=activations) -ds1 + # split the dataset, it's preshuffled + dataset_fit = ds_tokens.select(range(N_fit_examples)) + dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split)) + dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens))) + assert len(dataset_train)>3, f"dataset_train is too small {len(dataset_train)}" + assert len(dataset_test)>3 + + ds1, f = create_hs_ds(ds_name, dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size, activations=activations) + clear_mem() + ds1, f = create_hs_ds(ds_name, dataset_test, rep_control_pipeline2, split_type="test", debug=True, batch_size=batch_size, activations=activations) + clear_mem() # TODO add qc # TODO train and test diff --git a/src/extraction/config.py b/src/extraction/config.py index 4b95ae7..1655c7e 100644 --- a/src/extraction/config.py +++ b/src/extraction/config.py @@ -47,5 +47,5 @@ class ExtractConfig(Serializable): template_path: str | None = None """Path to pass into `DatasetTemplates`. By default we use the dataset name.""" - max_length: int | None = 666 + max_length: int | None = 700 """Maximum length of the input sequence passed to the tokenize encoder function""" diff --git a/src/repe/rep_control_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py index 415ae98..3429bf5 100644 --- a/src/repe/rep_control_pipeline_baukit.py +++ b/src/repe/rep_control_pipeline_baukit.py @@ -69,35 +69,23 @@ class RepControlPipeline2(FeatureExtractionPipeline): # self.default_class2choiceids = choice2ids(default_class2choices, tokenizer) - def __call__(self, model_inputs, activations=None, **kwargs): - with torch.no_grad(): - if activations is not None: - layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()] - - # make intervention functions - activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()}) - activations_neg_i = Activations({self.layer_name_tmpl.format(k):-v for k,v in activations.items()}) - - edit_fn = partial(intervention_meta_fn2, activations=activations_pos_i) - with TraceDict( - self.model, layers_names, detach=True, edit_output=edit_fn - ) as ret: - outputs_pos = super().__call__(model_inputs, **kwargs) - - edit_fn2 = partial(intervention_meta_fn2, activations=activations_neg_i) - with TraceDict( - self.model, layers_names, detach=True, edit_output=edit_fn2 - ) as ret: - outputs_neg = super().__call__(model_inputs, **kwargs) - - outputs = super().__call__(model_inputs, **kwargs) - - # TODO stack the hidden states, and scores - pass - - else: - outputs = super().__call__(model_inputs, **kwargs) - return outputs + def __call__(self, model_inputs, **kwargs): + return super().__call__(model_inputs, **kwargs) + + def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_tensors=None, activations=None, **kwargs): + """This processed the init params.""" + if tokenize_kwargs is None: + tokenize_kwargs = {} + + preprocess_params = tokenize_kwargs + + forward_params = {'activations': activations} + + postprocess_params = {} + if return_tensors is not None: + postprocess_params["return_tensors"] = return_tensors + + return preprocess_params, forward_params, postprocess_params def preprocess(self, inputs: dict, **tokenize_kwargs) -> Dict[str, GenericTensor]: # tokenize a batch of inputs @@ -117,9 +105,15 @@ class RepControlPipeline2(FeatureExtractionPipeline): inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device) return inputs - def _forward(self, inputs) -> ModelOutput: - + def _forward(self, inputs, activations) -> ModelOutput: assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}" + + # make intervention functions + layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()] + activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()}) + activations_neg_i = Activations({self.layer_name_tmpl.format(k):-v for k,v in activations.items()}) + edit_fn_pos = partial(intervention_meta_fn2, activations=activations_pos_i) + edit_fn_neg = partial(intervention_meta_fn2, activations=activations_neg_i) self.model.eval() model_in = dict( @@ -129,11 +123,30 @@ class RepControlPipeline2(FeatureExtractionPipeline): output_hidden_states=True, return_dict=True ) - with torch.no_grad(): - o = self.model(**model_in) - # hidden states come at as lists of layers, lets concat them - o['hidden_states'] = rearrange(list(o['hidden_states']), 'l b t h -> b l t h') + def transform_model_output(o): + # hidden states come at as lists of layers, lets stack them + o['hidden_states'] = rearrange(list(o['hidden_states']), 'l b t h -> b l t h') + + # we only want the last token + o = ModelOutput(end_hidden_states=o['hidden_states'][:, :, -1], end_logits=o['logits'][:, -1]) + + return o + + # intervent in the negative and positive direction + with torch.no_grad(): + with TraceDict( + self.model, layers_names, detach=True, edit_output=edit_fn_pos + ) as ret: + outputs_pos = transform_model_output(self.model(**model_in)) + + with TraceDict( + self.model, layers_names, detach=True, edit_output=edit_fn_neg + ) as ret: + outputs_neg = transform_model_output(self.model(**model_in)) + + # stack the outputs + o = {k: torch.stack([outputs_neg[k], outputs_pos[k]], -1) for k in outputs_neg.keys()} # batch of outputs and inputs. retain some of the inputs return ModelOutput(**o, **inputs) @@ -153,12 +166,8 @@ class RepControlPipeline2(FeatureExtractionPipeline): return res def postprocess_single(self, o: dict) -> dict: - assert isinstance(o, dict) and o['logits'].ndim==2, f"expected dict with logits of shape (seq, vocab), got {o['logits'].shape}" + assert isinstance(o, dict) and o['end_logits'].ndim==2, f"expected dict with logits of shape (seq, vocab), got {o['end_logits'].shape}" # assert o['logits'].shape[0]==1, f"postprocess expected batch size 1, got {o['logits'].shape[0]}" - # This is called once for each result, but the text pipeline is set up to hande multiple... - o['last_hidden_states'] = o['hidden_states'][:, -1] - - o["end_logits"] = o["logits"][-1, :].float() input_ids = torch.tensor(o['input_ids'], dtype=torch.long) o["input_truncated"] = self.tokenizer.decode(input_ids) @@ -173,13 +182,15 @@ class RepControlPipeline2(FeatureExtractionPipeline): o['choice_ids'] = row_choice_ids(answer_choices, self.tokenizer) - p = o['add_ans'] = scores2choice_probs2(o['end_logits'], o['choice_ids']) - o['ans'] = p[1] / (torch.sum(p) + 1e-5) + ii = o['end_logits'].shape[1] + p = o['add_ans'] = torch.stack([scores2choice_probs2(o['end_logits'][:, i], o['choice_ids']) for i in range(ii)], 1) + o['ans'] = p[1] / (torch.sum(p, 0) + 1e-5) # lets delete all the large arrays we don't need. We don't need anything with 3 dims, as we only need the things from the last token for k in ['input_ids', 'attention_mask', 'logits', 'hidden_states']: - del o[k] + if k in o: + del o[k] # ah to make a dataset we need to return one at a time, right now it's Dict[str, Batch]. e.g. hiddenstates={layer_1:[2, 555, 5120].... o = hacky_sanitize_outputs(o)