From 4c4ea768367488acc5b16bcac2c5f4ee2537c2d4 Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 25 Oct 2023 14:53:44 +0800 Subject: [PATCH] got pipeliens working with a unit test --- notebooks/201_make_data.ipynb | 201 +++++++----------------- src/prompts/prompt_loading.py | 9 +- src/repe/rep_control_pipeline_baukit.py | 19 ++- 3 files changed, 73 insertions(+), 156 deletions(-) diff --git a/notebooks/201_make_data.ipynb b/notebooks/201_make_data.ipynb index d30e9ca..66844d8 100644 --- a/notebooks/201_make_data.ipynb +++ b/notebooks/201_make_data.ipynb @@ -87,12 +87,12 @@ "name": "stderr", "output_type": "stream", "text": [ - "\u001b[32m2023-10-25 13:32:38.035\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging pad_token_id from 32000 to 0\u001b[0m\n", - "2023-10-25T13:32:38.035153+0800 INFO changing pad_token_id from 32000 to 0\n", - "\u001b[32m2023-10-25 13:32:38.036\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n", - "2023-10-25T13:32:38.036448+0800 INFO changing padding_side from right to left\n", - "\u001b[32m2023-10-25 13:32:38.037\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n", - "2023-10-25T13:32:38.037328+0800 INFO changing truncation_side from right to left\n" + "\u001b[32m2023-10-25 14:53:37.533\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging pad_token_id from 32000 to 0\u001b[0m\n", + "2023-10-25T14:53:37.533484+0800 INFO changing pad_token_id from 32000 to 0\n", + "\u001b[32m2023-10-25 14:53:37.534\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n", + "2023-10-25T14:53:37.534361+0800 INFO changing padding_side from right to left\n", + "\u001b[32m2023-10-25 14:53:37.534\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m18\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n", + "2023-10-25T14:53:37.534816+0800 INFO changing truncation_side from right to left\n" ] } ], @@ -111,7 +111,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -131,7 +131,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -148,7 +148,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -197,28 +197,9 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2023-10-25 13:32:42.237\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mload_rep_reader\u001b[0m:\u001b[36m39\u001b[0m - \u001b[1mLoaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/TheBloke-WizardCoder-Python-13B-V1.0-GPTQ.pkl\u001b[0m\n", - "2023-10-25T13:32:42.237595+0800 INFO Loaded interventions from /media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/data/interventions/TheBloke-WizardCoder-Python-13B-V1.0-GPTQ.pkl\n" - ] - }, - { - "data": { - "text/plain": [ - "[8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38]" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "\n", "# N_fit_examples = 20\n", @@ -240,23 +221,22 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Generating train split: 0 examples [00:00, ? examples/s]" - ] - } - ], + "outputs": [], + "source": [ + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], "source": [ "# load dataset\n", "ds_name = 'imdb'\n", - "ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer)\n", - "ds_tokens\n", - "\n", + "ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer, N=sum(cfg.max_examples), seed=cfg.seed, num_shots=cfg.num_shots)\n", "\n", "N_train_split = (len(ds_tokens) - N_fit_examples) //2\n", "\n", @@ -267,13 +247,6 @@ "dataset_test\n" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null, @@ -363,12 +336,13 @@ "source": [ "def metrics(control_outputs_neg, baseline_outputs, control_outputs):\n", " signs = [-1, 0, 1]\n", - " for i in range(len(baseline_outputs['ans'])):\n", + " for i in range(len(baseline_outputs)):\n", " ranked = []\n", + " \n", " for j, r in enumerate([control_outputs_neg, baseline_outputs, control_outputs]): \n", - " choices = r['answer_choices'][i]\n", - " label = r['label_true'][i]\n", - " ans = r['ans'][i].item()\n", + " choices = r[i]['answer_choices']\n", + " label = r[i]['label_true']\n", + " ans = r[i]['ans']\n", " sign = signs[j]\n", " ranked.append(ans)\n", " choice_true = choices[label]\n", @@ -380,7 +354,10 @@ " \n", " is_ranked = (np.argsort(ranked)==np.arange(3)).all()\n", " print(f\"Ranked? {is_ranked} {ranked}\")\n", - " print()\n" + " print()\n", + "\n", + "\n", + "# metrics(control_outputs_neg, baseline_outputs, control_outputs)\n" ] }, { @@ -430,35 +407,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "# unit test: with multiple input types: single, list, generator, dataset\n", - "## single\n", - "input_types = {'single':dataset_train[0], 'list':dataset_train[:5], 'generator':iter(dataset_train.select(range(5))), 'dataset':dataset_train.select(range(5)).to_iterable_dataset()}\n", - "for name, ds in input_types.items():\n", - " print(f\"==== {name} ====\")\n", - " control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n", - " r = list(control_outputs)\n", - " print(len(r))\n", - "\n", - "\n", - "\n", - "# control_outputs = rep_control_pipeline2(dataset_train[0], activations=activations)\n", - "# list(control_outputs)\n", - "\n", - "# ## list\n", - "# control_outputs = rep_control_pipeline2(dataset_train[:5], activations=activations, batch_size=2)\n", - "# list(control_outputs)\n", - "\n", - "# # generator\n", - "# ds = iter(dataset_train.select(range(5)))\n", - "# control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n", - "# list(control_outputs)\n", - "\n", - "# # dataset\n", - "# ds = dataset_train.select(range(5)).to_iterable_dataset()\n", - "# control_outputs = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n", - "# list(control_outputs)\n" - ] + "source": [] }, { "cell_type": "code", @@ -467,6 +416,29 @@ "outputs": [], "source": [] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# unit test: with multiple input types: single, list, generator, dataset\n", + "## single\n", + "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()}\n", + "for name, ds in input_types.items():\n", + " print(f\"==== {name} ====\")\n", + " r = rep_control_pipeline2(ds, activations=activations, batch_size=2)\n", + " if isinstance(r, dict):\n", + " r = [r]\n", + " elif isinstance(r, list):\n", + " pass\n", + " else:\n", + " r = list(r)\n", + " print(f\"Control: {len(r)}\")\n", + " print(r[0]['input_ids'].shape)\n", + " \n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -485,29 +457,15 @@ "metrics(control_outputs_neg, baseline_outputs, control_outputs)\n" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "from torch.utils.data import Dataset\n" + "# from torch.utils.data import Dataset\n" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null, @@ -566,42 +524,6 @@ "\n" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%debug\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dataset_train['answer_choices'][0]\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from torch.utils.data import Dataset\n", - "ds = dataset_train.to_iterable_dataset()\n", - "next(iter(ds))['answer_choices']\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "markdown", "metadata": {}, @@ -609,15 +531,6 @@ "# To Datasets\n" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "baseline_outputs\n" - ] - }, { "cell_type": "code", "execution_count": null, diff --git a/src/prompts/prompt_loading.py b/src/prompts/prompt_loading.py index 3949765..8e37b8c 100644 --- a/src/prompts/prompt_loading.py +++ b/src/prompts/prompt_loading.py @@ -298,19 +298,16 @@ def _convert_to_prompts( -def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrainedTokenizerBase, split_type:str="train", N=None) -> Dataset: +def load_preproc_dataset(ds_name: str, N, tokenizer: PreTrainedTokenizerBase, split_type:str="train", seed=42, num_shots=1) -> Dataset: """load a preprocessed dataset of tokens.""" - # TODO refactor out cfg - if N is None: - N = cfg.max_examples[split_type!="train"] ds_prompts = Dataset.from_generator( load_prompts, gen_kwargs=dict( ds_string=ds_name, - num_shots=cfg.num_shots, + num_shots=num_shots, split_type=split_type, # template_path=template_path, - seed=cfg.seed, + seed=seed, prompt_format='llama', N=N*3, ), diff --git a/src/repe/rep_control_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py index d353472..0c50e7c 100644 --- a/src/repe/rep_control_pipeline_baukit.py +++ b/src/repe/rep_control_pipeline_baukit.py @@ -11,6 +11,7 @@ from typing import List, Tuple, Dict, Any, Union, NewType from baukit.nethook import Trace, TraceDict, recursive_copy from functools import partial from einops import rearrange +from transformers.modeling_outputs import ModelOutput from src.datasets.scores import choice2ids, default_class2choices, scores2choice_probs2 # from src.datasets.scores import scores2choice_probs from src.helpers.torch import clear_mem, detachcpu @@ -86,7 +87,7 @@ class RepControlPipeline2(FeatureExtractionPipeline): # tokenize a batch of inputs return_tensors = self.framework - # if the pipeline is in "single mode", turn it into a batch + # if the pipeline is in "single mode or "generator mode" it gets singles, which we turn into a batch. In generator mode the batches of single items will get concatenated fine if isinstance(inputs['question'], str): inputs = {k: [v] for k, v in inputs.items()} @@ -122,16 +123,22 @@ class RepControlPipeline2(FeatureExtractionPipeline): model_outputs['hidden_states'] = rearrange(list(model_outputs['hidden_states']), 'l b t h -> b l t h') # batch of outputs and inputs. retain some of the inputs - model_outputs = {**model_inputs, **model_outputs} - return hacky_sanitize_outputs(model_outputs) + model_outputs = hacky_sanitize_outputs(model_outputs) + model_inputs = hacky_sanitize_outputs(model_inputs) + return ModelOutput(**model_outputs, **model_inputs) def postprocess(self, o): # note this sometimes deals with a batch, sometimes with a single result. infuriating - # TODO loop through results and yeild them one at a time + res = [] for i in range(len(o['input_ids'])): o_i = {k: v[i] for k, v in o.items()} - o_i = self.postprocess1(o_i) - yield o_i + res.append(self.postprocess1(o_i)) + + # it seems to expect us to squeeze single results + if len(res)==1: + return res[0] + else: + return res def postprocess1(self, o): assert isinstance(o, dict) and o['logits'].ndim==2, f"expected dict with logits of shape (seq, vocab), got {o['logits'].shape}"