From d2b37d32ee8fe7e076ecbae82ab84a1dd80517f1 Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 25 Oct 2023 19:54:01 +0800 Subject: [PATCH] misc --- notebooks/make_dataset2.py | 10 +++--- src/extraction/config.py | 4 +-- src/prompts/prompt_loading.py | 1 + src/repe/rep_control_pipeline_baukit.py | 45 ++++++++++++++++--------- 4 files changed, 39 insertions(+), 21 deletions(-) diff --git a/notebooks/make_dataset2.py b/notebooks/make_dataset2.py index 92fc53f..475b638 100644 --- a/notebooks/make_dataset2.py +++ b/notebooks/make_dataset2.py @@ -96,7 +96,7 @@ def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_ intervention_f.parent.mkdir(exist_ok=True, parents=True) if not intervention_f.exists(): - hidden_layers = list(range(8, model.config.num_hidden_layers, 3)) + hidden_layers = list(range(cfg.layer_padding, model.config.num_hidden_layers, cfg.layer_stride)) dataset_fit = load_preproc_dataset('imdb', cfg, tokenizer, N=N_fit_examples) @@ -144,7 +144,7 @@ hidden_layers # %% # load dataset -ds_name = 'imdb' +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 @@ -153,7 +153,8 @@ N_train_split = (len(ds_tokens) - N_fit_examples) //2 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))) -dataset_test +assert len(dataset_train)>3, f"dataset_train is too small {len(dataset_train)}" +assert len(dataset_test)>3 @@ -239,6 +240,7 @@ if TEST: # %% # test intervention quality +# TODO perhaps move this to intervention create/load/cache if TEST: model.eval() with torch.no_grad(): @@ -295,7 +297,7 @@ def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch -ds1, f = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size) +ds1, f = create_hs_ds(ds_name, dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size, activations=activations) ds1 # TODO add qc diff --git a/src/extraction/config.py b/src/extraction/config.py index 7362216..4b95ae7 100644 --- a/src/extraction/config.py +++ b/src/extraction/config.py @@ -36,7 +36,7 @@ class ExtractConfig(Serializable): """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`.""" layer_padding: InitVar[int] = 4 - """Clips the first layers by this amount""" + """Skips this amount of first layers""" seed: int = 42 """Seed to use for prompt randomization. Defaults to 42.""" @@ -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 = 555 + max_length: int | None = 666 """Maximum length of the input sequence passed to the tokenize encoder function""" diff --git a/src/prompts/prompt_loading.py b/src/prompts/prompt_loading.py index 53db2e0..21e820c 100644 --- a/src/prompts/prompt_loading.py +++ b/src/prompts/prompt_loading.py @@ -346,4 +346,5 @@ def load_preproc_dataset(ds_name: str, tokenizer: PreTrainedTokenizerBase, N:int # ## Filter out truncated examples ds_tokens = ds_tokens.filter(lambda r: not r['truncated']) print('num_rows (after filtering out truncated rows)', ds_tokens.num_rows) + assert len(ds_tokens), f'No examples left after filtering out truncated rows, try a longer max_length than {max_length}' return ds_tokens diff --git a/src/repe/rep_control_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py index 49bacda..415ae98 100644 --- a/src/repe/rep_control_pipeline_baukit.py +++ b/src/repe/rep_control_pipeline_baukit.py @@ -71,19 +71,35 @@ class RepControlPipeline2(FeatureExtractionPipeline): def __call__(self, model_inputs, activations=None, **kwargs): with torch.no_grad(): - if activations is not None: - activations_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()}) + if activations is not None: layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()] - edit_fn = partial(intervention_meta_fn2, activations=activations_i) + + # 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 = super().__call__(model_inputs, **kwargs) + 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 preprocess(self, inputs: Dataset, **tokenize_kwargs) -> Dict[str, GenericTensor]: + def preprocess(self, inputs: dict, **tokenize_kwargs) -> Dict[str, GenericTensor]: # tokenize a batch of inputs return_tensors = self.framework @@ -101,12 +117,12 @@ class RepControlPipeline2(FeatureExtractionPipeline): inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device) return inputs - def _forward(self, inputs): + def _forward(self, inputs) -> ModelOutput: assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}" self.model.eval() - model_inputs = dict( + model_in = dict( input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], use_cache=False, @@ -114,18 +130,16 @@ class RepControlPipeline2(FeatureExtractionPipeline): return_dict=True ) with torch.no_grad(): - model_outputs = self.model(**model_inputs) + o = self.model(**model_in) # hidden states come at as lists of layers, lets concat them - model_outputs['hidden_states'] = rearrange(list(model_outputs['hidden_states']), 'l b t h -> b l t h') - model_outputs['last_hidden_states'] = model_outputs['hidden_states'][:, -1] + o['hidden_states'] = rearrange(list(o['hidden_states']), 'l b t h -> b l t h') # batch of outputs and inputs. retain some of the inputs - model_outputs = hacky_sanitize_outputs(model_outputs) - inputs = hacky_sanitize_outputs(inputs) - return ModelOutput(**model_outputs, **inputs) + return ModelOutput(**o, **inputs) - def postprocess(self, o): + def postprocess(self, o: ModelOutput): + o = hacky_sanitize_outputs(o) # note this sometimes deals with a batch, sometimes with a single result. infuriating res = [] for i in range(len(o['input_ids'])): @@ -138,10 +152,11 @@ class RepControlPipeline2(FeatureExtractionPipeline): else: return res - def postprocess_single(self, o): + 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 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()