From bd1229636c850ddc06229f4f45dac028109a45c3 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Sun, 4 May 2025 14:54:13 +0800 Subject: [PATCH] wi[ --- .gitignore | 1 + nbs/02b_TQA_regr_w_kv.ipynb | 773 +++++++++++++++++++++++++++++++----- 2 files changed, 664 insertions(+), 110 deletions(-) diff --git a/.gitignore b/.gitignore index fa5631b..6eb1488 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ wheels/ # Virtual environments .venv nbs/old/ +outputs/ diff --git a/nbs/02b_TQA_regr_w_kv.ipynb b/nbs/02b_TQA_regr_w_kv.ipynb index e67d229..6bcae01 100644 --- a/nbs/02b_TQA_regr_w_kv.ipynb +++ b/nbs/02b_TQA_regr_w_kv.ipynb @@ -13,7 +13,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 67, "metadata": {}, "outputs": [], "source": [ @@ -23,7 +23,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 68, "metadata": {}, "outputs": [], "source": [ @@ -35,7 +35,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 69, "metadata": {}, "outputs": [], "source": [ @@ -74,7 +74,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 70, "metadata": {}, "outputs": [], "source": [ @@ -101,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 71, "metadata": {}, "outputs": [], "source": [ @@ -125,21 +125,38 @@ "\n", "# model_name = \"unsloth/Phi-4-mini-instruct\" # 4b\n", "# model_name = \"stelterlab/phi-4-AWQ\"\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 72, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "68b61ba11b344abc87c84919121f02cb", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/2 [00:00<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|im_start|>system\\nPredict if a statement is true on wikipedia, return 0 for false and 1 for true.\\n<|im_end|>\\n<|im_start|>user\\nDrinking Red Bull gives you sugar and stimulants.<|im_end|>\\n<|im_start|>assistant\\n\\n\\n\\n\\nThe answer is '" + ] + }, + "execution_count": 49, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "tokenizer.batch_decode(ds2['input_ids'])[0]" ] @@ -214,7 +243,15 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], "source": [ "collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n", "ds = DataLoader(ds2, batch_size=batch_size, collate_fn=collate_fn)\n", @@ -233,17 +270,100 @@ "execution_count": null, "metadata": {}, "outputs": [], + "source": [ + "# # choose layers to cache\n", + "# n_layers = model.config.num_hidden_layers\n", + "# a = int(0.3*n_layers)\n", + "# b = n_layers-2\n", + "# layer_groups = {\n", + "# 'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][a:b],\n", + "# 'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][a:b],\n", + "# 'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')][a:b],\n", + "# }\n", + "# layer_groups" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'mlp.down_proj': ['model.layers.8.mlp.down_proj',\n", + " 'model.layers.9.mlp.down_proj',\n", + " 'model.layers.10.mlp.down_proj',\n", + " 'model.layers.11.mlp.down_proj',\n", + " 'model.layers.12.mlp.down_proj',\n", + " 'model.layers.13.mlp.down_proj',\n", + " 'model.layers.14.mlp.down_proj',\n", + " 'model.layers.15.mlp.down_proj',\n", + " 'model.layers.16.mlp.down_proj',\n", + " 'model.layers.17.mlp.down_proj',\n", + " 'model.layers.18.mlp.down_proj',\n", + " 'model.layers.19.mlp.down_proj',\n", + " 'model.layers.20.mlp.down_proj',\n", + " 'model.layers.21.mlp.down_proj',\n", + " 'model.layers.22.mlp.down_proj',\n", + " 'model.layers.23.mlp.down_proj',\n", + " 'model.layers.24.mlp.down_proj',\n", + " 'model.layers.25.mlp.down_proj'],\n", + " 'self_attn': ['model.layers.8.self_attn',\n", + " 'model.layers.9.self_attn',\n", + " 'model.layers.10.self_attn',\n", + " 'model.layers.11.self_attn',\n", + " 'model.layers.12.self_attn',\n", + " 'model.layers.13.self_attn',\n", + " 'model.layers.14.self_attn',\n", + " 'model.layers.15.self_attn',\n", + " 'model.layers.16.self_attn',\n", + " 'model.layers.17.self_attn',\n", + " 'model.layers.18.self_attn',\n", + " 'model.layers.19.self_attn',\n", + " 'model.layers.20.self_attn',\n", + " 'model.layers.21.self_attn',\n", + " 'model.layers.22.self_attn',\n", + " 'model.layers.23.self_attn',\n", + " 'model.layers.24.self_attn',\n", + " 'model.layers.25.self_attn'],\n", + " 'mlp.up_proj': ['model.layers.8.mlp.up_proj',\n", + " 'model.layers.9.mlp.up_proj',\n", + " 'model.layers.10.mlp.up_proj',\n", + " 'model.layers.11.mlp.up_proj',\n", + " 'model.layers.12.mlp.up_proj',\n", + " 'model.layers.13.mlp.up_proj',\n", + " 'model.layers.14.mlp.up_proj',\n", + " 'model.layers.15.mlp.up_proj',\n", + " 'model.layers.16.mlp.up_proj',\n", + " 'model.layers.17.mlp.up_proj',\n", + " 'model.layers.18.mlp.up_proj',\n", + " 'model.layers.19.mlp.up_proj',\n", + " 'model.layers.20.mlp.up_proj',\n", + " 'model.layers.21.mlp.up_proj',\n", + " 'model.layers.22.mlp.up_proj',\n", + " 'model.layers.23.mlp.up_proj',\n", + " 'model.layers.24.mlp.up_proj',\n", + " 'model.layers.25.mlp.up_proj']}" + ] + }, + "execution_count": 65, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# choose layers to cache\n", "n_layers = model.config.num_hidden_layers\n", - "a = int(0.3*n_layers)\n", + "a = int(0.5*n_layers)\n", "b = n_layers-2\n", + "select = slice(a, b, 3)\n", "layer_groups = {\n", - " 'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][a:b],\n", - " 'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][a:b],\n", - " 'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')][a:b],\n", + " 'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][select],\n", + " 'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][select],\n", + " 'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')][select],\n", "}\n", - "# layer_groups = []" + "layer_groups" ] }, { @@ -252,8 +372,44 @@ "metadata": {}, "outputs": [], "source": [ + "import os, unicodedata, string\n", + "from pathlib import Path\n", "\n", - "acts_outfile = Path(f'/tmp/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}.parquet')\n", + "def sanitize_path(path: Path | str, allow_period: bool = True) -> Path:\n", + " \"\"\"\n", + " Whitelist only ASCII letters, digits, dash, underscore,\n", + " optionally period, and forward‐slash. Replace others with '_'.\n", + " \"\"\"\n", + " s = unicodedata.normalize(\"NFKD\", str(path))\\\n", + " .encode(\"ascii\", \"ignore\")\\\n", + " .decode()\n", + " s = s.replace(os.sep, \"/\")\n", + " allowed = set(string.ascii_letters + string.digits + \"_-\")\n", + " if allow_period: allowed.add(\".\")\n", + " allowed.add(\"/\")\n", + " return Path(\"\".join(ch if ch in allowed else \"_\" for ch in s))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-110.parquet')" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "\n", + "acts_outfile = Path(f'/tmp/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}_v2.parquet')\n", + "acts_outfile = sanitize_path(acts_outfile)\n", "acts_outfile" ] }, @@ -261,7 +417,29 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "{'mlp.down_proj': ['model.layers.14.mlp.down_proj',\n", + " 'model.layers.17.mlp.down_proj',\n", + " 'model.layers.20.mlp.down_proj',\n", + " 'model.layers.23.mlp.down_proj'],\n", + " 'self_attn': ['model.layers.14.self_attn',\n", + " 'model.layers.17.self_attn',\n", + " 'model.layers.20.self_attn',\n", + " 'model.layers.23.self_attn'],\n", + " 'mlp.up_proj': ['model.layers.14.mlp.up_proj',\n", + " 'model.layers.17.mlp.up_proj',\n", + " 'model.layers.20.mlp.up_proj',\n", + " 'model.layers.23.mlp.up_proj']}" + ] + }, + "execution_count": 53, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "def collect_all_tokens(*args, **kwargs):\n", " return default_postprocess_result(*args, **kwargs, last_token=False)\n", @@ -277,7 +455,35 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "d4d0acec794f4102a361592f358ed92b", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading dataset shards: 0%| | 0/53 [00:00 or newline)\n" + ] + } + ], "source": [ "def get_uniq_token_ids(tokens):\n", " token_ids = tokenizer(\n", @@ -433,7 +763,21 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "Dataset({\n", + " features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts'],\n", + " num_rows: 316\n", + "})" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# now we map to 1) calc supressed activations 2) llm answer (prob of 0 vs prob of 1)\n", "\n", @@ -465,6 +809,7 @@ "\n", " o[\"hidden_states\"] = hs.half()\n", " o[\"supr_amounts\"] = supr_amounts.half()\n", + " o['logits'] = o['logits'][-1].half()\n", " return o\n", "\n", "\n", @@ -476,7 +821,19 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "before ['0', '0 ', '0\\n', 'false', 'False ']\n", + "after ['false', 'False', '0']\n", + "before ['1', '1 ', '1\\n', 'true', 'True ']\n", + "after ['1', 'True', 'true']\n", + "QC: manually check that these are equivilent (no or newline)\n" + ] + } + ], "source": [ "model = Wo = Wo_inv = tokenizer = None\n", "clear_mem()" @@ -486,7 +843,22 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "17b983edf4d64ca995dc90e063b4eb8a", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Map: 0%| | 0/316 [00:00 torch.Tensor:\n", + " \"\"\"Better torch.quantile for one SCALAR quantile.\n", + "\n", + " Using torch.kthvalue. Better than torch.quantile because:\n", + " - No 2**24 input size limit (pytorch/issues/67592),\n", + " - Much faster, at least on big input sizes.\n", + "\n", + " Arguments:\n", + " input (torch.Tensor): See torch.quantile.\n", + " q (float): See torch.quantile. Supports only scalar input\n", + " currently.\n", + " dim (int | None): See torch.quantile.\n", + " keepdim (bool): See torch.quantile. Supports only False\n", + " currently.\n", + " interpolation: {\"nearest\", \"lower\", \"higher\"}\n", + " See torch.quantile.\n", + " out (torch.Tensor | None): See torch.quantile. Supports only\n", + " None currently.\n", + " \"\"\"\n", + " # Sanitization: q\n", + " try:\n", + " q = float(q)\n", + " assert 0 <= q <= 1\n", + " except Exception:\n", + " raise ValueError(f\"Only scalar input 0<=q<=1 is currently supported (got {q})!\")\n", + "\n", + " # Sanitization: dim\n", + " # Because one cannot pass `dim=None` to `squeeze()` or `kthvalue()`\n", + " if dim_was_none := dim is None:\n", + " dim = 0\n", + " input = input.reshape((-1,) + (1,) * (input.ndim - 1))\n", + "\n", + " # Sanitization: inteporlation\n", + " if interpolation == \"nearest\":\n", + " inter = round\n", + " elif interpolation == \"lower\":\n", + " inter = floor\n", + " elif interpolation == \"higher\":\n", + " inter = ceil\n", + " else:\n", + " raise ValueError(\n", + " \"Supported interpolations currently are {'nearest', 'lower', 'higher'} \"\n", + " f\"(got '{interpolation}')!\"\n", + " )\n", + "\n", + " # Sanitization: out\n", + " if out is not None:\n", + " raise ValueError(f\"Only None value is currently supported for out (got {out})!\")\n", + "\n", + " # Logic\n", + " k = inter(q * (input.shape[dim] - 1)) + 1\n", + " out = torch.kthvalue(input, k, dim, keepdim=True, out=out)[0]\n", + "\n", + " # Rectification: keepdim\n", + " if keepdim:\n", + " return out\n", + " if dim_was_none:\n", + " return out.squeeze()\n", + " else:\n", + " return out.squeeze(dim)\n", + "\n", + " return out" + ] + }, { "cell_type": "code", "execution_count": null, @@ -682,10 +1185,11 @@ "def magnitude_filtered_post_softmax(x: FiltIn, quantile=0.9) -> FiltOut:\n", " \"\"\"Filter out tokens with abnormally high post-softmax values\"\"\"\n", " # Apply softmax to get attention-like weights\n", - " weights = torch.softmax(x, dim=-1)\n", + " weights = torch.softmax(x, dim=-1).max(dim=-1, keepdim=True).values\n", " \n", " # Create mask for tokens below threshold\n", - " threshold = torch.quantile(weights, quantile)\n", + " # FIXME should be independant of batch\n", + " threshold = torch_quantile(weights, quantile)\n", " mask = (weights <= threshold)\n", " \n", " # # Ensure we don't filter everything out\n", @@ -707,10 +1211,13 @@ "source": [ "def quantile_filtered(x: FiltIn, quantile=.9) -> FiltOut:\n", " \"\"\"Filter out tokens with attention weights above a quantile threshold\"\"\"\n", - " weights = torch.softmax(x, dim=-1)\n", - " threshold = torch.quantile(weights, quantile)\n", + " weights = torch.log_softmax(x, dim=-1)#.max(dim=-1, keepdim=True).values\n", + " # can't use max as there are too many attention sinks\n", + " # print(weights)\n", + " # FIXME should be independant of batch\n", + " threshold = torch_quantile(weights, quantile)\n", " mask = weights <= threshold\n", - " print(weights.shape, mask.shape, mask.float().mean(), weights)\n", + " # print(weights.shape, mask.shape, mask.float().mean(), threshold)\n", " if mask.sum() == 0:\n", " logger.warning(f\"All tokens filtered out quantile={quantile}.\")\n", " return x * mask, mask" @@ -722,10 +1229,10 @@ "metadata": {}, "outputs": [], "source": [ - "X = ds_a2['hidden_states']\n", - "print(X.shape)\n", - "X2, mask = quantile_filtered(X)\n", - "mask.float().mean()" + "# X = ds_a2['hidden_states']\n", + "# print(X.shape)\n", + "# X2, mask = quantile_filtered(X)\n", + "# mask.float().mean()" ] }, { @@ -756,7 +1263,8 @@ " entropy_per_token = entropy_per_token / denom\n", " \n", " # Get statistics across tokens\n", - " threshold = torch.quantile(entropy_per_token, quantile)\n", + " # FIXME should be independant of batch\n", + " threshold = torch_quantile(entropy_per_token, quantile)\n", " \n", " # Create mask for tokens below threshold\n", " mask = entropy_per_token <= threshold\n", @@ -774,7 +1282,39 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n", + "threshold: 0.02307834103703499\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n", + "threshold: 0.02307834103703499\n" + ] + }, + { + "data": { + "text/plain": [ + "(torch.Size([316, 13, 110, 2048]),\n", + " torch.Size([316, 13, 110, 2048]),\n", + " tensor(4856864.5000),\n", + " tensor(4856636.5000),\n", + " torch.Size([316, 13, 110, 1]),\n", + " tensor(0.9752))" + ] + }, + "execution_count": 81, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "# X = ds_a2['hidden_states']\n", "# X2, mask = entropy_guided_filter(X)\n", @@ -803,6 +1343,8 @@ "def filter_high_magnitude(x: FiltIn, threshold_factor=2.0) -> FiltIn:\n", " \"\"\"Filter out tokens with abnormally high magnitude (potential attention sinks)\"\"\"\n", " magnitudes = torch.norm(x, dim=-1, keepdim=True)\n", + "\n", + " # FIXME should be independant of batch\n", " mean_mag = magnitudes.mean()\n", " std_mag = magnitudes.std()\n", " threshold = mean_mag + threshold_factor * std_mag\n", @@ -840,7 +1382,19 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "[('supressed_hs_-0.5', 'magnitude', 'first'),\n", + " ('supressed_hs_-0.5', 'entropy', 'sum')]" + ] + }, + "execution_count": 49, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "X_cols = [\"supr_amounts\", \"hidden_states\",] + act_groups\n", "\n", @@ -866,19 +1420,30 @@ "for eps in [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]:\n", " filters[f'magnitude_{eps}'] = lambda x: filter_high_magnitude(x, eps)\n", "\n", - "# 2. token aggregators\n", + "# # 2. token aggregators\n", + "# token_level_funcs = {\n", + "# \"min:\": lambda x: x.min(0)[0],\n", + "# \"max\": lambda x: x.max(0)[0],\n", + "# \"mean\": lambda x: x.mean(0),\n", + "# \"sum\": lambda x: x.sum(0),\n", + "# \"first\": lambda x: x[0],\n", + "# \"last\": lambda x: x[-1],\n", + "# # \"none\": lambda x: x,\n", + "# \"std\": lambda x: x.std(0),\n", + "# }\n", + "# unit test agg\n", "token_level_funcs = {\n", - " \"min:\": lambda x: x.min(0)[0],\n", - " \"max\": lambda x: x.max(0)[0],\n", - " \"mean\": lambda x: x.mean(0),\n", - " \"sum\": lambda x: x.sum(0),\n", - " \"first\": lambda x: x[0],\n", - " \"last\": lambda x: x[-1],\n", - " \"none\": lambda x: x,\n", - " \"std\": lambda x: x.std(0),\n", + " \"min:\": lambda x: x.min(2)[0],\n", + " \"max\": lambda x: x.max(2)[0],\n", + " \"mean\": lambda x: x.mean(2),\n", + " \"sum\": lambda x: x.sum(2),\n", + " \"first\": lambda x: x[:, :, 0],\n", + " \"last\": lambda x: x[:, :, -1],\n", + " # \"none\": lambda x: x,\n", + " \"flatten\": lambda x: x.flatten(2),\n", + " \"std\": lambda x: x.std(2),\n", "}\n", "\n", - "\n", "# now get and shuffle all perms\n", "perms = []\n", "for k, _ in datasets.items():\n", @@ -898,16 +1463,17 @@ "metadata": {}, "outputs": [], "source": [ - "# unit test filter\n", - "X0 = ds_a2['hidden_states']\n", - "for k,v in filters.items():\n", - " X0_norm = X0.norm()\n", - " X2, mask = v(X0) # bad\n", - " print(f\"filter={k}\\n\\tmask.mean()={mask.float().mean():2.2%} \\n\\tmask={mask.shape}, \\n\\toutput={X0.shape}->{X2.shape}, \\n\\tnorm={X0_norm}->{X.norm():2.6f}={(X0_norm-X2.norm())/X0_norm:.6%}\")\n", - " assert mask.float().mean() < 1.0\n", - " assert mask.float().mean() > 0.0\n", - " assert X2.shape == X0.shape\n", - " assert X2.norm() < X0_norm" + "# # unit test filter\n", + "# X0 = ds_a2['hidden_states']\n", + "# for k,v in filters.items():\n", + "# X0_norm = X0.norm()\n", + "# X2, mask = v(X0) # bad\n", + "# print(f\"filter={k}\\n\\tmask.mean()={mask.float().mean():2.2%} \\n\\tmask={mask.shape}, \\n\\toutput={X0.shape}->{X2.shape}, \\n\\tnorm={X0_norm}->{X2.norm():2.6f}={(X0_norm-X2.norm())/X0_norm:.6%}\")\n", + "# assert mask.float().mean() < 1.0\n", + "# assert mask.float().mean() > 0.0\n", + "# assert X2.shape == X0.shape\n", + "# assert X2.norm() < X0_norm\n", + "\n" ] }, { @@ -916,27 +1482,8 @@ "metadata": {}, "outputs": [], "source": [ - "# unit test agg\n", - "token_level_funcs = {\n", - " \"min:\": lambda x: x.min(2)[0],\n", - " \"max\": lambda x: x.max(2)[0],\n", - " \"mean\": lambda x: x.mean(2),\n", - " \"sum\": lambda x: x.sum(2),\n", - " \"first\": lambda x: x[:, :, 0],\n", - " \"last\": lambda x: x[:, :, -1],\n", - " # \"none\": lambda x: x,\n", - " \"flatten\": lambda x: x.flatten(2),\n", - " \"std\": lambda x: x.std(2),\n", - "}\n", - "X = ds_a2['hidden_states']\n", - "for k,v in token_level_funcs.items():\n", - " X0_shape = X.shape\n", - " X2 = v(X)\n", - " print(f\"agg={k}\\n\\toutput={X0_shape}->{X2.shape}\")\n", - " assert X2.ndim == X.ndim - 1\n", - " assert X2.shape[0]==X.shape[0], 'same batch size'\n", - " assert X2.shape[1]==X.shape[1], 'same layers'\n", - " assert X2.shape[-1]>=X.shape[-1], 'same or greater hidden dim'" + "output_path = Path('../outputs')\n", + "output_path.mkdir(exist_ok=True, parents=True)" ] }, { @@ -948,9 +1495,13 @@ "results = []\n", "for i, (ds_key, filter_key, token_key) in tqdm(enumerate(perms), total=len(perms)):\n", "\n", - " res_f = Path(f\"{acts_outfile.stem}_{ds_key}_{filter_key}_{token_key}.json\")\n", + " name = f\"{ds_key}_{filter_key}_{token_key}\"\n", + " res_f = output_path / f\"{acts_outfile.stem}_{ds_key}_{filter_key}_{token_key}.json\"\n", + " res_f = sanitize_path(res_f)\n", " if res_f.exists():\n", - " logger.info(f\"Already processed {res_f}, skipping\")\n", + " d = json.load(res_f.open())\n", + " logger.info(f'Already processed {res_f}, skipping {d[\"score\"]:2.2f}')\n", + " results.append((name, d[\"score\"]))\n", " continue\n", "\n", " try:\n", @@ -980,6 +1531,10 @@ " \"score\": score,\n", " }\n", " json.dump(res, open(res_f, \"w\"))\n", + " results.append((name, score))\n", + " except KeyboardInterrupt:\n", + " logger.info(f\"KeyboardInterrupt, stopping {ds_key}, {filter_key}, {token_key}\")\n", + " break\n", " except Exception as e:\n", " logger.error(f\"Error processing {ds_key}, {filter_key}, {token_key}: {e}\")\n", " raise\n", @@ -993,7 +1548,9 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "X.view(len(X), -1)" + ] }, { "cell_type": "code", @@ -1047,11 +1604,11 @@ "metadata": {}, "outputs": [], "source": [ - "name = \"logits\"\n", - "X = ds_a2['logits']\n", - "score = train_linear_prob_on_dataset(X, name)\n", - "results.append((name, score))\n", - "score" + "# name = \"logits\"\n", + "# X = ds_a2['logits']\n", + "# score = train_linear_prob_on_dataset(X, name)\n", + "# results.append((name, score))\n", + "# score" ] }, { @@ -1061,12 +1618,14 @@ "outputs": [], "source": [ "\n", - "X = ds_a2['llm_ans']\n", + "X = ds_a2['llm_ans'].exp()\n", + "# X = X[:, 1] / (X[:, 0] + X[:, 1])\n", + "X = X.argmax(1)\n", "y = ds_a2['label']\n", "\n", "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_fraction, shuffle=False)\n", "\n", - "score = roc_auc_score(y_test, X_test[:, 0]).item()\n", + "score = roc_auc_score(y_test, X_test).item()\n", "if score<0.5:\n", " score = 1-score\n", "results.append(('llm_ans', score))\n", @@ -1094,13 +1653,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "# # another llm score? FIXME do I need it\n", - "# X, y = ds_a2[\"llm_log_prob_true\"] > 0, ds_a2[\"label\"]\n", - "# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_fraction, shuffle=False)\n", - "# score = roc_auc_score(X_test, y_test)\n", - "# print(f\"LLM score: {score:.2f} roc auc, n={len(X_test)}\")" - ] + "source": [] }, { "cell_type": "markdown",