{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "%reload_ext autoreload\n", "%autoreload 2" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "from datasets import load_dataset\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "from datasets import Dataset\n", "import torch\n", "\n", "from activation_store.collect import activation_store\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load model" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "model_name = \"Qwen/Qwen2.5-0.5B-Instruct\"\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " model_name,\n", " torch_dtype=\"auto\",\n", " device_map=\"auto\",\n", " attn_implementation=\"eager\", # flex_attention flash_attention_2 sdpa eager\n", ")\n", "tokenizer = AutoTokenizer.from_pretrained(model_name)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load data and tokenize" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Dataset({\n", " features: ['attention_mask', 'input_ids'],\n", " num_rows: 10\n", "})" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "N = 10\n", "max_length = 128\n", "\n", "imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)\n", "\n", "\n", "def proc(row):\n", " messages = [\n", " {\"role\":\"user\", \"content\": row['prompt'] },\n", " {\"role\":\"assistant\", \"content\": row['chosen'] }\n", " ]\n", " return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_dict=True, max_length=max_length)\n", "\n", "ds2 = imdb.map(proc).with_format(\"torch\")\n", "new_cols = set(ds2.column_names) - set(imdb.column_names)\n", "ds2 = ds2.select_columns(new_cols)\n", "ds2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Data loader" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "from torch.utils.data import DataLoader\n", "def collate_fn(examples):\n", " # Pad the batch to max length within this batch\n", " return tokenizer.pad(\n", " examples,\n", " padding=True,\n", " return_tensors=\"pt\",\n", " )\n", "ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)\n", "print(ds)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Collect activations" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "['model.layers.0.mlp.down_proj',\n", " 'model.layers.1.mlp.down_proj',\n", " 'model.layers.2.mlp.down_proj',\n", " 'model.layers.3.mlp.down_proj',\n", " 'model.layers.4.mlp.down_proj',\n", " 'model.layers.5.mlp.down_proj',\n", " 'model.layers.6.mlp.down_proj',\n", " 'model.layers.7.mlp.down_proj',\n", " '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']" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# choose layers to cache\n", "layers = [k for k,v in model.named_modules() if 'mlp.down_proj' in k]\n", "layers" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2025-02-16 09:16:55.292\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m122\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet\u001b[0m\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "90a9936ab9f94893a77fc79bf972a04f", "version_major": 2, "version_minor": 0 }, "text/plain": [ "collecting activations: 0%| | 0/3 [00:00 1\u001b[0m \u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\n", "\u001b[0;31mZeroDivisionError\u001b[0m: division by zero" ] } ], "source": [ "1/0" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## With dtypes compression - wip" ] }, { "cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "lost 0.01%\n" ] } ], "source": [ "def float_to_int16(x: torch.Tensor) -> torch.Tensor:\n", " \"\"\"Converts a floating point tensor to float16, then reinterprets as int16.\"\"\"\n", " downcast = x.type(torch.float16)\n", " # if not downcast.isfinite().all():\n", " # raise ValueError(\"Cannot convert to 16 bit: values are not finite\")\n", "\n", " return downcast.view(torch.int16)\n", "\n", "def int8_to_float32(x: torch.Tensor) -> torch.Tensor:\n", " \"\"\"Converts int16 to float16, then reinterprets as float32.\"\"\"\n", " return x.view(torch.float16).type(torch.float32)\n", "\n", "\n", "x = torch.randn(2, 3, 4)\n", "x2 = float_to_int16(x)\n", "x3 = int8_to_float32(x2)\n", "assert torch.isfinite(x3).all()\n", "assert torch.allclose(x, x3, rtol=1e-1)\n", "d = ((x-x3)/x).abs().mean()\n", "print(f'lost {d:.2%}')" ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [], "source": [ "from activation_store.collect import default_postprocess_result\n", "from datasets.features.features import cast_to_python_objects\n", "# o = cast_to_python_objects(o, only_1d_for_numpy=True, optimize_list_casting=False)\n", "\n", "def float16_postprocess_result(\n", " input, trace, output, model\n", "):\n", " o = default_postprocess_result(input, trace, output, model)\n", " # o = cast_to_python_objects(o, only_1d_for_numpy=False, optimize_list_casting=False)\n", "\n", " for k, v in o.items():\n", " if k=='attention_mask':\n", " o[k] = v.to(torch.int8)\n", " if isinstance(v, torch.Tensor) and torch.is_floating_point(v):\n", " print(k, v.dtype, v.shape, 'to int16')\n", " o[k] = float_to_int16(v)\n", " else:\n", " print('no conv', k, type(v))\n", " # o = {k: float_to_int8(v) if isinstance(v, torch.Tensor) else v\n", " # for k, v in o.items()}\n", " return o" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),\n", " 'act-model.layers.0.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.1.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.2.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.3.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.4.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.5.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.6.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.7.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.8.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.9.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.10.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.11.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.12.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.13.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.14.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.15.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.16.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.17.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.18.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.19.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.20.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.21.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.22.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'act-model.layers.23.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),\n", " 'logits': Array3D(shape=(-1, 453, 151936), dtype='int16', id=None),\n", " 'hidden_states': Array4D(shape=(-1, 25, 453, 896), dtype='int16', id=None)}" ] }, "execution_count": 42, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from datasets.arrow_writer import OptimizedTypedSequence, _ArrayXDExtensionType\n", "from datasets.features.features import Features, Array2D, Array3D, Array4D, Array5D\n", "\n", "# manually build features\n", "optimized_int_type_by_col = {\n", " \"attention_mask\": \"int8\", # binary tensor\n", " \"special_tokens_mask\": \"int8\",\n", " \"input_ids\": \"int32\", # typical vocab size: 0-50k (max ~500k, never > 1M)\n", " \"token_type_ids\": \"int8\", # binary mask; some (XLNetModel) use an additional token represented by a 2\n", "}\n", "\n", "def build_schema(d):\n", " inferred_features = Features()\n", " cols = d.keys()\n", " for col in cols:\n", " x = d[col]\n", " if col in optimized_int_type_by_col:\n", " dtype = optimized_int_type_by_col[col]\n", " typed_sequence = OptimizedTypedSequence(x, col=col)\n", " inferred_features[col] = typed_sequence.get_inferred_type()\n", " else:\n", " if x.ndim == 1:\n", " inferred_features[col] = OptimizedTypedSequence(x, col=col)\n", " inferred_features[col] = typed_sequence.get_inferred_type()\n", " shape=(-1,)+x.shape[1:]\n", " dtype = 'int16' if x.dtype == torch.float32 else x.dtype\n", " if x.ndim == 2:\n", " cls = Array2D\n", " elif x.ndim == 3:\n", " cls = Array3D\n", " elif x.ndim == 4:\n", " cls = Array4D\n", " elif x.ndim == 5:\n", " cls = Array5D\n", " else:\n", " raise ValueError(f\"Unsupported number of dimensions: {x.ndim}\")\n", " inferred_features[col] = cls(dtype=dtype, shape=shape)\n", " return inferred_features.arrow_schema\n", " # Features.from_arrow_schema(schema)\n", "\n", "d = ds_a[0:2]\n", "schema = build_schema(d)\n", "schema\n", "Features.from_arrow_schema(schema)" ] }, { "cell_type": "code", "execution_count": 43, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2025-02-16 09:35:12.432\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m155\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__c6184d05bf03be61.parquet\u001b[0m\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "1a2162645d9142c183936834d6815831", "version_major": 2, "version_minor": 0 }, "text/plain": [ "collecting activations: 0%| | 0/3 [00:00\n", "act-model.layers.0.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.1.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.2.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.3.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.4.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.5.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.6.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.7.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.8.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.9.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.10.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.11.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.12.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.13.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.14.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.15.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.16.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.17.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.18.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.19.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.20.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.21.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.22.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "act-model.layers.23.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16\n", "logits torch.float16 torch.Size([4, 453, 151936]) to int16\n", "hidden_states torch.float32 torch.Size([4, 25, 453, 896]) to int16\n" ] }, { "ename": "TypeError", "evalue": "Incompatible storage type list> for extension type extension>", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[43], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m f2 \u001b[38;5;241m=\u001b[39m \u001b[43mactivation_store\u001b[49m\u001b[43m(\u001b[49m\u001b[43mds\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlayers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlayers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mFeatures\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_arrow_schema\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 4\u001b[0m \u001b[43m \u001b[49m\u001b[43mpostprocess_result\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfloat16_postprocess_result\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 5\u001b[0m f2\n\u001b[1;32m 6\u001b[0m ds_a2 \u001b[38;5;241m=\u001b[39m Dataset\u001b[38;5;241m.\u001b[39mfrom_parquet(\u001b[38;5;28mstr\u001b[39m(f2))\u001b[38;5;241m.\u001b[39mwith_format(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/activation_store/collect.py:175\u001b[0m, in \u001b[0;36mactivation_store\u001b[0;34m(loader, model, dataset_name, layers, dataset_dir, writer_batch_size, postprocess_result, features, schema)\u001b[0m\n\u001b[1;32m 171\u001b[0m bs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(bo\u001b[38;5;241m.\u001b[39mvalues())))\n\u001b[1;32m 172\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mall\u001b[39m(\u001b[38;5;28mlen\u001b[39m(v) \u001b[38;5;241m==\u001b[39m bs \u001b[38;5;28;01mfor\u001b[39;00m v \u001b[38;5;129;01min\u001b[39;00m bo\u001b[38;5;241m.\u001b[39mvalues()), (\n\u001b[1;32m 173\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmust return Dict[str,Tensor] and all tensors with same batch size a first dimension\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 174\u001b[0m )\n\u001b[0;32m--> 175\u001b[0m \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 176\u001b[0m writer\u001b[38;5;241m.\u001b[39mfinalize()\n\u001b[1;32m 177\u001b[0m writer\u001b[38;5;241m.\u001b[39mclose()\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:605\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[0;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[1;32m 603\u001b[0m col_try_type \u001b[38;5;241m=\u001b[39m try_features[col] \u001b[38;5;28;01mif\u001b[39;00m try_features \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m col \u001b[38;5;129;01min\u001b[39;00m try_features \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 604\u001b[0m typed_sequence \u001b[38;5;241m=\u001b[39m OptimizedTypedSequence(col_values, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mcol_type, try_type\u001b[38;5;241m=\u001b[39mcol_try_type, col\u001b[38;5;241m=\u001b[39mcol)\n\u001b[0;32m--> 605\u001b[0m arrays\u001b[38;5;241m.\u001b[39mappend(\u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtyped_sequence\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m 606\u001b[0m inferred_features[col] \u001b[38;5;241m=\u001b[39m typed_sequence\u001b[38;5;241m.\u001b[39mget_inferred_type()\n\u001b[1;32m 607\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:252\u001b[0m, in \u001b[0;36mpyarrow.lib.array\u001b[0;34m()\u001b[0m\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:114\u001b[0m, in \u001b[0;36mpyarrow.lib._handle_arrow_array_protocol\u001b[0;34m()\u001b[0m\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:219\u001b[0m, in \u001b[0;36mTypedSequence.__arrow_array__\u001b[0;34m(self, type)\u001b[0m\n\u001b[1;32m 217\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(pa_type, _ArrayXDExtensionType):\n\u001b[1;32m 218\u001b[0m storage \u001b[38;5;241m=\u001b[39m to_pyarrow_listarray(data, pa_type)\n\u001b[0;32m--> 219\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mExtensionArray\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_storage\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstorage\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 221\u001b[0m \u001b[38;5;66;03m# efficient np array to pyarrow array\u001b[39;00m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, np\u001b[38;5;241m.\u001b[39mndarray):\n", "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:4354\u001b[0m, in \u001b[0;36mpyarrow.lib.ExtensionArray.from_storage\u001b[0;34m()\u001b[0m\n", "\u001b[0;31mTypeError\u001b[0m: Incompatible storage type list> for extension type extension>" ] } ], "source": [ "f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, \n", " schema=schema,\n", " features=Features.from_arrow_schema(schema),\n", " postprocess_result=float16_postprocess_result)\n", "f2\n", "ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n", "ds_a2.info" ] }, { "cell_type": "code", "execution_count": 44, "metadata": {}, "outputs": [], "source": [ "# %debug" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, postprocess_result=float8_postprocess_result)\n", "f2" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from datasets import Dataset\n", "# load\n", "ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n", "for c in ds_a2.column_names[1:]:\n", " print(c)\n", " ds_a2[c] = int8_to_float32(ds_a2[0:-1][c])\n", "# ds_a2 = int8_to_float32(ds_a)\n", "ds_a2.info" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "d = ds_a2[:][c]\n", "print(c)\n", "d.shape" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ds_a2.info" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ds_a[0:2]['logits'].shape" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 2 }