saprmarks / dictionary_learning

MIT License
142 stars 37 forks source link

`EmptyStream` isn't used #2

Closed joelburget closed 10 months ago

joelburget commented 10 months ago

Hi there, thanks for putting this library together! I noticed that you've defined EmptyStream and even handle it in ActivationBuffer.__next__, but it's never thrown. As an experiment I changed my dataset to have only a single point and this triggered an exception (which from my reading of the code it shouldn't have).

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
saprmarks commented 10 months ago

Thanks for pointing this out, and sorry it took so long to resolve. As of the last update, I think this should be fixed.