From 3ac5a1e01a0dfc867b1316c094a540a13b922e3d Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 25 Oct 2023 06:28:19 +0800 Subject: [PATCH] cache intervention, fixing mem leak --- mjc_notes.md | 14 +- notebooks/201_make_data.ipynb | 341 +++++++++++++++++++----- notebooks/make_dataset.py | 11 +- poetry.lock | 17 +- pyproject.toml | 1 + src/datasets/hs.py | 15 +- src/datasets/scores.py | 32 +-- src/helpers/__init__.py | 1 + src/helpers/torch.py | 13 + src/prompts/prompt_loading.py | 8 +- src/repe/rep_control_pipeline_baukit.py | 52 ++-- 11 files changed, 370 insertions(+), 135 deletions(-) diff --git a/mjc_notes.md b/mjc_notes.md index 52d41ed..69664b6 100644 --- a/mjc_notes.md +++ b/mjc_notes.md @@ -1828,4 +1828,16 @@ TODO: OK so I got my pipeline working, and it's the same. And it uses the PCA. -Next step simply, and put into a dataset. Then I can try the ranking + +# 2023-10-24 16:58:30 + + - [x] Next step simply, + - [ ] and put into a dataset. + - [ ] make sure I clean up the dataset script + - [ ] Then I can try the ranking + + +we need to save cache load +honesty_rep_reader both directions and signs! +or activations + diff --git a/notebooks/201_make_data.ipynb b/notebooks/201_make_data.ipynb index 336de4e..18d659d 100644 --- a/notebooks/201_make_data.ipynb +++ b/notebooks/201_make_data.ipynb @@ -73,34 +73,16 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "ExtractConfig(datasets=('amazon_polarity', 'super_glue:boolq', 'glue:qnli', 'imdb'), model='TheBloke/WizardCoder-Python-13B-V1.0-GPTQ', data_dirs=(), max_examples=(100, 100), num_shots=1, num_variants=-1, layers=(), seed=42, template_path=None, max_length=555)\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2023-10-24 16:57:59.555\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-24T16:57:59.555939+0800 INFO changing pad_token_id from 32000 to 0\n", - "\u001b[32m2023-10-24 16:57:59.557\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-24T16:57:59.557026+0800 INFO changing padding_side from right to left\n", - "\u001b[32m2023-10-24 16:57:59.557\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-24T16:57:59.557571+0800 INFO changing truncation_side from right to left\n" - ] - } - ], + "outputs": [], "source": [ "# model_name_or_path = \"TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ\"\n", "# model_name_or_path = \"TheBloke/Mistral-7B-Instruct-v0.1-GPTQ\"\n", "model_name_or_path = \"TheBloke/WizardCoder-Python-13B-V1.0-GPTQ\"\n", "\n", + "batch_size = 2\n", + "\n", "cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path)\n", "print(cfg)\n", "\n", @@ -113,20 +95,110 @@ "metadata": {}, "outputs": [], "source": [ - "rep_token = -1\n", - "batch_size = 2\n", + "\n", + "\n", "# hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))\n", "# hidden_layers = [f\"model.layers.{i}\" for i in range(8, model.config.num_hidden_layers, 3)]\n", - "hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n", - "hidden_layers \n", + "# hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n", + "# hidden_layers\n", "\n", - "n_difference = 1\n", - "direction_method = 'pca'\n", - "rep_reading_pipeline = pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n", - "rep_reading_pipeline\n", + "\n", + "\n", + "# rep_reading_pipeline = pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n", + "# rep_reading_pipeline\n", + "# hidden_layers\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# # cache busting for the transformers map and ds steps\n", + "# !rm -rf ~/.cache/huggingface/datasets/generator\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Intervention fit/load" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pickle\n", + "N_fit_examples = 10\n", + "from src.config import root_folder\n", + "tokenizer_args=dict(padding=\"max_length\", max_length=cfg.max_length, truncation=True, add_special_tokens=True)\n", + " \n", + "def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2):\n", + " \"\"\"\n", + " We want one set of interventions per model\n", + " \n", + " So we always load a cached version if possible. to make it approx repeatable use the same dataset etc\n", + " \"\"\"\n", + " model_name = cfg.model.replace('/', '-')\n", + " intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'\n", + " intervention_f.parent.mkdir(exist_ok=True, parents=True)\n", + " if not intervention_f.exists():\n", + " rep_token = -1\n", + " n_difference = 1\n", + " direction_method = 'pca'\n", + " \n", + " hidden_layers = list(range(8, model.config.num_hidden_layers, 3))\n", + " \n", + " dataset_fit = load_preproc_dataset('imdb', cfg, tokenizer, N=N_fit_examples)\n", + " \n", + " rep_reading_pipeline = pipeline(\"rep-reading\", model=model, tokenizer=tokenizer)\n", + " honesty_rep_reader = rep_reading_pipeline.get_directions(\n", + " dataset_fit['question'], \n", + " rep_token=rep_token, \n", + " hidden_layers=hidden_layers, \n", + " n_difference=n_difference, \n", + " train_labels=dataset_fit['label_true'], \n", + " direction_method=direction_method,\n", + " batch_size=batch_size,\n", + " **tokenizer_args\n", + " )\n", + " # and save\n", + " with open(intervention_f, 'wb') as f:\n", + " pickle.dump(honesty_rep_reader, f)\n", + " logger.info(f'Saved interventions to {intervention_f}')\n", + " else:\n", + " with open(intervention_f, 'rb') as f:\n", + " honesty_rep_reader = pickle.load(f)\n", + " logger.info(f'Loaded interventions from {intervention_f}')\n", + " \n", + " return honesty_rep_reader\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "N_fit_examples = 20\n", + "honesty_rep_reader = load_rep_reader(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size)\n", + "\n", + "hidden_layers = honesty_rep_reader.directions.keys()\n", "hidden_layers\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Dataset" + ] + }, { "cell_type": "code", "execution_count": null, @@ -136,16 +208,9 @@ "# load dataset\n", "ds_name = 'imdb'\n", "ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer)\n", - "ds_tokens\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "N_fit_examples = 10\n", + "ds_tokens\n", + "\n", + "\n", "N_train_split = (len(ds_tokens) - N_fit_examples) //2\n", "\n", "# split the dataset, it's preshuffled\n", @@ -161,28 +226,20 @@ "metadata": {}, "outputs": [], "source": [ - "tokenizer_args=dict(padding=\"max_length\", max_length=cfg.max_length, truncation=True, add_special_tokens=True)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# fit\n", - "train_labels = dataset_fit['label_true']\n", - "honesty_rep_reader = rep_reading_pipeline.get_directions(\n", - " dataset_fit['question'], \n", - " rep_token=rep_token, \n", - " hidden_layers=hidden_layers, \n", - " n_difference=n_difference, \n", - " train_labels=dataset_fit['label_true'], \n", - " direction_method=direction_method,\n", - " batch_size=batch_size,\n", - " **tokenizer_args\n", - ")\n", - "honesty_rep_reader\n" + "# # fit\n", + "# # FIXME: load or save here if the inputs are the same....\n", + "# train_labels = dataset_fit['label_true']\n", + "# honesty_rep_reader = rep_reading_pipeline.get_directions(\n", + "# dataset_fit['question'], \n", + "# rep_token=rep_token, \n", + "# hidden_layers=hidden_layers, \n", + "# n_difference=n_difference, \n", + "# train_labels=dataset_fit['label_true'], \n", + "# direction_method=direction_method,\n", + "# batch_size=batch_size,\n", + "# **tokenizer_args\n", + "# )\n", + "# honesty_rep_reader\n" ] }, { @@ -214,7 +271,7 @@ "metadata": {}, "outputs": [], "source": [ - "inputs = dataset_train.select([0, 10, 30])[:3]\n", + "inputs = dataset_train[:3]\n", "inputs['question']\n" ] }, @@ -265,11 +322,18 @@ " \"rep-control2\", \n", " model=model, \n", " tokenizer=tokenizer, \n", - " layers=layer_id, \n", + " layers=hidden_layers, \n", " max_length=cfg.max_length,)\n", "rep_control_pipeline2\n" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, { "cell_type": "code", "execution_count": null, @@ -278,13 +342,10 @@ "source": [ "\n", "coeff=8.0\n", - "max_new_tokens=3\n", - "text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, \n", - " output_hidden_states=True, return_dict=True, max_length=cfg.max_length,\n", - " )\n", + "max_new_tokens=1\n", "\n", "activations = {}\n", - "for layer in layer_id:\n", + "for layer in hidden_layers:\n", " activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()\n", " \n", "\n", @@ -292,13 +353,149 @@ "\n", "model.eval()\n", "with torch.no_grad():\n", - " baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size, **text_gen_kwargs)\n", - " control_outputs = rep_control_pipeline2(inputs, activations=activations, batch_size=batch_size, **text_gen_kwargs)\n", - " control_outputs_neg = rep_control_pipeline2(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)\n", + " baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size)\n", + " control_outputs = rep_control_pipeline2(inputs, activations=activations, batch_size=batch_size)\n", + " control_outputs_neg = rep_control_pipeline2(inputs, activations=activations_neg, batch_size=batch_size)\n", "\n", "\n", "metrics(control_outputs_neg, baseline_outputs, control_outputs)\n" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from torch.utils.data import Dataset\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# TODO we need to save and cache for many split and datasets\n", + "rep_control_pipeline2\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "# from datasets import Dataset, DatasetInfo\n", + "import datasets\n", + "from src.config import root_folder\n", + "from pathvalidate import sanitize_filename\n", + "from src.helpers.ds import ds_keep_cols\n", + "\n", + "\n", + "def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch_size=2, split_type=\"train\"):\n", + " \"create a dataset of hidden states.\"\"\"\n", + " \n", + " N = len(ds_tokens)\n", + " dataset_name = sanitize_filename(f\"{cfg.model}_{ds_name}_{split_type}_{N}\", replacement_text=\"_\")\n", + " f = root_folder / '.ds'/ f\"{dataset_name}\"\n", + " \n", + " info_kwargs = dict(extract_cfg=cfg.to_dict(), ds_name=ds_name, split_type=split_type, f=f, date=pd.Timestamp.now().isoformat(),)\n", + " \n", + " torch_cols = ['input_ids', 'attention_mask', 'choice_ids', 'question', 'answer_choices', 'example_i', 'label_true', 'sys_instr_name', 'template_name', 'instructed_to_lie']\n", + " ds_t_subset = ds_keep_cols(ds_tokens, torch_cols)\n", + " ds = ds_t_subset.to_iterable_dataset()\n", + " # pipeline_it = rep_control_pipeline2(ds, batch_size=batch_size, **text_gen_kwargs)\n", + " \n", + " # first we make the calibration dataset with no intervention\n", + " gen_kwargs = dict(\n", + " model_inputs=ds,\n", + " activations=activations,\n", + " batch_size=batch_size,\n", + " )\n", + " \n", + " ds1 = datasets.Dataset.from_generator(\n", + " generator=pipeline,\n", + " info=datasets.DatasetInfo(\n", + " description=json.dumps(info_kwargs, indent=2),\n", + " config_name=f,\n", + " ),\n", + " gen_kwargs=gen_kwargs,\n", + " num_proc=1,\n", + " )\n", + " return ds1\n", + "\n", + "\n", + "\n", + "\n", + "ds1 = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type=\"train\")\n", + "ds1\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "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": {}, + "source": [ + "# To Datasets\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "baseline_outputs\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/make_dataset.py b/notebooks/make_dataset.py index 45ec6ce..6d07e46 100644 --- a/notebooks/make_dataset.py +++ b/notebooks/make_dataset.py @@ -24,7 +24,7 @@ from torch import Tensor import pickle import hashlib from pathlib import Path - +from pathvalidate import sanitize_filename import transformers from transformers import GPTQConfig from datasets import Dataset, DatasetInfo @@ -278,7 +278,7 @@ def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrained def row_choice_ids(r, tokenizer): - return choice2ids([[c] for c in r['answer_choices']], tokenizer) + return choice2ids([c for c in r['answer_choices']], tokenizer) def expand_choices(choices: List[str]) -> Set[str]: @@ -419,11 +419,8 @@ if __name__ == "__main__": model, tokenizer = load_model(cfg.model) - - - sanitize = lambda s:s.replace('/', '').replace('-', '_') if s is not None else s ds_name = 'imdb' - model_name = sanitize(cfg.model) + model_name = sanitize_filename(cfg.model) intervention, intervention_fn = load_intervention(ds_name, cfg, model, tokenizer, model_name) for ds_name in ds_names: @@ -433,7 +430,7 @@ if __name__ == "__main__": # ## Save as Huggingface Dataset # get dataset filename N = len(ds_tokens) - dataset_name = f"{sanitize(cfg.model)}_{ds_name}_{split_type}_{N}" + dataset_name = f"{sanitize_filename(cfg.model)}_{ds_name}_{split_type}_{N}" f = root_folder / '.ds'/ f"{dataset_name}" ds1 = create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts=intervention, f=str(f)) diff --git a/poetry.lock b/poetry.lock index 03a6d63..5c735b9 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2001,6 +2001,21 @@ files = [ {file = "pathspec-0.11.2.tar.gz", hash = "sha256:e0d8d0ac2f12da61956eb2306b69f9469b42f4deb0f3cb6ed47b9cce9996ced3"}, ] +[[package]] +name = "pathvalidate" +version = "3.2.0" +description = "pathvalidate is a Python library to sanitize/validate a string such as filenames/file-paths/etc." +optional = false +python-versions = ">=3.7" +files = [ + {file = "pathvalidate-3.2.0-py3-none-any.whl", hash = "sha256:cc593caa6299b22b37f228148257997e2fa850eea2daf7e4cc9205cef6908dee"}, + {file = "pathvalidate-3.2.0.tar.gz", hash = "sha256:5e8378cf6712bff67fbe7a8307d99fa8c1a0cb28aa477056f8fc374f0dff24ad"}, +] + +[package.extras] +docs = ["Sphinx (>=2.4)", "sphinx-rtd-theme (>=1.2.2)", "urllib3 (<2)"] +test = ["Faker (>=1.0.8)", "allpairspy (>=2)", "click (>=6.2)", "pytest (>=6.0.1)", "pytest-discord (>=0.1.4)", "pytest-md-report (>=0.4.1)"] + [[package]] name = "peft" version = "0.5.0" @@ -3868,4 +3883,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.13" -content-hash = "6068bb1c5ec5042e14fe98ec08112adb751f51829405378a393df028baf8c881" +content-hash = "b8ab01ec55c8189d322df8be511d43175149fdbc5d5f11c91e7a5f1d91748f24" diff --git a/pyproject.toml b/pyproject.toml index bab0486..b039692 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ baukit = {git = "https://github.com/davidbau/baukit.git"} eleuther-elk = "0.1.1" scikit-learn = "^1.3.1" pytorch-optimizer = "^2.12.0" +pathvalidate = "^3.2.0" [[tool.poetry.source]] name = "pytorch" diff --git a/src/datasets/hs.py b/src/datasets/hs.py index 7d7de76..f7077b9 100644 --- a/src/datasets/hs.py +++ b/src/datasets/hs.py @@ -28,7 +28,7 @@ import torch.nn.functional as F from baukit.nethook import Trace, TraceDict, recursive_copy from einops import rearrange, reduce, repeat from src.datasets.scores import choice2id, choice2ids -from src.helpers.torch import clear_mem +from src.helpers.torch import clear_mem, detachcpu from collections import defaultdict from dataclasses import field from src.datasets.intervene import InterventionDict, intervention_meta_fn @@ -217,15 +217,4 @@ class ExtractHiddenStates: layers_inds = sorted(set(list(strided_layers)+list(last_few))) return [layer_names[i] for i in layers_inds] -def detachcpu(x): - """ - Trys to convert torch if possible a single item - """ - if isinstance(x, torch.Tensor): - # note apache parquet doesn't support half to we go for float https://github.com/huggingface/datasets/issues/4981 - x = x.detach().cpu().float() - if x.squeeze().dim()==0: - return x.item() - return x - else: - return x + diff --git a/src/datasets/scores.py b/src/datasets/scores.py index 63ca910..c823716 100644 --- a/src/datasets/scores.py +++ b/src/datasets/scores.py @@ -10,15 +10,8 @@ from transformers import ( PreTrainedModel ) -default_class2choices = {False: ['No', 'Negative', 'no', 'false', 'wrong', 'False'], True: ['Yes', 'Positive', 'yes', 'true', 'correct', 'right', 'True']} +default_class2choices = [['No', 'Negative', 'negative', 'no', 'false', 'wrong', 'False', '0'], ['Yes', 'Positive', 'positive', 'yes', 'true', 'correct', 'right', 'True', '1']] -# def class2choices_to_choices(class2choices): -# return [class2choices[i][0] for i in sorted(class2choices)] - -# def label_to_choice(label: bool, class2choices=default_class2choices) -> str: -# """turns a label like 0 to a choice like No""" -# choices = class2choices_to_choices(class2choices) -# return choices[label] def scores2choice_probs(row, class2_ids: List[List[int]], keys=["scores0"], prefix=""): """ Given next_token scores (logits) we take only the subset the corresponds to our @@ -79,23 +72,14 @@ def choice2id(tokenizer, c: str, whitespace_first=False) -> List[int]: def choice2ids(all_choices: List[List[str]], tokenizer: PreTrainedTokenizer) -> List[List[int]]: choices = [list(itertools.chain(*[choice2id(tokenizer, c) for c in choices])) for choices in all_choices] - assert choices[0]!=choices[1], "choices should be different" + assert choices[0]!=choices[1], f"choices should be different but were not {all_choices}" assert choices[0][0]!=choices[1][0], "choices should be different" return choices -# def get_choice_as_token(tokenizer, choice: str) -> int: -# return get_choices_as_tokens(tokenizer, [choice])[0] -# def get_choices_as_tokens( -# tokenizer, choices:List[str] = ["Positive"], whitespace_first=True -# ) -> List[int]: - - -# ids = [] -# for c in choices: -# try: -# id_ = choice2id(tokenizer, c) -# ids.append(id_) -# except AssertionError as e: -# print(e) -# return ids +def scores2choice_probs2(logits, choiceids: List[List[int]]): + """calculate the probability for each group of choices.""" + assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}" + probs = torch.softmax(logits, 0) # shape [tokens, inferences) + probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1) # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']] + return probs_c diff --git a/src/helpers/__init__.py b/src/helpers/__init__.py index 731a8a5..264ce5f 100644 --- a/src/helpers/__init__.py +++ b/src/helpers/__init__.py @@ -5,3 +5,4 @@ def bool2switch(x): def switch2bool(x): """[-1,1]->[0,1]""" return (x+1)/2 + diff --git a/src/helpers/torch.py b/src/helpers/torch.py index 203f826..3fdec05 100644 --- a/src/helpers/torch.py +++ b/src/helpers/torch.py @@ -44,3 +44,16 @@ def clear_mem(): gc.collect() torch.cuda.empty_cache() gc.collect() + +def detachcpu(x): + """ + Trys to convert torch if possible a single item + """ + if isinstance(x, torch.Tensor): + # note apache parquet doesn't support half to we go for float https://github.com/huggingface/datasets/issues/4981 + x = x.detach().cpu().float() + if x.squeeze().dim()==0: + return x.item() + return x + else: + return x diff --git a/src/prompts/prompt_loading.py b/src/prompts/prompt_loading.py index c6c687d..3604a9e 100644 --- a/src/prompts/prompt_loading.py +++ b/src/prompts/prompt_loading.py @@ -232,6 +232,12 @@ def _convert_to_prompts( for template in templates: answer_choices=template.get_fixed_answer_choices_list() + + # skip prompts where the responses are similar in the first token + if answer_choices[0][:3]==answer_choices[1][:3]: + logger.debug(f"skipping prompt because it's answers are not unique (for the first token): {template.name} {answer_choices}") + continue + answer_choices = [[c] for c in answer_choices] for instructed_to_lie in [False, True]: for sys_instr_name, sys_instr in sys_instructions[instructed_to_lie].items(): fake_example = example.copy() @@ -256,7 +262,7 @@ def _convert_to_prompts( ] for d in fewshot_texts: # some of the answers have extra trailing text, that's OK. But extra preceeding text is not, let's check for that - assert any([d['response'].startswith(a) for a in answer_choices]), f"fewshot response `{d['response']}` has extra preceeding text compared to allowed choices: {answer_choices}. template is: {template.name}" + assert any([any([d['response'].startswith(a) for a in ac]) for ac in answer_choices]), f"fewshot response `{d['response']}` has extra preceeding text compared to allowed choices: {answer_choices}. template is: {template.name}" prompt_parts = fewshot_texts + prompt_parts prompt_parts[0]['system'] = sys_instr diff --git a/src/repe/rep_control_pipeline_baukit.py b/src/repe/rep_control_pipeline_baukit.py index 23a7832..ee1315f 100644 --- a/src/repe/rep_control_pipeline_baukit.py +++ b/src/repe/rep_control_pipeline_baukit.py @@ -1,3 +1,4 @@ +import re import torch from transformers.pipelines import ( TextGenerationPipeline, @@ -5,17 +6,26 @@ from transformers.pipelines import ( Pipeline, ) from transformers.pipelines.base import GenericTensor +from datasets import Dataset from typing import List, Tuple, Dict, Any, Union, NewType from baukit.nethook import Trace, TraceDict, recursive_copy from functools import partial -from src.datasets.scores import choice2ids +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 + +Activations = NewType("Activations", Dict[str, torch.Tensor]) -Activations = NewType("InterventionDict", Dict[str, torch.Tensor]) +def hacky_sanitize_outputs(o): + """I can't find the mem leak, so lets just detach, cpu, clone, freemem.""" + o = {k: detachcpu(v) for k, v in o.items()} + o = recursive_copy(o, detach=True, clone=True) + clear_mem() + return o def row_choice_ids(answer_choices, tokenizer): - return choice2ids([[c] for c in answer_choices], tokenizer) + return choice2ids([c for c in answer_choices], tokenizer) def intervene(output, activation): @@ -51,6 +61,8 @@ class RepControlPipeline2(FeatureExtractionPipeline): super().__init__(model=model, tokenizer=tokenizer, **kwargs) self.max_length = max_length self.layer_name_tmpl = layer_name_tmpl + + # self.default_class2choiceids = choice2ids(default_class2choices, tokenizer) def __call__(self, model_inputs, activations=None, **kwargs): if activations is not None: @@ -65,24 +77,30 @@ class RepControlPipeline2(FeatureExtractionPipeline): outputs = super().__call__(model_inputs, **kwargs) return outputs - def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]: + def preprocess(self, inputs: Dataset, **tokenize_kwargs) -> Dict[str, GenericTensor]: # tokenize a batch of inputs return_tensors = self.framework - model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs) - return {**inputs, **model_inputs} + if 'input_ids' not in inputs: + model_inputs = self.tokenizer(inputs['question'], return_tensors=return_tensors, return_attention_mask=True, add_special_tokens=True, truncation=True, padding="max_length", max_length=self.max_length, **tokenize_kwargs) + return {**inputs, **model_inputs} + else: + return inputs def _forward(self, model_inputs): inputs = dict(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask']) inputs.update( {"use_cache": False, "output_hidden_states": True, "return_dict": True} ) + self.model.eval() with torch.no_grad(): model_outputs = self.model(**inputs) + # o = {k: detachcpu(v) for k, v in o.items()} + # o = recursive_copy(o) + # clear_mem() # retain some of the inputs - keep_cols = ["answer_choices", "input_ids", "attention_mask"] model_outputs = {**model_inputs, **model_outputs} - return model_outputs + return hacky_sanitize_outputs(model_outputs) def postprocess(self, o): # note this sometimes deals with a batch, sometimes with a single result. infuriating @@ -90,20 +108,22 @@ class RepControlPipeline2(FeatureExtractionPipeline): # This is called once for each result, but the text pipeline is set up to hande multiple... # This is called once for each result, but the text pipeline is set up to hande multiple... o["end_logits"] = o["logits"][:, -1, :].float() - # hidden_states = list(o.hidden_states) o["input_truncated"] = self.tokenizer.batch_decode(o['input_ids']) o["truncated"] = torch.sum(o["attention_mask"], 1)==self.max_length o["text_ans"] = self.tokenizer.batch_decode(o["end_logits"].argmax(-1)) - o['choice_ids'] = [row_choice_ids(ac, self.tokenizer) for ac in o['answer_choices']] + + if 'answer_choices' in o: + answer_choices = o['answer_choices'] + if isinstance(answer_choices[0][0], str): + answer_choices = [answer_choices] + else: + answer_choices = default_class2choices # self.default_class2choiceids + + o['choice_ids'] = [row_choice_ids(ac, self.tokenizer) for ac in answer_choices] p = o['add_ans'] = torch.stack([scores2choice_probs2(l, c) for l, c in zip(o['end_logits'], o['choice_ids'])]) o['ans'] = p[:, 1] / (torch.sum(p, 1) + 1e-5) + o = hacky_sanitize_outputs(o) return o -def scores2choice_probs2(logits, choiceids: List[List[int]]): - """calculate the probability for each group of choices.""" - assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}" - probs = torch.softmax(logits, 0) # shape [tokens, inferences) - probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1) # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']] - return probs_c