mirror of
https://github.com/wassname/activation_store.git
synced 2026-07-25 13:00:35 +08:00
34 KiB
34 KiB
In [1]:
%reload_ext autoreload
%autoreload 2In [2]:
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import Dataset
import torch
from activation_store.collect import activation_store
In [3]:
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="auto",
device_map="auto",
attn_implementation="eager", # flex_attention flash_attention_2 sdpa eager
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
In [4]:
N = 10
max_length = 128
imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)
def proc(row):
messages = [
{"role":"user", "content": row['prompt'] },
{"role":"assistant", "content": row['chosen'] }
]
return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_dict=True, max_length=max_length)
ds2 = imdb.map(proc).with_format("torch")
new_cols = set(ds2.column_names) - set(imdb.column_names)
ds2 = ds2.select_columns(new_cols)
ds2Out [4]:
Dataset({
features: ['attention_mask', 'input_ids'],
num_rows: 10
})In [5]:
from torch.utils.data import DataLoader
def collate_fn(examples):
# Pad the batch to max length within this batch
return tokenizer.pad(
examples,
padding=True,
return_tensors="pt",
)
ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)
print(ds)
<torch.utils.data.dataloader.DataLoader object at 0x7089fb69f6e0>
In [6]:
# choose layers to cache
layers = [k for k,v in model.named_modules() if 'mlp.down_proj' in k]
layersOut [6]:
['model.layers.0.mlp.down_proj', 'model.layers.1.mlp.down_proj', 'model.layers.2.mlp.down_proj', 'model.layers.3.mlp.down_proj', 'model.layers.4.mlp.down_proj', 'model.layers.5.mlp.down_proj', 'model.layers.6.mlp.down_proj', 'model.layers.7.mlp.down_proj', 'model.layers.8.mlp.down_proj', 'model.layers.9.mlp.down_proj', 'model.layers.10.mlp.down_proj', 'model.layers.11.mlp.down_proj', 'model.layers.12.mlp.down_proj', 'model.layers.13.mlp.down_proj', 'model.layers.14.mlp.down_proj', 'model.layers.15.mlp.down_proj', 'model.layers.16.mlp.down_proj', 'model.layers.17.mlp.down_proj', 'model.layers.18.mlp.down_proj', 'model.layers.19.mlp.down_proj', 'model.layers.20.mlp.down_proj', 'model.layers.21.mlp.down_proj', 'model.layers.22.mlp.down_proj', 'model.layers.23.mlp.down_proj']
In [7]:
f = activation_store(ds, model, layers=layers, writer_batch_size=10)
fOut [7]:
[32m2025-02-16 09:16:55.292[0m | [1mINFO [0m | [36mactivation_store.collect[0m:[36mactivation_store[0m:[36m122[0m - [1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet[0m
collecting activations: 0%| | 0/3 [00:00<?, ?it/s]
You're using a Qwen2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
PosixPath('/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet')In [ ]:
In [ ]:
In [8]:
# load
ds_a = Dataset.from_parquet(str(f)).with_format("torch")
ds_aOut [8]:
Generating train split: 0 examples [00:00, ? examples/s]
Dataset({
features: ['attention_mask', 'act-model.layers.0.mlp.down_proj', 'act-model.layers.1.mlp.down_proj', 'act-model.layers.2.mlp.down_proj', 'act-model.layers.3.mlp.down_proj', 'act-model.layers.4.mlp.down_proj', 'act-model.layers.5.mlp.down_proj', 'act-model.layers.6.mlp.down_proj', 'act-model.layers.7.mlp.down_proj', 'act-model.layers.8.mlp.down_proj', 'act-model.layers.9.mlp.down_proj', 'act-model.layers.10.mlp.down_proj', 'act-model.layers.11.mlp.down_proj', 'act-model.layers.12.mlp.down_proj', 'act-model.layers.13.mlp.down_proj', 'act-model.layers.14.mlp.down_proj', 'act-model.layers.15.mlp.down_proj', 'act-model.layers.16.mlp.down_proj', 'act-model.layers.17.mlp.down_proj', 'act-model.layers.18.mlp.down_proj', 'act-model.layers.19.mlp.down_proj', 'act-model.layers.20.mlp.down_proj', 'act-model.layers.21.mlp.down_proj', 'act-model.layers.22.mlp.down_proj', 'act-model.layers.23.mlp.down_proj', 'logits', 'hidden_states'],
num_rows: 10
})In [10]:
ds_a.infoOut [10]:
DatasetInfo(description='', citation='', homepage='', license='', features={'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None), 'act-model.layers.0.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.1.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.2.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.3.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.4.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.5.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.6.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.7.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.8.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.9.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.10.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.11.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.12.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.13.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.14.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.15.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.16.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.17.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.18.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.19.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.20.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.21.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.22.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.23.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'logits': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'hidden_states': Sequence(feature=Sequence(feature=Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None), length=-1, id=None), length=-1, id=None)}, post_processed=None, supervised_keys=None, builder_name='parquet', dataset_name='parquet', config_name='default', version=0.0.0, splits={'train': SplitInfo(name='train', num_bytes=1391398926, num_examples=10, shard_lengths=[4, 6], dataset_name='parquet')}, download_checksums={'/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet': {'num_bytes': 1363203837, 'checksum': None}}, download_size=1363203837, post_processing_size=None, dataset_size=1391398926, size_in_bytes=2754602763)In [11]:
ds_a[0:2]['hidden_states'].shapeOut [11]:
torch.Size([2, 25, 453, 896])
In [12]:
ds_a[0:2]['act-model.layers.0.mlp.down_proj'].shapeOut [12]:
torch.Size([2, 453, 896])
In [9]:
1/0[0;31m---------------------------------------------------------------------------[0m [0;31mZeroDivisionError[0m Traceback (most recent call last) Cell [0;32mIn[9], line 1[0m [0;32m----> 1[0m [38;5;241;43m1[39;49m[38;5;241;43m/[39;49m[38;5;241;43m0[39;49m [0;31mZeroDivisionError[0m: division by zero
In [32]:
def float_to_int16(x: torch.Tensor) -> torch.Tensor:
"""Converts a floating point tensor to float16, then reinterprets as int16."""
downcast = x.type(torch.float16)
# if not downcast.isfinite().all():
# raise ValueError("Cannot convert to 16 bit: values are not finite")
return downcast.view(torch.int16)
def int8_to_float32(x: torch.Tensor) -> torch.Tensor:
"""Converts int16 to float16, then reinterprets as float32."""
return x.view(torch.float16).type(torch.float32)
x = torch.randn(2, 3, 4)
x2 = float_to_int16(x)
x3 = int8_to_float32(x2)
assert torch.isfinite(x3).all()
assert torch.allclose(x, x3, rtol=1e-1)
d = ((x-x3)/x).abs().mean()
print(f'lost {d:.2%}')lost 0.01%
In [41]:
from activation_store.collect import default_postprocess_result
from datasets.features.features import cast_to_python_objects
# o = cast_to_python_objects(o, only_1d_for_numpy=True, optimize_list_casting=False)
def float16_postprocess_result(
input, trace, output, model
):
o = default_postprocess_result(input, trace, output, model)
# o = cast_to_python_objects(o, only_1d_for_numpy=False, optimize_list_casting=False)
for k, v in o.items():
if k=='attention_mask':
o[k] = v.to(torch.int8)
if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
print(k, v.dtype, v.shape, 'to int16')
o[k] = float_to_int16(v)
else:
print('no conv', k, type(v))
# o = {k: float_to_int8(v) if isinstance(v, torch.Tensor) else v
# for k, v in o.items()}
return oIn [42]:
from datasets.arrow_writer import OptimizedTypedSequence, _ArrayXDExtensionType
from datasets.features.features import Features, Array2D, Array3D, Array4D, Array5D
# manually build features
optimized_int_type_by_col = {
"attention_mask": "int8", # binary tensor
"special_tokens_mask": "int8",
"input_ids": "int32", # typical vocab size: 0-50k (max ~500k, never > 1M)
"token_type_ids": "int8", # binary mask; some (XLNetModel) use an additional token represented by a 2
}
def build_schema(d):
inferred_features = Features()
cols = d.keys()
for col in cols:
x = d[col]
if col in optimized_int_type_by_col:
dtype = optimized_int_type_by_col[col]
typed_sequence = OptimizedTypedSequence(x, col=col)
inferred_features[col] = typed_sequence.get_inferred_type()
else:
if x.ndim == 1:
inferred_features[col] = OptimizedTypedSequence(x, col=col)
inferred_features[col] = typed_sequence.get_inferred_type()
shape=(-1,)+x.shape[1:]
dtype = 'int16' if x.dtype == torch.float32 else x.dtype
if x.ndim == 2:
cls = Array2D
elif x.ndim == 3:
cls = Array3D
elif x.ndim == 4:
cls = Array4D
elif x.ndim == 5:
cls = Array5D
else:
raise ValueError(f"Unsupported number of dimensions: {x.ndim}")
inferred_features[col] = cls(dtype=dtype, shape=shape)
return inferred_features.arrow_schema
# Features.from_arrow_schema(schema)
d = ds_a[0:2]
schema = build_schema(d)
schema
Features.from_arrow_schema(schema)Out [42]:
{'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),
'act-model.layers.0.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.1.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.2.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.3.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.4.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.5.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.6.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.7.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.8.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.9.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.10.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.11.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.12.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.13.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.14.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.15.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.16.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.17.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.18.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.19.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.20.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.21.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.22.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'act-model.layers.23.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int16', id=None),
'logits': Array3D(shape=(-1, 453, 151936), dtype='int16', id=None),
'hidden_states': Array4D(shape=(-1, 25, 453, 896), dtype='int16', id=None)}In [43]:
f2 = activation_store(ds, model, layers=layers, writer_batch_size=10,
schema=schema,
features=Features.from_arrow_schema(schema),
postprocess_result=float16_postprocess_result)
f2
ds_a2 = Dataset.from_parquet(str(f2)).with_format("torch")
ds_a2.info[32m2025-02-16 09:35:12.432[0m | [1mINFO [0m | [36mactivation_store.collect[0m:[36mactivation_store[0m:[36m155[0m - [1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__c6184d05bf03be61.parquet[0m
collecting activations: 0%| | 0/3 [00:00<?, ?it/s]
no conv attention_mask <class 'torch.Tensor'> act-model.layers.0.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.1.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.2.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.3.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.4.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.5.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.6.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.7.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.8.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.9.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.10.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.11.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.12.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.13.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.14.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.15.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.16.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.17.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.18.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.19.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.20.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.21.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.22.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 act-model.layers.23.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int16 logits torch.float16 torch.Size([4, 453, 151936]) to int16 hidden_states torch.float32 torch.Size([4, 25, 453, 896]) to int16
[0;31m---------------------------------------------------------------------------[0m [0;31mTypeError[0m Traceback (most recent call last) Cell [0;32mIn[43], line 1[0m [0;32m----> 1[0m f2 [38;5;241m=[39m [43mactivation_store[49m[43m([49m[43mds[49m[43m,[49m[43m [49m[43mmodel[49m[43m,[49m[43m [49m[43mlayers[49m[38;5;241;43m=[39;49m[43mlayers[49m[43m,[49m[43m [49m[43mwriter_batch_size[49m[38;5;241;43m=[39;49m[38;5;241;43m10[39;49m[43m,[49m[43m [49m [1;32m 2[0m [43m [49m[43mschema[49m[38;5;241;43m=[39;49m[43mschema[49m[43m,[49m [1;32m 3[0m [43m [49m[43mfeatures[49m[38;5;241;43m=[39;49m[43mFeatures[49m[38;5;241;43m.[39;49m[43mfrom_arrow_schema[49m[43m([49m[43mschema[49m[43m)[49m[43m,[49m [1;32m 4[0m [43m [49m[43mpostprocess_result[49m[38;5;241;43m=[39;49m[43mfloat16_postprocess_result[49m[43m)[49m [1;32m 5[0m f2 [1;32m 6[0m ds_a2 [38;5;241m=[39m Dataset[38;5;241m.[39mfrom_parquet([38;5;28mstr[39m(f2))[38;5;241m.[39mwith_format([38;5;124m"[39m[38;5;124mtorch[39m[38;5;124m"[39m) File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/activation_store/collect.py:175[0m, in [0;36mactivation_store[0;34m(loader, model, dataset_name, layers, dataset_dir, writer_batch_size, postprocess_result, features, schema)[0m [1;32m 171[0m bs [38;5;241m=[39m [38;5;28mlen[39m([38;5;28mnext[39m([38;5;28miter[39m(bo[38;5;241m.[39mvalues()))) [1;32m 172[0m [38;5;28;01massert[39;00m [38;5;28mall[39m([38;5;28mlen[39m(v) [38;5;241m==[39m bs [38;5;28;01mfor[39;00m v [38;5;129;01min[39;00m bo[38;5;241m.[39mvalues()), ( [1;32m 173[0m [38;5;124m"[39m[38;5;124mmust return Dict[str,Tensor] and all tensors with same batch size a first dimension[39m[38;5;124m"[39m [1;32m 174[0m ) [0;32m--> 175[0m [43mwriter[49m[38;5;241;43m.[39;49m[43mwrite_batch[49m[43m([49m[43mbo[49m[43m)[49m [1;32m 176[0m writer[38;5;241m.[39mfinalize() [1;32m 177[0m writer[38;5;241m.[39mclose() File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:605[0m, in [0;36mArrowWriter.write_batch[0;34m(self, batch_examples, writer_batch_size)[0m [1;32m 603[0m col_try_type [38;5;241m=[39m try_features[col] [38;5;28;01mif[39;00m try_features [38;5;129;01mis[39;00m [38;5;129;01mnot[39;00m [38;5;28;01mNone[39;00m [38;5;129;01mand[39;00m col [38;5;129;01min[39;00m try_features [38;5;28;01melse[39;00m [38;5;28;01mNone[39;00m [1;32m 604[0m typed_sequence [38;5;241m=[39m OptimizedTypedSequence(col_values, [38;5;28mtype[39m[38;5;241m=[39mcol_type, try_type[38;5;241m=[39mcol_try_type, col[38;5;241m=[39mcol) [0;32m--> 605[0m arrays[38;5;241m.[39mappend([43mpa[49m[38;5;241;43m.[39;49m[43marray[49m[43m([49m[43mtyped_sequence[49m[43m)[49m) [1;32m 606[0m inferred_features[col] [38;5;241m=[39m typed_sequence[38;5;241m.[39mget_inferred_type() [1;32m 607[0m schema [38;5;241m=[39m inferred_features[38;5;241m.[39marrow_schema [38;5;28;01mif[39;00m [38;5;28mself[39m[38;5;241m.[39mpa_writer [38;5;129;01mis[39;00m [38;5;28;01mNone[39;00m [38;5;28;01melse[39;00m [38;5;28mself[39m[38;5;241m.[39mschema File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:252[0m, in [0;36mpyarrow.lib.array[0;34m()[0m File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:114[0m, in [0;36mpyarrow.lib._handle_arrow_array_protocol[0;34m()[0m File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:219[0m, in [0;36mTypedSequence.__arrow_array__[0;34m(self, type)[0m [1;32m 217[0m [38;5;28;01mif[39;00m [38;5;28misinstance[39m(pa_type, _ArrayXDExtensionType): [1;32m 218[0m storage [38;5;241m=[39m to_pyarrow_listarray(data, pa_type) [0;32m--> 219[0m [38;5;28;01mreturn[39;00m [43mpa[49m[38;5;241;43m.[39;49m[43mExtensionArray[49m[38;5;241;43m.[39;49m[43mfrom_storage[49m[43m([49m[43mpa_type[49m[43m,[49m[43m [49m[43mstorage[49m[43m)[49m [1;32m 221[0m [38;5;66;03m# efficient np array to pyarrow array[39;00m [1;32m 222[0m [38;5;28;01mif[39;00m [38;5;28misinstance[39m(data, np[38;5;241m.[39mndarray): File [0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:4354[0m, in [0;36mpyarrow.lib.ExtensionArray.from_storage[0;34m()[0m [0;31mTypeError[0m: Incompatible storage type list<item: list<item: int16>> for extension type extension<datasets.features.features.Array3DExtensionType<Array3DExtensionType>>
In [44]:
# %debugIn [ ]:
f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, postprocess_result=float8_postprocess_result)
f2In [ ]:
from datasets import Dataset
# load
ds_a2 = Dataset.from_parquet(str(f2)).with_format("torch")
for c in ds_a2.column_names[1:]:
print(c)
ds_a2[c] = int8_to_float32(ds_a2[0:-1][c])
# ds_a2 = int8_to_float32(ds_a)
ds_a2.infoIn [ ]:
d = ds_a2[:][c]
print(c)
d.shapeIn [ ]:
ds_a2.infoIn [ ]:
ds_a[0:2]['logits'].shape