Skip to content

Restore tokenizer special-id patches in attribution tests - #96

Merged
hannamw merged 1 commit into
decoderesearch:mainfrom
Andrey170170:pr/test-tokenizer-state-cleanup
May 21, 2026
Merged

Restore tokenizer special-id patches in attribution tests#96
hannamw merged 1 commit into
decoderesearch:mainfrom
Andrey170170:pr/test-tokenizer-state-cleanup

Conversation

@Andrey170170

Copy link
Copy Markdown
Contributor

Summary

This PR prevents TransformerLens attribution tests from leaking tokenizer class state across tests.
The Gemma and Llama dummy-model tests temporarily patch tokenizer.all_special_ids so attribution treats token 0 as special. Previously, the tests assigned directly to the tokenizer class property without restoring the original value.
This change:

  • adds a small patch_tokenizer_special_ids(...) context manager
  • removes permanent all_special_ids mutation from dummy model loaders
  • wraps the relevant Gemma/Llama dummy attribution tests so the original tokenizer class property is restored after each test
  • reuses the helper in the Llama tests via the existing Gemma test imports

Why

This avoids leaking global tokenizer class state across tests and follows the safer restore pattern already used by the NNSight attribution tests.

Tests

Run:
python -m ruff check tests/test_attributions_gemma.py tests/test_attributions_llama.py
python -m ruff format --check tests/test_attributions_gemma.py tests/test_attributions_llama.py
python -m pytest tests/test_attributions_gemma.py::test_small_gemma_model tests/test_attributions_llama.py::test_small_llama_model

Results:

ruff check: passed
ruff format --check: passed
pytest: 2 passed

Not run:

  • full pytest suite
  • full pyright

@Andrey170170

Copy link
Copy Markdown
Contributor Author

CI note: the failing build check appears unrelated to this PR.

The failure is in the check types step (python -m pyright). Current main is failing the same step as well:
https://github.com/decoderesearch/circuit-tracer/actions/runs/26090319535

I reproduced locally that latest pyright (1.1.409) reports many torch export errors, while the currently declared lower-bound version passes:

  • uv run python -m pyright → fails with reportPrivateImportUsage on torch symbols
  • uv run --with pyright==1.1.403 python -m pyright → passes

Likely cause is the unpinned dev dependency pyright>=1.1.403 resolving to a newer pyright version.

@hannamw

hannamw commented May 20, 2026

Copy link
Copy Markdown
Collaborator

Yes, that's the conclusion I reached as well - the issue seems to have appeared in a pyright past 1.1.406. The torch issues appear to be noise, and can been suppressed, but some of the issues are legitimate. I've kept pyright unpinned, but suppressed the torch issues, while fixing the other ones, in #99 . I think it should be possible to merge this now!

@Andrey170170
Andrey170170 force-pushed the pr/test-tokenizer-state-cleanup branch from a5daa00 to 367dd3c Compare May 21, 2026 00:03

@hannamw hannamw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good, and seems like a sensible change. Thanks for contributing!

@hannamw
hannamw merged commit 80ca62b into decoderesearch:main May 21, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants