saprmarks / dictionary_learning Goto Github PK

Traceback (most recent call last):
  File "/home/joel/code/mamba-train/analyze_sae.py", line 180, in <module>
    run(args)
  File "/home/joel/code/mamba-train/analyze_sae.py", line 154, in run
    result = evaluate(model, submodule, ae, buffer)
  File "/home/joel/code/mamba-train/dictionary_learning/evaluation.py", line 82, in evaluate
    acts = next(activations).to(device)
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 103, in __next__
    self.refresh()
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 223, in refresh
    self._refresh_std()
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 156, in _refresh_std
    tokens = self.tokenized_batch()["input_ids"]
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 129, in tokenized_batch
    texts = self.text_batch(batch_size=batch_size)
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 123, in text_batch
    return [next(self.data) for _ in range(batch_size)]
  File "/home/joel/code/mamba-train/dictionary_learning/buffer.py", line 123, in <listcomp>
    return [next(self.data) for _ in range(batch_size)]
StopIteration

from nnsight import LanguageModel
from dictionary_learning import ActivationBuffer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate

model = LanguageModel("EleutherAI/pythia-70m-deduped")
submodule = model.gpt_neox.layers[1].mlp
activation_dim = 512
dictionary_size = 16 * activation_dim

data = iter(
    [
        "This is some example data",
        "In real life, for training a dictionary",
        "you would need much more data than this",
    ]
)
buffer = ActivationBuffer(
    data,
    model,
    submodule,
    out_feats=activation_dim,
    n_ctxs=3e4,
)

ae = trainSAE(
    buffer,
    activation_dim,
    dictionary_size,
    lr=3e-4,
    sparsity_penalty=1e-3,
)

buffer = ActivationBuffer(
    iter(["hello", "world"]),
    model,
    submodule,
    out_feats=activation_dim,
    n_ctxs=300,
)

print(evaluate(model, submodule, ae, buffer))

joel@simplex ~/c/mamba-sae (main)> python3.11 test.py
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
0it [00:00, ?it/s]
Traceback (most recent call last):
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 94, in text_batch
    return [
           ^
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 95, in <listcomp>
    next(self.data) for _ in range(batch_size)
    ^^^^^^^^^^^^^^^
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/joel/code/mamba-sae/test.py", line 46, in <module>
    print(evaluate(model, submodule, ae, buffer))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/joel/code/mamba-sae/dictionary_learning/evaluation.py", line 119, in evaluate
    acts = next(activations).to(device)
           ^^^^^^^^^^^^^^^^^
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 76, in __next__
    self.refresh()
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 164, in refresh
    self._refresh_std()
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 121, in _refresh_std
    with self.model.invoke(self.text_batch(), truncation=True, max_length=self.ctx_len) as invoker:
                           ^^^^^^^^^^^^^^^^^
  File "/Users/joel/code/mamba-sae/dictionary_learning/buffer.py", line 98, in text_batch
    raise StopIteration("End of data stream reached")
StopIteration: End of data stream reached

from nnsight import LanguageModel
from dictionary_learning import ActivationBuffer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate

model = LanguageModel("EleutherAI/pythia-70m-deduped")
submodule = model.gpt_neox.layers[1].mlp
activation_dim = 512
dictionary_size = 16 * activation_dim

data = iter(
    [
        "This is some example data",
        "In real life, for training a dictionary",
        "you would need much more data than this",
    ]
)
buffer = ActivationBuffer(
    data,
    model,
    submodule,
    out_feats=activation_dim,
    n_ctxs=3e4,
)

ae = trainSAE(
    buffer,
    activation_dim,
    dictionary_size,
    lr=3e-4,
    sparsity_penalty=1e-3,
)

(env) root@12702e9f3980:/workspace/mamba-sae# python3 test.py
Resolving data files: 100%|██████████████████████████████████████████| 30/30 [00:06<00:00,  4.89it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
0it [00:00, ?it/s]You're using a GPTNeoXTokenizerFast 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.
0it [00:17, ?it/s]
Traceback (most recent call last):
  File "/workspace/mamba-sae/test.py", line 33, in <module>
    ae = trainSAE(
  File "/workspace/mamba-sae/dictionary_learning/training.py", line 180, in trainSAE
    for step, acts in enumerate(tqdm(activations, total=steps)):
  File "/workspace/mamba-sae/env/lib/python3.10/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
  File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 76, in __next__
    self.refresh()
  File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 164, in refresh
    self._refresh_std()
  File "/workspace/mamba-sae/dictionary_learning/buffer.py", line 130, in _refresh_std
    self.activations = t.cat([self.activations, hidden_states.value.to(self.device)], dim=0)
RuntimeError: Tensors must have same number of dimensions: got 2 and 3