diff --git a/docs/source/content/natural_language_autoencoders.md b/docs/source/content/natural_language_autoencoders.md
new file mode 100644
index 000000000..47bda6119
--- /dev/null
+++ b/docs/source/content/natural_language_autoencoders.md
@@ -0,0 +1,223 @@
+# Natural language autoencoder inference
+
+## Availability and validation boundary
+
+`NaturalLanguageAutoencoder` is available from `transformer_lens.tools.analysis`.
+The supported registry name is `qwen2.5-7b-instruct-L20`, with immutable target,
+AV, and AR revisions. The released pair passed direct-HF AV-logit and AR-vector
+checks on CPU in fp32 with eager attention, with zero maximum and mean absolute
+error for two inputs per path. The sequential released-model round trip also
+completed. These finite test inputs do not guarantee explanation quality.
+
+Only raw fp32/eager single-device execution is implemented. CPU is the validated
+configuration; CUDA/MPS, quantization, arbitrary dispatch, and other released
+pairs are not advertised as verified. Local synthetic tests complement, but do
+not replace, slow released-artifact checks.
+
+## What the round trip does
+
+An Activation Verbalizer (AV) translates one residual-stream activation into
+explanation text. An Activation Reconstructor (AR) translates that text back into
+a vector using a truncated backbone and its separate trained linear value head.
+TransformerBridge hooks provide local execution without an inference server.
+Training, explanation-edit steering, corpus execution, and dataset-level FVE are
+not part of this inference interface.
+
+The candidate Qwen artifacts expect the raw output of zero-based block 20:
+`blocks.20.hook_out`, corresponding to HF `hidden_states[21]`. Their hidden size
+is 3584. This block is artifact-fixed, not a freely interchangeable layer.
+Compatibility weight processing changes the expected residual basis and is
+rejected. Merely matching architecture and dimensions does not establish that
+an activation belongs to the training distribution.
+
+The AV sidecar records its mandatory injection L2 norm (150 for the inspected
+Qwen candidate). The vector replaces one reserved token's post-embedding output
+in the canonical chat template. The tokenizer encoding, placeholder count, and
+neighbor IDs are checked rather than guessed. Cached execution injects on
+prefill only; uncached execution reapplies at the original position on each
+prefix recomputation. One-token prefill is valid. On ordinary fp32 inputs,
+injection normalization matches the released client's fp32 norm and division
+order. A scale-safe fallback handles norm underflow/overflow outside that input
+domain; this fallback is not claimed as reference-client equivalence.
+
+AR takes a raw prompt, not a chat template. For the inspected Qwen candidate it
+is `Summary of the following text: {explanation} `.
+Special-token and suffix conventions must match the sidecar. AR reads the raw
+final-token output of block 20 before final normalization, then applies the
+released bias-free linear head. The private loader validates the omitted final
+norm/LM-head weights and runs only the backbone under Bridge hooks. The omitted
+LM head is never used to reconstruct a vector.
+
+## Explanation and scoring contracts
+
+An explanation result preserves raw generated continuation text and continuation
+token IDs separately from canonical inner explanation content. Exactly one
+complete `...` region is required. Missing, multiple,
+nested, malformed, incomplete, and empty regions yield `text=None` with an
+explicit parse status; raw text is not substituted as a success.
+
+Reconstruction accepts nonempty untagged explanation content. Explicit or malformed
+explanation wrappers are rejected. It preserves authored interior whitespace and
+does not claim semantic detection of arbitrary untagged full responses. Do not
+pass raw failed generation to AR automatically.
+
+Scoring normalizes both finite nonzero vectors independently and reports cosine
+and normalized per-coordinate MSE:
+
+```text
+normalized_mse = 2 * (1 - cosine)
+```
+
+Equivalently, normalize each vector to L2 norm `sqrt(d)` before taking the
+per-coordinate squared-error mean. There is no additional dimension multiplier.
+CPU float64 accumulation with scale-safe norms handles finite extreme inputs;
+only roundoff-sized cosine overshoot is clamped. Unsupported AR normalization
+modes, including null/raw-MSE metadata, are rejected.
+
+These scores are directional diagnostics, not raw vector MSE, dataset FVE, or
+proof of human-readable faithfulness. A low reconstruction error can coexist
+with misleading text. Do not infer causal validity from fluent explanations.
+
+## Execution examples
+
+The registry pins the target and both role checkpoints. Metadata and tokenizer
+contracts are checked before full role-model loading. Use the context manager to
+release owned AV/AR resources, including on exceptions.
+
+### Convenience workflow
+
+This retains the caller-owned target while AV and AR execute and may require
+multiple models to fit in memory:
+
+```python
+from transformer_lens.tools.analysis import NaturalLanguageAutoencoder
+
+with NaturalLanguageAutoencoder.from_pretrained(
+ "qwen2.5-7b-instruct-L20", device="cpu"
+) as nla:
+ target = nla.load_target()
+ try:
+ activation = nla.extract(target, prompt="The auditor asked about reasoning.")
+ explanation = nla.verbalize(activation, max_new_tokens=512)
+ if explanation.text is None:
+ raise ValueError(f"Explanation parse failed: {explanation.parse_status}")
+ reconstruction = nla.reconstruct(explanation.text)
+ score = nla.score(activation, reconstruction)
+ print(explanation.text, score.cosine, score.normalized_mse)
+ finally:
+ del target
+```
+
+### Sequential low-memory workflow
+
+Extraction and reconstruction return detached owned vectors on CPU by default.
+Release the target before loading AV, and unload AV before AR:
+
+```python
+from transformer_lens.tools.analysis import NaturalLanguageAutoencoder
+
+with NaturalLanguageAutoencoder.from_pretrained(
+ "qwen2.5-7b-instruct-L20", device="cpu"
+) as nla:
+ target = nla.load_target()
+ try:
+ activation = nla.extract(
+ target, prompt="The auditor asked about reasoning.", output_device="cpu"
+ )
+ finally:
+ del target
+ try:
+ explanation = nla.verbalize(activation)
+ finally:
+ nla.unload("av")
+ if explanation.text is None:
+ raise ValueError(f"Explanation parse failed: {explanation.parse_status}")
+ try:
+ reconstruction = nla.reconstruct(explanation.text)
+ finally:
+ nla.unload("ar")
+ score = nla.score(activation, reconstruction.vector)
+```
+
+The context owns AV/AR only, not returned targets or caller-held tensors. Other
+references can keep a target alive after `del target`. Unload releases owned
+references, not all allocator-reserved memory or externally held allocations.
+An unloaded role can reload lazily; `close()` permanently closes the instance.
+Concurrent operations, including unload during inference, are rejected.
+
+## Inputs, provenance, and bounded generation
+
+- One unpadded string or explicit integer token sequence is supported. Negative
+ positions resolve against the original prompt. Oversized inputs raise rather
+ than silently truncate. Target strings are not automatically chat-templated.
+- `max_new_tokens=512`, temperature 1, and KV caching are the finite AV defaults.
+ Greedy decoding uses `do_sample=False`; sampling filters require sampling.
+ The prompt plus generation budget must fit the context even if EOS could stop
+ earlier. The default length is a convenience policy, not a paper replication.
+- `source_verified` means recorded pinned-loading provenance, not continuously
+ verified live weights. In-place weight edits can invalidate artifact assumptions.
+ Unknown target origins fail by default; explicit `allow_unverified_target=True`
+ warns and marks results unverified without bypassing known incompatibilities.
+- Bare vectors have unknown source provenance. Structured extraction results record
+ identity, layer, position, basis, and tokenization policy without retaining models.
+- Cache/offline options reuse HF snapshots. Metadata/tokenizers are checked before
+ full AV/AR weight loading. Remote-code execution is disabled. Import itself does
+ not initiate NLA downloads.
+
+## Resource requirements and verification
+
+The workflow involves three separate 7B-class checkpoints. A pinned-artifact file
+inventory measured 38.5 GiB of missing safetensor files for target, AV, and AR on
+an empty-weight cache, excluding tokenizer assets and temporary/runtime overhead.
+The slow checker adds a 2 GiB free-space reserve. Sequential execution reduces
+peak residency, not retained checkpoint storage. A cached serial CPU/fp32 run of
+all three released tests, including public `from_pretrained()` and the round trip,
+completed in 281.10 seconds (284.33 seconds wall time including process overhead).
+The child process maximum resident set size reported by macOS `time -l` was
+36,417,159,168 bytes, approximately 33.92 GiB, on a machine with 48 GiB RAM.
+This measures the complete test process, not a per-role minimum or accelerator
+VRAM requirement. Downloads are excluded from this cached-run timing; no CUDA/MPS
+fit or timing claim is made.
+
+Run validation serially on a machine with sufficient storage and memory. Source
+`.env` before Hub access. Set `HF_HOME` to a larger cache volume before starting
+Python if needed:
+
+```bash
+set -a; source .env; set +a
+uv run pytest tests/integration/test_nla_oracle_parity.py -m slow -s
+```
+
+These tests compare AV prefill logits and AR vectors to pinned direct-HF reference
+paths and exercise a complete sequential round trip. They do not execute downloaded
+reference Python code. Exact fp32 equality is the initial strict criterion; any
+nonzero drift requires investigation and measured tolerance justification.
+Current standard test targets deselect `slow`, so their passing does not imply
+released-artifact parity. Normal-CI synthetic/cached-model tests remain necessary
+for maintaining hook execution and lifecycle behavior.
+
+## Package-resource check
+
+The registry must be present in a built wheel even while it is empty. A local
+build/check uses no model downloads:
+
+```bash
+uv build
+uv run pytest tests/integration/test_nla_packaging.py
+```
+
+The packaging test builds a wheel with a fresh uv cache, allowing build dependency
+resolution, and reads its resources. A separate isolated subprocess imports the
+actual extracted-wheel package and analysis initializers with network calls
+forbidden. Build dependencies may require network access; the package import
+check does not. These checks verify package/resource delivery, not model
+compatibility.
+
+## References
+
+- [Natural Language Autoencoders paper](https://transformer-circuits.pub/2026/nla/)
+- [Authors' reference repository](https://github.com/kitft/natural_language_autoencoders)
+
+The test oracle records its pinned reference revision and artifact identities.
+The integration consumes trained artifacts and does not import the authors'
+training or serving stack.
diff --git a/docs/source/index.md b/docs/source/index.md
index e81b4cf4c..fae4cc375 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -62,6 +62,7 @@ content/compatibility_mode
content/ssm_interpretability
content/projection_kernel
content/jacobian_lens_fitting
+content/natural_language_autoencoders
generated/demos/Jacobian_Lens_Decomposition_Demo
generated/demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo
content/backward_lens
diff --git a/pyproject.toml b/pyproject.toml
index fc7142c95..20a1c923f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,8 @@
"pandas>=1.1.5",
"protobuf>=3.20.0",
"rich>=12.6.0",
+ "PyYAML>=6.0",
+ "safetensors>=0.4.3",
"sentencepiece",
"torch>=2.6",
"tqdm>=4.64.1",
@@ -83,6 +85,7 @@
"pytest-rerunfailures>=12.0",
"pytest-xdist>=3.8.0",
"pytest>=7.2.0",
+ "types-PyYAML>=6.0",
]
docs=[
"furo>=2023.3.27,<2024.0.0",
diff --git a/tests/integration/test_nla_bridge_execution.py b/tests/integration/test_nla_bridge_execution.py
new file mode 100644
index 000000000..b592b2010
--- /dev/null
+++ b/tests/integration/test_nla_bridge_execution.py
@@ -0,0 +1,332 @@
+"""Real cached-model checks of private NLA extraction mechanics."""
+
+import gc
+import weakref
+
+import pytest
+import torch
+
+from transformer_lens.tools.analysis._nla_artifacts import (
+ ArtifactRef,
+ NLAMetadata,
+ NLAPairConfig,
+)
+from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ NLAActivation,
+ NLASamplingConfig,
+ SourceProvenance,
+)
+from transformer_lens.tools.analysis._nla_target import (
+ extract_activation,
+ load_target,
+ validate_target_structure,
+)
+from transformer_lens.tools.analysis._nla_verbalizer import generate_explanation
+
+
+@pytest.fixture(scope="module")
+def target_pair():
+ from huggingface_hub import model_info
+
+ repo = "hf-internal-testing/tiny-random-gpt2"
+ revision = model_info(repo).sha
+ assert revision is not None
+ return NLAPairConfig(
+ ArtifactRef(repo, revision),
+ ArtifactRef("example/av", "a" * 40),
+ ArtifactRef("example/ar", "b" * 40),
+ "gpt2",
+ 32,
+ 1,
+ 5,
+ )
+
+
+@pytest.fixture(scope="module")
+def target(target_pair):
+ return load_target(target_pair, device="cpu")
+
+
+def test_real_extraction_matches_block_cache(target, target_pair):
+ tokens = torch.tensor([[1, 2, 3]])
+ result = extract_activation(target, target_pair, "mechanical-fixture", tokens, position=-1)
+ _, cache = target.run_with_cache(tokens, names_filter="blocks.1.hook_out", return_type=None)
+ torch.testing.assert_close(result.vector, cache["blocks.1.hook_out"][0, -1])
+ assert result.vector.device.type == "cpu" and not result.vector.requires_grad
+ assert result.provenance.source_verified and result.provenance.position == 2
+ assert result.vector.data_ptr() != cache["blocks.1.hook_out"].data_ptr()
+
+
+@pytest.mark.parametrize("position", [-4, 3, True])
+def test_invalid_position_does_not_run_forward(target, target_pair, position):
+ with pytest.raises(ValueError, match="position"):
+ extract_activation(
+ target, target_pair, "mechanical-fixture", torch.tensor([1, 2, 3]), position=position
+ )
+
+
+def test_string_input_and_repeated_extraction(target, target_pair):
+ a = extract_activation(
+ target, target_pair, "mechanical-fixture", "Hello world", prepend_bos=False
+ )
+ b = extract_activation(
+ target, target_pair, "mechanical-fixture", "Hello world", prepend_bos=False
+ )
+ torch.testing.assert_close(a.vector, b.vector)
+ assert "prepend_bos=False" in a.provenance.tokenization_policy
+
+
+def test_result_does_not_keep_target_alive(target_pair):
+ model = load_target(target_pair, device="cpu")
+ ref = weakref.ref(model)
+ result = extract_activation(model, target_pair, "mechanical-fixture", torch.tensor([1, 2]))
+ del model
+ gc.collect()
+ assert ref() is None
+ assert result.vector.shape == (32,)
+
+
+def test_processed_basis_rejected(target_pair):
+ model = load_target(target_pair, device="cpu")
+ model.compatibility_mode = True
+ with pytest.raises(ValueError, match="raw"):
+ validate_target_structure(model, target_pair)
+
+
+@pytest.mark.parametrize(
+ "attribute,value,message",
+ [
+ ("_weights_processed", True, "raw"),
+ ("training", True, "eval"),
+ ],
+)
+def test_mutated_target_guards(target_pair, attribute, value, message):
+ model = load_target(target_pair, device="cpu")
+ setattr(model, attribute, value)
+ with pytest.raises(ValueError, match=message):
+ validate_target_structure(model, target_pair)
+
+
+def test_precision_guard(target_pair):
+ model = load_target(target_pair, device="cpu")
+ model.original_model.to(dtype=torch.float16)
+ with pytest.raises(ValueError, match="float32"):
+ validate_target_structure(model, target_pair)
+
+
+def test_caller_hooks_preserved(target, target_pair):
+ calls = []
+
+ def observe(value, hook):
+ calls.append(hook.name)
+ return value
+
+ with target.hooks(fwd_hooks=[("blocks.1.hook_out", observe)]):
+ extract_activation(target, target_pair, "mechanical-fixture", torch.tensor([1, 2]))
+ target(torch.tensor([[1, 2]]))
+ assert calls == ["blocks.1.hook_out", "blocks.1.hook_out"]
+
+
+@pytest.fixture
+def actor_metadata():
+ return NLAMetadata(
+ "av",
+ 32,
+ 1,
+ "marker",
+ 2,
+ "{injection_char}",
+ "{explanation}",
+ 32**0.5,
+ 5.0,
+ (1, 2, 3),
+ 1,
+ None,
+ )
+
+
+@pytest.mark.parametrize("cached", [True, False])
+def test_actor_generation_hooks_and_continuation(target, actor_metadata, cached):
+ observations = []
+ activation = NLAActivation(
+ torch.ones(32), ActivationProvenance("mechanical-fixture", SourceProvenance.UNKNOWN_VECTOR)
+ )
+
+ def observe(value, hook):
+ observations.append(value.detach().clone())
+ return value
+
+ with target.hooks(fwd_hooks=[("hook_embed", observe)]):
+ first = generate_explanation(
+ target,
+ actor_metadata,
+ activation,
+ NLASamplingConfig(max_new_tokens=3, do_sample=False, use_past_kv_cache=cached),
+ )
+ second = generate_explanation(
+ target,
+ actor_metadata,
+ activation,
+ NLASamplingConfig(max_new_tokens=3, do_sample=False, use_past_kv_cache=cached),
+ )
+ target(torch.tensor([[1, 2, 3]]))
+ assert first.token_ids == second.token_ids
+ assert 0 < len(first.token_ids) <= 3
+ assert first.raw_text == target.tokenizer.decode(
+ list(first.token_ids), skip_special_tokens=True
+ )
+ assert len(observations) == 7
+ assert [x.shape[1] for x in observations[:3]] == ([3, 1, 1] if cached else [3, 4, 5])
+
+
+def test_actor_exception_cleanup(target, actor_metadata):
+ activation = NLAActivation(
+ torch.ones(32), ActivationProvenance("mechanical-fixture", SourceProvenance.UNKNOWN_VECTOR)
+ )
+ calls = []
+
+ def fail(value, hook):
+ calls.append(1)
+ raise RuntimeError("controlled failure")
+
+ with target.hooks(fwd_hooks=[("blocks.1.hook_out", fail)]):
+ with pytest.raises(RuntimeError, match="controlled failure"):
+ generate_explanation(
+ target, actor_metadata, activation, NLASamplingConfig(max_new_tokens=1)
+ )
+ assert calls == [1]
+ assert not target.hook_dict["embed.hook_out"].fwd_hooks
+ target(torch.tensor([[1, 2, 3]]))
+
+
+def test_embedding_injection_logits_match_direct_hf(target):
+ from transformer_lens.tools.analysis._nla_verbalizer import embedding_overwrite
+
+ tokens = torch.tensor([[1, 2, 3]])
+ vector = torch.arange(1.0, 33.0)
+ replacement = vector.double() / vector.double().norm() * 5
+ with torch.no_grad():
+ embeddings = target.original_model.get_input_embeddings()(tokens).clone()
+ embeddings[0, 1] = replacement.float()
+ expected = target.original_model(inputs_embeds=embeddings).logits
+ with target.hooks(
+ fwd_hooks=[("hook_embed", embedding_overwrite(vector, 5.0, 1, 3, cached=True))]
+ ):
+ actual = target(tokens)
+ torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
+
+
+def test_actor_budget_rejected(target, actor_metadata):
+ activation = NLAActivation(
+ torch.ones(32), ActivationProvenance("mechanical-fixture", SourceProvenance.UNKNOWN_VECTOR)
+ )
+ with pytest.raises(ValueError, match="context"):
+ generate_explanation(
+ target, actor_metadata, activation, NLASamplingConfig(max_new_tokens=target.cfg.n_ctx)
+ )
+
+
+def test_reconstruction_uses_raw_block_not_final_norm(target_pair):
+ from dataclasses import replace
+
+ from transformer_lens.tools.analysis._nla_reconstructor import reconstruct_vector
+
+ pair = replace(target_pair, layer=4)
+ model = load_target(pair)
+ tokens = torch.tensor([[1, 2, 3]])
+ head = torch.nn.Linear(32, 32, bias=False)
+ head.weight.data.copy_(torch.eye(32))
+ head.eval()
+ metadata = NLAMetadata(
+ "ar",
+ 32,
+ 4,
+ "marker",
+ 2,
+ "{injection_char}",
+ "{explanation}",
+ 32**0.5,
+ None,
+ None,
+ None,
+ (3,),
+ )
+ _, cache = model.run_with_cache(tokens, names_filter="blocks.4.hook_out", return_type=None)
+ expected = cache["blocks.4.hook_out"][0, -1].detach().clone()
+
+ class ForbiddenOutput(torch.nn.Module):
+ def forward(self, *args, **kwargs):
+ raise AssertionError("final output module must not execute")
+
+ # The backbone is owned by this test; shared fixture models remain untouched.
+ model.original_model.transformer.ln_f = torch.nn.Identity()
+ model.original_model.lm_head = ForbiddenOutput()
+ result = reconstruct_vector(model, head, tokens, metadata, pair)
+ torch.testing.assert_close(result.vector, expected)
+ assert not result.vector.requires_grad
+
+
+def test_qwen_truncated_loader_with_local_checkpoint(tmp_path, target):
+ from safetensors.torch import save_file
+ from transformers import Qwen2Config, Qwen2ForCausalLM
+
+ from transformer_lens.tools.analysis._nla_reconstructor import (
+ load_critic_model,
+ reconstruct_vector,
+ )
+
+ config = Qwen2Config(
+ hidden_size=32,
+ intermediate_size=64,
+ num_hidden_layers=2,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ vocab_size=1000,
+ max_position_embeddings=32,
+ tie_word_embeddings=False,
+ )
+ original = Qwen2ForCausalLM(config).eval()
+ state = {
+ key: value
+ for key, value in original.state_dict().items()
+ if key not in {"model.norm.weight", "lm_head.weight"}
+ }
+ config.save_pretrained(tmp_path)
+ save_file(state, tmp_path / "model.safetensors")
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 32,
+ 1,
+ 3,
+ )
+ model = load_critic_model(tmp_path, pair, target.tokenizer, "cpu")
+ head = torch.nn.Linear(32, 32, bias=False).eval()
+ metadata = NLAMetadata(
+ "ar",
+ 32,
+ 1,
+ "marker",
+ 2,
+ "{injection_char}",
+ "{explanation}",
+ 32**0.5,
+ None,
+ None,
+ None,
+ (3,),
+ )
+ tokens = torch.tensor([[1, 2, 3]])
+ original.model.norm = torch.nn.Identity()
+ with torch.no_grad():
+ expected = head(original.model(tokens, use_cache=False).last_hidden_state[0, -1])
+ actual = reconstruct_vector(model, head, tokens, metadata, pair)
+ torch.testing.assert_close(actual.vector, expected, atol=1e-6, rtol=1e-5)
+
+
+def test_non_bridge_rejected(target_pair):
+ with pytest.raises(TypeError, match="TransformerBridge"):
+ validate_target_structure(object(), target_pair)
diff --git a/tests/integration/test_nla_facade_execution.py b/tests/integration/test_nla_facade_execution.py
new file mode 100644
index 000000000..4b11a856b
--- /dev/null
+++ b/tests/integration/test_nla_facade_execution.py
@@ -0,0 +1,235 @@
+"""Facade execution with real local synthetic Qwen artifacts, without Hub mocks."""
+
+import gc
+import json
+import weakref
+
+import pytest
+import torch
+import yaml
+from safetensors.torch import save_file
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.trainers import BpeTrainer
+from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla import NaturalLanguageAutoencoder
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig
+from transformer_lens.tools.analysis._nla_contracts import NLASamplingConfig
+
+
+@pytest.fixture
+def local_artifacts(tmp_path):
+ backend = Tokenizer(BPE(unk_token="UNK"))
+ backend.pre_tokenizer = ByteLevel(add_prefix_space=False)
+ backend.train_from_iterator(
+ ["left MARK right Summary x END objects in space"] * 10,
+ BpeTrainer(
+ vocab_size=300, special_tokens=["UNK", "MARK"], initial_alphabet=ByteLevel.alphabet()
+ ),
+ )
+ tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="UNK")
+ tokenizer.chat_template = "{{ messages[0]['content'] }}"
+ config = Qwen2Config(
+ hidden_size=32,
+ intermediate_size=64,
+ num_hidden_layers=3,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ vocab_size=tokenizer.vocab_size + 1,
+ max_position_embeddings=1024,
+ tie_word_embeddings=False,
+ eos_token_id=None,
+ bos_token_id=None,
+ )
+ pair = NLAPairConfig(
+ ArtifactRef("synthetic/target", "a" * 40),
+ ArtifactRef("synthetic/av", "b" * 40),
+ ArtifactRef("synthetic/ar", "c" * 40),
+ "qwen2",
+ 32,
+ 1,
+ 3,
+ )
+ for role, ref in (("av", pair.av), ("ar", pair.ar)):
+ directory = (
+ tmp_path / ("models--" + ref.repo_id.replace("/", "--")) / "snapshots" / ref.revision
+ )
+ directory.mkdir(parents=True)
+ role_config = Qwen2Config.from_dict(config.to_dict())
+ if role == "ar":
+ role_config.num_hidden_layers = 2
+ role_config.layer_types = role_config.layer_types[:2]
+ model = Qwen2ForCausalLM(role_config).eval()
+ role_config.save_pretrained(directory)
+ state = model.state_dict()
+ if role == "ar":
+ state = {
+ key: value
+ for key, value in state.items()
+ if key not in {"model.norm.weight", "lm_head.weight"}
+ }
+ save_file({"weight": torch.eye(32)}, directory / "value_head.safetensors")
+ save_file(state, directory / "model.safetensors")
+ tokenizer.save_pretrained(directory)
+ tokenizer_config_path = directory / "tokenizer_config.json"
+ tokenizer_config = json.loads(tokenizer_config_path.read_text())
+ tokenizer_config["tokenizer_class"] = "TokenizersBackend"
+ tokenizer_config_path.write_text(json.dumps(tokenizer_config))
+ from transformers import AutoTokenizer
+
+ live = AutoTokenizer.from_pretrained(directory, local_files_only=True)
+ marker = live.encode("MARK", add_special_tokens=False)[0]
+ actor_ids = live.apply_chat_template(
+ [{"role": "user", "content": "left MARK right"}],
+ tokenize=True,
+ add_generation_prompt=True,
+ )
+ if not isinstance(actor_ids, list):
+ actor_ids = actor_ids["input_ids"]
+ position = actor_ids.index(marker)
+ suffix = live.encode("Summary x END", add_special_tokens=True)[-1:]
+ metadata = {
+ "kind": "nla_model",
+ "schema_version": 2,
+ "role": role,
+ "d_model": 32,
+ "extraction_layer_index": 1,
+ "extraction": {
+ "injection_scale": 5.0 if role == "av" else None,
+ "mse_scale": 32**0.5,
+ },
+ "tokens": {
+ "injection_char": "MARK",
+ "injection_token_id": marker,
+ "injection_left_neighbor_id": actor_ids[position - 1],
+ "injection_right_neighbor_id": actor_ids[position + 1],
+ "critic_suffix_ids": suffix if role == "ar" else None,
+ },
+ "prompt_templates": {
+ "av": "left {injection_char} right",
+ "ar": "Summary {explanation} END",
+ },
+ }
+ (directory / "nla_meta.yaml").write_text(yaml.safe_dump(metadata))
+ target_directory = tmp_path / "models--synthetic--target" / "snapshots" / pair.target.revision
+ target_directory.mkdir(parents=True)
+ original = Qwen2ForCausalLM(config).eval()
+ original.save_pretrained(target_directory)
+ tokenizer.save_pretrained(target_directory)
+ original.config.dtype = torch.float32
+ original.config._name_or_path = pair.target.repo_id
+ original.config._commit_hash = pair.target.revision
+ target = TransformerBridge.boot_transformers(
+ pair.target.repo_id, hf_model=original, tokenizer=tokenizer, device="cpu"
+ )
+ target.eval()
+ return pair, target, tmp_path
+
+
+def test_facade_local_round_trip_and_reload(local_artifacts):
+ pair, target, cache = local_artifacts
+ with NaturalLanguageAutoencoder(
+ pair, "synthetic", cache_dir=str(cache), local_files_only=True
+ ) as nla:
+ loaded_target = nla.load_target()
+ activation = nla.extract(loaded_target, torch.tensor([1, 2, 3]))
+ del loaded_target
+ explanation = nla.verbalize(activation, max_new_tokens=2, do_sample=False)
+ assert explanation.provenance.source_verified
+ assert 0 < len(explanation.token_ids) <= 2
+ nla.unload("av")
+ assert nla._av._model is None
+ again = nla.verbalize(activation, max_new_tokens=2, do_sample=False)
+ assert again.token_ids == explanation.token_ids
+ nla.unload("av")
+ reconstruction = nla.reconstruct("x")
+ assert reconstruction.vector.shape == (32,) and not reconstruction.vector.requires_grad
+ nla.unload("ar")
+ again_vector = nla.reconstruct("x").vector
+ torch.testing.assert_close(again_vector, reconstruction.vector)
+ assert -1 <= nla.score(activation, reconstruction).cosine <= 1
+ combined = nla.explain(
+ target,
+ torch.tensor([1, 2, 3]),
+ sampling=NLASamplingConfig(max_new_tokens=2, do_sample=False),
+ )
+ assert combined.provenance.position == 2
+ with pytest.raises(ValueError, match="wrappers"):
+ nla.reconstruct("x")
+ assert nla._closed and nla._av._model is None and nla._ar._head is None
+
+
+def test_generated_explanation_reaches_reconstructor(local_artifacts):
+ pair, _, cache = local_artifacts
+ with NaturalLanguageAutoencoder(
+ pair, "synthetic", cache_dir=str(cache), local_files_only=True
+ ) as nla:
+ target = nla.load_target()
+ target_ref = weakref.ref(target)
+ activation = nla.extract(target, torch.tensor([1, 2, 3]))
+ del target
+ gc.collect()
+ assert target_ref() is None
+ nla.verbalize(activation, max_new_tokens=1, do_sample=False)
+ actor = nla._av._model
+ expected_text = "objects in space"
+ continuation = actor.tokenizer.encode(
+ f"{expected_text}", add_special_tokens=False
+ )
+ step = 0
+
+ def force_continuation(logits, hook):
+ nonlocal step
+ output = torch.full_like(logits, -1e9)
+ output[:, -1, continuation[step]] = 1e9
+ step += 1
+ return output
+
+ with actor.hooks(fwd_hooks=[("unembed.hook_out", force_continuation)]):
+ explanation = nla.verbalize(
+ activation, max_new_tokens=len(continuation), do_sample=False
+ )
+ assert explanation.text == expected_text
+ assert explanation.token_ids == tuple(continuation)
+ actor_ref = weakref.ref(actor)
+ del actor
+ nla.unload("av")
+ gc.collect()
+ assert actor_ref() is None
+ assert nla._ar._model is None
+ reconstruction = nla.reconstruct(explanation.text)
+ head = nla._ar._head
+ model = nla._ar._model
+ raw_prompt = f"Summary {expected_text} END"
+ ids = model.tokenizer.encode(raw_prompt, add_special_tokens=True)
+ with torch.no_grad():
+ residual = model.original_model.model(
+ torch.tensor([ids]), use_cache=False
+ ).last_hidden_state[0, -1]
+ expected = head(residual)
+ torch.testing.assert_close(reconstruction.vector, expected)
+ del head, model
+ nla.unload("ar")
+ assert nla._ar._model is None and nla._ar._head is None
+
+
+def test_facade_does_not_own_target(local_artifacts):
+ pair, target, cache = local_artifacts
+ nla = NaturalLanguageAutoencoder(pair, "synthetic", cache_dir=str(cache), local_files_only=True)
+ original = Qwen2ForCausalLM(target.original_model.config).eval()
+ original.config._name_or_path = pair.target.repo_id
+ original.config._commit_hash = pair.target.revision
+ owned = TransformerBridge.boot_transformers(
+ pair.target.repo_id, hf_model=original, tokenizer=target.tokenizer, device="cpu"
+ )
+ owned.eval()
+ result = nla.extract(owned, torch.tensor([1, 2, 3]))
+ ref = weakref.ref(owned)
+ del owned, original
+ gc.collect()
+ assert ref() is None
+ assert result.vector.device.type == "cpu"
+ nla.close()
diff --git a/tests/integration/test_nla_oracle_parity.py b/tests/integration/test_nla_oracle_parity.py
new file mode 100644
index 000000000..0cc4f96bf
--- /dev/null
+++ b/tests/integration/test_nla_oracle_parity.py
@@ -0,0 +1,343 @@
+"""Deterministic local and slow released-Qwen NLA reference comparisons.
+
+The direct-HF AR path follows NLACritic in kitft/natural_language_autoencoders
+at 0577769b55ad4fdd96d159e983361b97fa4e7331 (Apache-2.0): replace final
+normalization, execute the truncated backbone, and apply the separate value head.
+Oracle code is independently written; no downloaded Python client is executed.
+Exact fp32 equality is the initial acceptance criterion, not a calibrated drift
+allowance. Any nonzero drift requires investigation and recorded measurements
+before introducing a justified tolerance. Released tests are deliberately slow
+and run on CPU, one model at a time, without skips for unavailable resources.
+"""
+
+import gc
+import json
+import shutil
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+from huggingface_hub import HfApi
+from safetensors.torch import load_file
+from transformers import (
+ AutoModelForCausalLM,
+ PreTrainedTokenizerFast,
+ Qwen2Config,
+ Qwen2ForCausalLM,
+)
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla_artifacts import (
+ ArtifactRef,
+ NLAPairConfig,
+ acquire_snapshot,
+ load_local_metadata,
+ load_snapshot_tokenizer,
+ validate_value_head_file,
+)
+from transformer_lens.tools.analysis._nla_reconstructor import (
+ load_critic_model,
+ prepare_critic_tokens,
+ reconstruct_vector,
+ validate_loading_info,
+)
+from transformer_lens.tools.analysis._nla_target import extract_activation, load_target
+from transformer_lens.tools.analysis._nla_verbalizer import (
+ VerbalizerRuntime,
+ embedding_overwrite,
+)
+from transformer_lens.utilities.hf_utils import CACHE_DIR, call_hf_with_retry
+
+QWEN_PAIR = NLAPairConfig(
+ ArtifactRef("Qwen/Qwen2.5-7B-Instruct", "a09a35458c702b33eeacc393d103063234e8bc28"),
+ ArtifactRef("kitft/nla-qwen2.5-7b-L20-av", "b88469162777ae6553bc14208eb0cb579336f8f4"),
+ ArtifactRef("kitft/nla-qwen2.5-7b-L20-ar", "e2c9e57eac213d37a31612087f645ab6332c1bb6"),
+ "qwen2",
+ 3584,
+ 20,
+ 28,
+)
+_FIXED_TEXTS = (
+ "The model is considering whether its answer will be evaluated.",
+ "A description of ordinary objects and their spatial relationships.",
+)
+
+
+def assert_tensor_parity(actual: torch.Tensor, expected: torch.Tensor, label: str) -> None:
+ """Report finite tensor errors and require deterministic exact equivalence."""
+ assert actual.shape == expected.shape, label
+ assert torch.isfinite(actual).all() and torch.isfinite(expected).all(), label
+ errors = (actual.double() - expected.double()).abs()
+ print(f"{label}: max_abs={errors.max().item():.9g}, mean_abs={errors.mean().item():.9g}")
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0, msg=label)
+
+
+def actor_reference(
+ model: Any, tokens: torch.Tensor, vector: torch.Tensor, position: int, scale: float
+) -> torch.Tensor:
+ """Construct post-embedding inputs explicitly, independently of the TL hook."""
+ source = vector.detach().cpu().float()
+ norm = source.norm().clamp_min(1e-12)
+ source = source / (norm / scale).to(source.dtype)
+ with torch.inference_mode():
+ embeddings = model.get_input_embeddings()(tokens).clone()
+ embeddings[0, position] = source.to(embeddings)
+ return model(inputs_embeds=embeddings, use_cache=False).logits.detach().cpu().clone()
+
+
+def critic_reference(model: Any, head: torch.nn.Linear, tokens: torch.Tensor) -> torch.Tensor:
+ """Direct-HF reference: final norm bypass plus raw final-token value head."""
+ with torch.inference_mode():
+ residual = model.model(tokens, use_cache=False).last_hidden_state[0, -1]
+ return head(residual).detach().cpu().clone()
+
+
+def require_storage(refs: tuple[ArtifactRef, ...]) -> None:
+ """Fail before large downloads if missing artifact bytes exceed free cache space."""
+ required = 0
+ for ref in refs:
+ info = call_hf_with_retry(
+ HfApi().model_info, ref.repo_id, revision=ref.revision, files_metadata=True
+ )
+ assert info.sha == ref.revision
+ snapshot = (
+ Path(CACHE_DIR)
+ / ("models--" + ref.repo_id.replace("/", "--"))
+ / "snapshots"
+ / ref.revision
+ )
+ for file in info.siblings or ():
+ if file.rfilename.endswith(".safetensors"):
+ assert file.size is not None, "Artifact size must be known before download"
+ cached = snapshot / file.rfilename
+ if not cached.is_file() or cached.stat().st_size != file.size:
+ required += file.size
+ cache = Path(CACHE_DIR)
+ ancestor = cache
+ while not ancestor.exists():
+ ancestor = ancestor.parent
+ free = shutil.disk_usage(ancestor).free
+ reserve = 2 * 1024**3
+ if required + reserve > free:
+ raise RuntimeError(
+ f"Insufficient cache storage: {required / 1024 ** 3:.1f} GiB of missing weights "
+ f"plus 2 GiB reserve required, {free / 1024 ** 3:.1f} GiB free. "
+ "Use HF_HOME on a larger volume and rerun the slow tests."
+ )
+
+
+def _normal_activation() -> torch.Tensor:
+ target = load_target(QWEN_PAIR, device="cpu")
+ try:
+ result = extract_activation(
+ target,
+ QWEN_PAIR,
+ "qwen-parity",
+ "The auditor asked about the model's internal reasoning.",
+ position=-1,
+ prepend_bos=False,
+ )
+ assert result.provenance.source_verified
+ return result.vector.clone()
+ finally:
+ del target
+ gc.collect()
+
+
+def test_parity_assertion_detects_drift():
+ with pytest.raises(AssertionError):
+ assert_tensor_parity(torch.tensor([1.001]), torch.tensor([1.0]), "deliberate drift")
+ with pytest.raises(AssertionError):
+ assert_tensor_parity(torch.tensor([float("nan")]), torch.tensor([1.0]), "nonfinite")
+
+
+def test_direct_hf_oracle_mechanics(tmp_path):
+ """Normal-CI sanity check of oracle construction on a tiny local Qwen model."""
+ config = Qwen2Config(
+ hidden_size=32,
+ intermediate_size=64,
+ num_hidden_layers=2,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ vocab_size=32,
+ max_position_embeddings=64,
+ tie_word_embeddings=False,
+ )
+ original = Qwen2ForCausalLM(config).eval()
+ original.config._attn_implementation = "eager"
+ original.save_pretrained(tmp_path)
+ original.config.dtype = torch.float32
+ tokens = torch.tensor([[1, 2, 3]])
+ vector = torch.arange(1.0, 33.0)
+ expected = actor_reference(original, tokens, vector, 1, 5.0)
+ from tokenizers import Tokenizer
+ from tokenizers.models import WordLevel
+
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=Tokenizer(WordLevel({"UNK": 0}, unk_token="UNK")), unk_token="UNK"
+ )
+ bridge = TransformerBridge.boot_transformers(
+ "local-qwen", hf_model=original, tokenizer=tokenizer, device="cpu"
+ )
+ bridge.eval()
+ with torch.inference_mode(), bridge.hooks(
+ fwd_hooks=[("hook_embed", embedding_overwrite(vector, 5.0, 1, 3, cached=True))]
+ ):
+ actual = bridge(tokens, use_cache=False).cpu()
+ assert_tensor_parity(actual, expected, "tiny AV logits")
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 32,
+ 1,
+ 3,
+ )
+ from transformer_lens.tools.analysis._nla_artifacts import NLAMetadata
+
+ metadata = NLAMetadata(
+ "ar",
+ 32,
+ 1,
+ "marker",
+ 2,
+ "{injection_char}",
+ "{explanation}",
+ 32**0.5,
+ None,
+ None,
+ None,
+ (3,),
+ )
+ original.model.norm = torch.nn.Identity()
+ head = torch.nn.Linear(32, 32, bias=False).eval()
+ expected_vector = critic_reference(original, head, tokens)
+ actual_vector = reconstruct_vector(bridge, head, tokens, metadata, pair).vector
+ assert_tensor_parity(actual_vector, expected_vector, "tiny AR vector")
+
+
+@pytest.mark.slow
+def test_released_qwen_actor_prefill_parity():
+ require_storage((QWEN_PAIR.target, QWEN_PAIR.av))
+ normal = _normal_activation()
+ vectors = (normal, torch.linspace(-1.0, 1.0, QWEN_PAIR.d_model))
+ directory = acquire_snapshot(QWEN_PAIR.av, include_weights=True)
+ tokenizer = load_snapshot_tokenizer(directory)
+ metadata = load_local_metadata(directory, QWEN_PAIR, "av", tokenizer)
+ assert metadata.actor_token_ids is not None and metadata.injection_position is not None
+ assert metadata.injection_scale is not None
+ tokens = torch.tensor([metadata.actor_token_ids], dtype=torch.int64)
+ direct: Any = AutoModelForCausalLM.from_pretrained(
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ )
+ direct.eval()
+ expected = [
+ actor_reference(direct, tokens, v, metadata.injection_position, metadata.injection_scale)
+ for v in vectors
+ ]
+ del direct
+ gc.collect()
+ runtime = VerbalizerRuntime(QWEN_PAIR, "qwen-parity", device="cpu", local_files_only=True)
+ model = runtime._load(directory, tokenizer)
+ try:
+ for index, (vector, reference) in enumerate(zip(vectors, expected)):
+ hook = embedding_overwrite(
+ vector,
+ metadata.injection_scale,
+ metadata.injection_position,
+ tokens.shape[1],
+ cached=True,
+ )
+ with torch.inference_mode(), model.hooks(fwd_hooks=[("hook_embed", hook)]):
+ actual = model(tokens, use_cache=False).detach().cpu()
+ assert_tensor_parity(actual, reference, f"released AV input {index}")
+ finally:
+ del model, runtime
+ gc.collect()
+
+
+@pytest.mark.slow
+def test_released_qwen_critic_vector_parity():
+ require_storage((QWEN_PAIR.ar,))
+ directory = acquire_snapshot(QWEN_PAIR.ar, include_weights=True)
+ tokenizer = load_snapshot_tokenizer(directory)
+ metadata = load_local_metadata(directory, QWEN_PAIR, "ar", tokenizer)
+ config = json.loads((directory / "config.json").read_text())
+ head_path = directory / "value_head.safetensors"
+ validate_value_head_file(head_path, QWEN_PAIR.d_model)
+ head = torch.nn.Linear(QWEN_PAIR.d_model, QWEN_PAIR.d_model, bias=False)
+ head.load_state_dict(load_file(str(head_path)), strict=True)
+ head.eval()
+ tokens_list = []
+ for text in _FIXED_TEXTS:
+ tokens = prepare_critic_tokens(tokenizer, metadata, text, config["max_position_embeddings"])
+ independent = tokenizer(
+ metadata.critic_template.format(explanation=text),
+ return_tensors="pt",
+ add_special_tokens=True,
+ )["input_ids"]
+ assert torch.equal(tokens, independent)
+ tokens_list.append(tokens)
+ loaded: Any = AutoModelForCausalLM.from_pretrained(
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ output_loading_info=True,
+ )
+ direct, info = loaded
+ del loaded
+ validate_loading_info(info)
+ direct.model.norm = torch.nn.Identity()
+ direct.lm_head = torch.nn.Identity()
+ direct.eval()
+ expected = [critic_reference(direct, head, tokens) for tokens in tokens_list]
+ del direct
+ gc.collect()
+ model = load_critic_model(directory, QWEN_PAIR, tokenizer, "cpu")
+ try:
+ for index, (tokens, reference) in enumerate(zip(tokens_list, expected)):
+ actual = reconstruct_vector(model, head, tokens, metadata, QWEN_PAIR).vector
+ assert_tensor_parity(actual, reference, f"released AR text {index}")
+ finally:
+ del model
+ gc.collect()
+
+
+@pytest.mark.slow
+def test_released_qwen_sequential_round_trip():
+ from transformer_lens.tools.analysis import NaturalLanguageAutoencoder
+
+ require_storage((QWEN_PAIR.target, QWEN_PAIR.av, QWEN_PAIR.ar))
+ with NaturalLanguageAutoencoder.from_pretrained(
+ "qwen2.5-7b-instruct-L20", device="cpu", local_files_only=True
+ ) as nla:
+ target = nla.load_target()
+ activation = nla.extract(
+ target,
+ "The auditor asked about the model's internal reasoning.",
+ prepend_bos=False,
+ )
+ del target
+ gc.collect()
+ explanation = nla.verbalize(activation, max_new_tokens=512)
+ nla.unload("av")
+ gc.collect()
+ assert (
+ explanation.text is not None
+ ), f"AV parse failure: {explanation.parse_status}; {explanation.raw_text!r}"
+ result = nla.reconstruct(explanation.text)
+ nla.unload("ar")
+ assert result.vector.shape == activation.vector.shape
+ assert torch.isfinite(result.vector).all()
+ assert torch.any(result.vector != 0)
+ score = nla.score(activation, result)
+ assert -1 <= score.cosine <= 1
+ assert 0 <= score.normalized_mse <= 4
diff --git a/tests/integration/test_nla_packaging.py b/tests/integration/test_nla_packaging.py
new file mode 100644
index 000000000..c01f6e93d
--- /dev/null
+++ b/tests/integration/test_nla_packaging.py
@@ -0,0 +1,83 @@
+"""Built-wheel NLA resources without Hub or model downloads."""
+
+import json
+import os
+import subprocess
+import sys
+import zipfile
+from pathlib import Path
+
+_REQUIRED = (
+ "transformer_lens/tools/analysis/nla_registry.json",
+ "transformer_lens/tools/analysis/_nla_artifacts.py",
+ "transformer_lens/tools/analysis/_nla_contracts.py",
+ "transformer_lens/tools/analysis/_nla_target.py",
+ "transformer_lens/tools/analysis/_nla_verbalizer.py",
+ "transformer_lens/tools/analysis/_nla_reconstructor.py",
+ "transformer_lens/tools/analysis/_nla.py",
+ "transformer_lens/tools/analysis/nla.py",
+)
+
+
+def test_built_wheel_nla_resources(tmp_path):
+ repository = Path(__file__).resolve().parents[2]
+ output = tmp_path / "dist"
+ build_environment = os.environ.copy()
+ build_environment["UV_CACHE_DIR"] = str(tmp_path / "fresh-uv-cache")
+ subprocess.run(
+ ["uv", "build", "--wheel", "--out-dir", str(output)],
+ cwd=repository,
+ env=build_environment,
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ wheels = list(output.glob("*.whl"))
+ assert len(wheels) == 1
+ extracted = tmp_path / "installed"
+ with zipfile.ZipFile(wheels[0]) as wheel:
+ for name in _REQUIRED:
+ assert name in wheel.namelist(), f"Missing wheel resource: {name}"
+ registry = json.loads(wheel.read(_REQUIRED[0]))
+ assert registry["schema_version"] == 1
+ assert isinstance(registry["models"], dict)
+ wheel.extractall(extracted)
+ script = r"""
+import pathlib
+import socket
+import sys
+
+root = pathlib.Path(sys.argv[1])
+
+def forbid_network(*args, **kwargs):
+ raise AssertionError("Wheel resource imports must not initiate network access")
+
+socket.socket.connect = forbid_network
+socket.create_connection = forbid_network
+sys.path.insert(0, str(root))
+import transformer_lens
+import transformer_lens.tools.analysis
+assert root in pathlib.Path(transformer_lens.__file__).parents
+assert root in pathlib.Path(transformer_lens.tools.analysis.__file__).parents
+from transformer_lens.tools.analysis._nla_artifacts import load_registry
+from transformer_lens.tools.analysis._nla_contracts import NLASamplingConfig
+assert "qwen2.5-7b-instruct-L20" in load_registry()
+from transformer_lens.tools.analysis import NaturalLanguageAutoencoder, NLAActivation
+assert NaturalLanguageAutoencoder is not None and NLAActivation is not None
+assert NLASamplingConfig().max_new_tokens == 512
+from transformer_lens.tools.analysis import _nla_artifacts
+assert root in pathlib.Path(_nla_artifacts.__file__).parents
+print("Isolated wheel NLA resource import passed")
+"""
+ environment = os.environ.copy()
+ environment["HF_HUB_OFFLINE"] = "1"
+ environment["TRANSFORMERS_OFFLINE"] = "1"
+ result = subprocess.run(
+ [sys.executable, "-I", "-c", script, str(extracted)],
+ cwd=tmp_path,
+ env=environment,
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ assert "Isolated wheel NLA resource import passed" in result.stdout
diff --git a/tests/unit/tools/test_nla.py b/tests/unit/tools/test_nla.py
new file mode 100644
index 000000000..b5b2d9e7a
--- /dev/null
+++ b/tests/unit/tools/test_nla.py
@@ -0,0 +1,382 @@
+"""Local artifact validation without model or Hub loading."""
+
+import copy
+import json
+import math
+
+import pytest
+import torch
+import yaml
+from safetensors.torch import save_file
+
+from transformer_lens.tools.analysis._nla_artifacts import (
+ ArtifactRef,
+ NLAPairConfig,
+ load_local_metadata,
+ load_registry,
+ parse_registry,
+ validate_metadata,
+ validate_snapshot_path,
+ validate_value_head,
+ validate_value_head_file,
+)
+
+
+class LocalTokenizer:
+ unk_token_id = 0
+
+ def __init__(self):
+ self.injection_ids = [7]
+ self.actor_ids = [1, 2, 7, 3, 4]
+ self.critic_ids = [1, 8, 9]
+
+ def encode(self, text, *, add_special_tokens):
+ if text == "MARK":
+ assert add_special_tokens is False
+ return self.injection_ids
+ assert text == "Summary: x END"
+ assert add_special_tokens is True
+ return self.critic_ids
+
+ def apply_chat_template(self, conversation, *, tokenize, add_generation_prompt):
+ assert conversation == [{"role": "user", "content": "Explain MARK now"}]
+ assert tokenize is True and add_generation_prompt is True
+ return self.actor_ids
+
+
+@pytest.fixture
+def pair():
+ return NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 4,
+ 1,
+ 3,
+ )
+
+
+@pytest.fixture
+def actor():
+ return {
+ "kind": "nla_model",
+ "schema_version": 2,
+ "role": "av",
+ "d_model": 4,
+ "extraction_layer_index": 1,
+ "extraction": {"injection_scale": 150.0, "mse_scale": 2.0},
+ "tokens": {
+ "injection_char": "MARK",
+ "injection_token_id": 7,
+ "injection_left_neighbor_id": 2,
+ "injection_right_neighbor_id": 3,
+ "critic_suffix_ids": None,
+ },
+ "prompt_templates": {
+ "av": "Explain {injection_char} now",
+ "ar": "Summary: {explanation} END",
+ },
+ }
+
+
+@pytest.fixture
+def critic(actor):
+ result = copy.deepcopy(actor)
+ result["role"] = "ar"
+ result["extraction"]["injection_scale"] = None
+ result["tokens"]["critic_suffix_ids"] = [8, 9]
+ result["critic"] = {"extraction_layer_index": 1}
+ return result
+
+
+def config(role):
+ return {"model_type": "qwen2", "hidden_size": 4, "num_hidden_layers": 3 if role == "av" else 2}
+
+
+def test_valid_roles(actor, critic, pair):
+ av = validate_metadata(actor, config("av"), pair, "av", LocalTokenizer())
+ ar = validate_metadata(critic, config("ar"), pair, "ar", LocalTokenizer())
+ assert av.actor_token_ids == (1, 2, 7, 3, 4)
+ assert av.injection_position == 2 and av.injection_scale == 150
+ assert av.critic_suffix_ids is None
+ assert ar.critic_suffix_ids == (8, 9)
+ assert ar.actor_token_ids is None and ar.injection_scale is None
+
+
+@pytest.mark.parametrize(
+ "path,value",
+ [
+ (("kind",), "nla_dataset"),
+ (("schema_version",), 1),
+ (("schema_version",), True),
+ (("role",), "ar"),
+ (("d_model",), 5),
+ (("d_model",), False),
+ (("d_model",), 0),
+ (("extraction_layer_index",), 0),
+ (("extraction_layer_index",), True),
+ (("extraction", "injection_scale"), None),
+ (("extraction", "injection_scale"), 0),
+ (("extraction", "injection_scale"), -1),
+ (("extraction", "injection_scale"), float("inf")),
+ (("extraction", "injection_scale"), True),
+ (("extraction", "mse_scale"), None),
+ (("extraction", "mse_scale"), float("nan")),
+ (("extraction", "mse_scale"), 1.99),
+ (("tokens", "injection_token_id"), True),
+ (("tokens", "injection_left_neighbor_id"), 10),
+ (("tokens", "injection_right_neighbor_id"), 10),
+ (("tokens", "injection_char"), ""),
+ (("tokens", "critic_suffix_ids"), [8]),
+ (("prompt_templates", "av"), "{unknown}"),
+ (("prompt_templates", "av"), "{injection_char}{injection_char}"),
+ (("prompt_templates", "av"), "{injection_char!r}"),
+ (("prompt_templates", "ar"), "{explanation:>10}"),
+ (("prompt_templates", "ar"), "{"),
+ (("tokens",), []),
+ (("extraction",), None),
+ ],
+)
+def test_invalid_actor_metadata(actor, pair, path, value):
+ mapping = actor
+ for key in path[:-1]:
+ mapping = mapping[key]
+ mapping[path[-1]] = value
+ with pytest.raises(ValueError):
+ validate_metadata(actor, config("av"), pair, "av", LocalTokenizer())
+
+
+@pytest.mark.parametrize(
+ "key",
+ [
+ "kind",
+ "schema_version",
+ "role",
+ "d_model",
+ "tokens",
+ "extraction",
+ "prompt_templates",
+ "extraction_layer_index",
+ ],
+)
+def test_missing_metadata_fields(actor, pair, key):
+ actor.pop(key)
+ with pytest.raises(ValueError):
+ validate_metadata(actor, config("av"), pair, "av", LocalTokenizer())
+
+
+@pytest.mark.parametrize(
+ "field,value",
+ [
+ ("model_type", "llama"),
+ ("hidden_size", 8),
+ ("hidden_size", True),
+ ("num_hidden_layers", 2),
+ ("num_hidden_layers", False),
+ ],
+)
+def test_invalid_model_config(actor, pair, field, value):
+ model = config("av")
+ model[field] = value
+ with pytest.raises(ValueError):
+ validate_metadata(actor, model, pair, "av", LocalTokenizer())
+
+
+@pytest.mark.parametrize("ids", [[], [7, 7], [0], [8], [True]])
+def test_reserved_token_drift(actor, pair, ids):
+ tokenizer = LocalTokenizer()
+ tokenizer.injection_ids = ids
+ with pytest.raises(ValueError):
+ validate_metadata(actor, config("av"), pair, "av", tokenizer)
+
+
+def test_unknown_reserved_token(actor, pair):
+ tokenizer = LocalTokenizer()
+ tokenizer.unk_token_id = 7
+ with pytest.raises(ValueError, match="UNK"):
+ validate_metadata(actor, config("av"), pair, "av", tokenizer)
+
+
+@pytest.mark.parametrize("ids", [[1, 2, 3], [2, 7, 7, 3], [7, 3], [2, 7], [2, 7, True]])
+def test_placeholder_occurrence_and_position(actor, pair, ids):
+ tokenizer = LocalTokenizer()
+ tokenizer.actor_ids = ids
+ with pytest.raises(ValueError):
+ validate_metadata(actor, config("av"), pair, "av", tokenizer)
+
+
+@pytest.mark.parametrize(
+ "path,value",
+ [
+ (("tokens", "critic_suffix_ids"), None),
+ (("tokens", "critic_suffix_ids"), []),
+ (("tokens", "critic_suffix_ids"), [9, 8]),
+ (("tokens", "critic_suffix_ids"), [True]),
+ (("extraction", "injection_scale"), 150),
+ (("critic", "extraction_layer_index"), 0),
+ ],
+)
+def test_invalid_critic_metadata(critic, pair, path, value):
+ critic[path[0]][path[1]] = value
+ with pytest.raises(ValueError):
+ validate_metadata(critic, config("ar"), pair, "ar", LocalTokenizer())
+
+
+def test_wrong_critic_depth(critic, pair):
+ with pytest.raises(ValueError, match="depth"):
+ validate_metadata(critic, config("av"), pair, "ar", LocalTokenizer())
+
+
+def test_qwen_scale_serialization(actor, pair):
+ large = NLAPairConfig(pair.target, pair.av, pair.ar, "qwen2", 3584, 20, 28)
+ actor["d_model"] = 3584
+ actor["extraction_layer_index"] = 20
+ actor["extraction"]["mse_scale"] = 59.86651818838306
+ av_config = {"model_type": "qwen2", "hidden_size": 3584, "num_hidden_layers": 28}
+ assert math.isclose(
+ validate_metadata(actor, av_config, large, "av", LocalTokenizer()).mse_scale,
+ math.sqrt(3584),
+ )
+ actor["extraction"]["mse_scale"] += 1e-8
+ with pytest.raises(ValueError, match="mse_scale"):
+ validate_metadata(actor, av_config, large, "av", LocalTokenizer())
+
+
+def test_local_metadata_files(tmp_path, actor, pair):
+ (tmp_path / "nla_meta.yaml").write_text(yaml.safe_dump(actor))
+ (tmp_path / "config.json").write_text(json.dumps(config("av")))
+ assert load_local_metadata(tmp_path, pair, "av", LocalTokenizer()).d_model == 4
+ (tmp_path / "nla_meta.yaml").write_text("!!python/object:builtins.object {}")
+ with pytest.raises(ValueError, match="syntax"):
+ load_local_metadata(tmp_path, pair, "av", LocalTokenizer())
+
+
+@pytest.mark.parametrize("revision", ["main", "a" * 39, "A" * 40, "../snapshot", ""])
+def test_immutable_revisions(revision):
+ with pytest.raises(ValueError, match="revision"):
+ ArtifactRef("example/av", revision)
+
+
+@pytest.mark.parametrize(
+ "repo", ["../av", "example/../av", "av", "/example/av", "./av", "example/..", "example/."]
+)
+def test_repo_identifiers(repo):
+ with pytest.raises(ValueError, match="repo_id"):
+ ArtifactRef(repo, "a" * 40)
+
+
+def registry_entry():
+ return {
+ "target": {"repo_id": "example/target", "revision": "a" * 40},
+ "av": {"repo_id": "example/av", "revision": "b" * 40},
+ "ar": {"repo_id": "example/ar", "revision": "c" * 40},
+ "model_type": "qwen2",
+ "d_model": 4,
+ "layer": 1,
+ "target_layers": 3,
+ }
+
+
+def test_registry_resource_and_parser(pair):
+ registered = load_registry()["qwen2.5-7b-instruct-L20"]
+ assert registered.d_model == 3584 and registered.layer == 20
+ assert registered.target.repo_id == "Qwen/Qwen2.5-7B-Instruct"
+ assert parse_registry({"schema_version": 1, "models": {"tiny": registry_entry()}}) == {
+ "tiny": pair
+ }
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ None,
+ [],
+ {"schema_version": True, "models": {}},
+ {"schema_version": 2, "models": {}},
+ {"schema_version": 1, "models": []},
+ ],
+)
+def test_invalid_registry(value):
+ with pytest.raises(ValueError):
+ parse_registry(value)
+
+
+@pytest.mark.parametrize(
+ "key,value", [("d_model", True), ("layer", -1), ("layer", 3), ("target_layers", 0)]
+)
+def test_invalid_pair_contract(key, value):
+ entry = registry_entry()
+ entry[key] = value
+ with pytest.raises(ValueError):
+ parse_registry({"schema_version": 1, "models": {"tiny": entry}})
+
+
+def test_distinct_artifact_identities():
+ entry = registry_entry()
+ entry["ar"] = entry["av"]
+ with pytest.raises(ValueError, match="distinct"):
+ parse_registry({"schema_version": 1, "models": {"tiny": entry}})
+
+
+@pytest.mark.parametrize(
+ "state",
+ [
+ {},
+ {"weight": torch.eye(4), "bias": torch.zeros(4)},
+ {"weight": torch.eye(4), "other": torch.zeros(1)},
+ {"weight": torch.eye(3)},
+ {"weight": torch.zeros(4, 4, dtype=torch.int64)},
+ {"weight": torch.full((4, 4), float("nan"))},
+ ],
+)
+def test_invalid_value_heads(state):
+ with pytest.raises(ValueError):
+ validate_value_head(state, 4)
+
+
+def test_snapshot_path_identity(tmp_path, pair):
+ path = tmp_path / "models--example--av" / "snapshots" / pair.av.revision
+ assert validate_snapshot_path(path, pair.av) == path
+ for invalid in (
+ path.with_name(pair.ar.revision),
+ tmp_path / "models--example--ar" / "snapshots" / pair.av.revision,
+ tmp_path / pair.av.revision,
+ ):
+ with pytest.raises(ValueError, match="repository and revision"):
+ validate_snapshot_path(invalid, pair.av)
+
+
+def test_chat_template_batch_encoding(actor, pair):
+ from transformers import BatchEncoding
+
+ tokenizer = LocalTokenizer()
+ tokenizer.actor_ids = BatchEncoding({"input_ids": [1, 2, 7, 3, 4]})
+ result = validate_metadata(actor, config("av"), pair, "av", tokenizer)
+ assert result.injection_position == 2
+
+
+@pytest.mark.parametrize(
+ "role,section,key",
+ [
+ ("av", "tokens", "critic_suffix_ids"),
+ ("ar", "extraction", "injection_scale"),
+ ],
+)
+def test_explicit_nullable_fields(actor, critic, pair, role, section, key):
+ meta = actor if role == "av" else critic
+ del meta[section][key]
+ with pytest.raises(ValueError, match="must record"):
+ validate_metadata(meta, config(role), pair, role, LocalTokenizer())
+
+
+def test_value_head_safetensors(tmp_path):
+ path = tmp_path / "value_head.safetensors"
+ save_file({"weight": torch.eye(4)}, path)
+ validate_value_head_file(path, 4)
+ save_file({"weight": torch.eye(3)}, path)
+ with pytest.raises(ValueError, match="shape"):
+ validate_value_head_file(path, 4)
+ save_file({"weight": torch.eye(4), "bias": torch.zeros(4)}, path)
+ with pytest.raises(ValueError, match="bias"):
+ validate_value_head_file(path, 4)
diff --git a/tests/unit/tools/test_nla_contracts.py b/tests/unit/tools/test_nla_contracts.py
new file mode 100644
index 000000000..2e4d1c970
--- /dev/null
+++ b/tests/unit/tools/test_nla_contracts.py
@@ -0,0 +1,283 @@
+"""Synthetic result, parsing, scoring, and length contract checks."""
+
+import math
+
+import pytest
+import torch
+from jaxtyping import TypeCheckError
+
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef
+from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ ExplanationStatus,
+ NLAActivation,
+ NLAExplanation,
+ NLAReconstruction,
+ NLASamplingConfig,
+ SourceProvenance,
+ detached_vector,
+ parse_explanation,
+ score_reconstruction,
+ validate_context_length,
+ validate_reconstruction_text,
+)
+
+
+@pytest.fixture
+def provenance():
+ return ActivationProvenance("synthetic", SourceProvenance.UNKNOWN_VECTOR)
+
+
+@pytest.fixture
+def artifact():
+ return ArtifactRef("example/ar", "a" * 40)
+
+
+def test_loading_evidence_is_not_live_weight_integrity(artifact):
+ value = ActivationProvenance(
+ "synthetic", SourceProvenance.PINNED_LOAD, artifact, 20, 0, "raw tokens, no BOS"
+ )
+ assert value.source_verified is True
+ assert (
+ ActivationProvenance("synthetic", SourceProvenance.CALLER_ATTESTED).source_verified is False
+ )
+ with pytest.raises(ValueError, match="requires"):
+ ActivationProvenance("synthetic", SourceProvenance.PINNED_LOAD)
+
+
+@pytest.mark.parametrize(
+ "kwargs",
+ [
+ {"artifact_name": ""},
+ {"source": "pinned_load"},
+ {"position": -1},
+ {"layer": True},
+ {"basis": "processed"},
+ {"tokenization_policy": ""},
+ ],
+)
+def test_invalid_provenance(kwargs):
+ args = {"artifact_name": "synthetic", "source": SourceProvenance.UNKNOWN_VECTOR}
+ args.update(kwargs)
+ with pytest.raises((ValueError, TypeCheckError)):
+ ActivationProvenance(**args)
+
+
+@pytest.mark.parametrize("kind", ["activation", "reconstruction"])
+def test_result_detaches_and_owns_storage(kind, provenance, artifact):
+ source = torch.arange(1.0, 5.0, requires_grad=True)
+ vector = source * 2
+ result = (
+ NLAActivation(vector, provenance)
+ if kind == "activation"
+ else NLAReconstruction(vector, artifact)
+ )
+ assert result.vector.requires_grad is False and result.vector.grad_fn is None
+ assert result.vector.data_ptr() != vector.data_ptr()
+ vector.detach().fill_(99)
+ torch.testing.assert_close(result.vector, torch.tensor([2.0, 4.0, 6.0, 8.0]))
+
+
+@pytest.mark.parametrize(
+ "vector",
+ [
+ torch.empty(0),
+ torch.ones(2, 2),
+ torch.zeros(3),
+ torch.ones(2, dtype=torch.int64),
+ torch.tensor([float("nan")]),
+ torch.tensor([float("inf")]),
+ torch.tensor([1j]),
+ ],
+)
+def test_invalid_vectors(vector, provenance, artifact):
+ with pytest.raises(ValueError):
+ detached_vector(vector)
+ with pytest.raises(ValueError):
+ NLAActivation(vector, provenance)
+ with pytest.raises(ValueError):
+ NLAReconstruction(vector, artifact)
+
+
+@pytest.mark.parametrize(
+ "raw,status,text",
+ [
+ (
+ "prefix first\nsecond suffix",
+ ExplanationStatus.SUCCESS,
+ "first\nsecond",
+ ),
+ ("no tags", ExplanationStatus.MISSING, None),
+ ("", ExplanationStatus.EMPTY, None),
+ (" \n ", ExplanationStatus.EMPTY, None),
+ ("x", ExplanationStatus.INCOMPLETE, None),
+ ("x", ExplanationStatus.INCOMPLETE, None),
+ ("x", ExplanationStatus.INCOMPLETE, None),
+ (
+ "ab",
+ ExplanationStatus.MULTIPLE,
+ None,
+ ),
+ (
+ "ab",
+ ExplanationStatus.NESTED,
+ None,
+ ),
+ ("x", ExplanationStatus.MALFORMED, None),
+ ("x", ExplanationStatus.MALFORMED, None),
+ ("a", ExplanationStatus.INCOMPLETE, None),
+ ],
+)
+def test_parser(raw, status, text, provenance):
+ result = parse_explanation(raw, (1, 2), provenance)
+ assert result.raw_text == raw and result.token_ids == (1, 2)
+ assert result.parse_status is status and result.text == text
+ assert result.provenance is provenance
+
+
+def test_reconstruction_preserves_authored_whitespace():
+ text = " authored\n explanation "
+ assert validate_reconstruction_text(text) == text
+ assert validate_reconstruction_text("An untagged full response cannot be detected semantically")
+ assert validate_reconstruction_text("an ordinary tag")
+
+
+@pytest.mark.parametrize(
+ "text",
+ [
+ None,
+ "",
+ " \n",
+ "x",
+ "",
+ "x",
+ ],
+)
+def test_invalid_reconstruction_text(text):
+ with pytest.raises(ValueError):
+ validate_reconstruction_text(text)
+
+
+def test_explanation_result_invariants(provenance):
+ with pytest.raises(ValueError):
+ NLAExplanation("text", "raw", (), ExplanationStatus.MISSING, provenance)
+ with pytest.raises(ValueError):
+ NLAExplanation(None, "raw", (), ExplanationStatus.SUCCESS, provenance)
+ with pytest.raises(ValueError):
+ NLAExplanation(" text ", "raw", (), ExplanationStatus.SUCCESS, provenance)
+ with pytest.raises(ValueError):
+ parse_explanation("x", (-1,), provenance)
+
+
+@pytest.mark.parametrize(
+ "b,cosine,mse", [([1.0, 0.0], 1.0, 0.0), ([0.0, 1.0], 0.0, 2.0), ([-1.0, 0.0], -1.0, 4.0)]
+)
+def test_metric_identities(b, cosine, mse, provenance, artifact):
+ result = score_reconstruction(
+ NLAActivation(torch.tensor([2.0, 0.0]), provenance),
+ NLAReconstruction(torch.tensor(b), artifact),
+ )
+ assert result.cosine == pytest.approx(cosine)
+ assert result.normalized_mse == pytest.approx(mse)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
+def test_metric_matches_normalized_coordinate_mse(dtype):
+ a = torch.tensor([1.0, 2.0, -3.0, 4.0], dtype=dtype)
+ b = torch.tensor([-2.0, 3.0, 4.0, 1.0], dtype=dtype)
+ result = score_reconstruction(a, b)
+ a64, b64 = a.double(), b.double()
+ scale = math.sqrt(a.numel())
+ expected = ((scale * a64 / a64.norm() - scale * b64 / b64.norm()) ** 2).mean().item()
+ assert result.normalized_mse == pytest.approx(expected, abs=1e-14)
+ assert 0 <= result.normalized_mse <= 4
+
+
+@pytest.mark.parametrize("scale", [1e-300, 1e300])
+def test_finite_extremes_do_not_overflow_norms(scale):
+ a = torch.tensor([scale, scale], dtype=torch.float64)
+ result = score_reconstruction(a, a)
+ assert result.cosine == pytest.approx(1.0)
+ assert result.normalized_mse == pytest.approx(0.0, abs=1e-14)
+
+
+def test_metric_rejects_dimension_mismatch():
+ with pytest.raises(ValueError, match="dimensions"):
+ score_reconstruction(torch.ones(2), torch.ones(3))
+
+
+@pytest.mark.parametrize(
+ "vector", [torch.zeros(2), torch.tensor([float("nan")]), torch.tensor([float("inf")])]
+)
+def test_metric_rejects_invalid_input(vector):
+ with pytest.raises(ValueError):
+ score_reconstruction(vector, torch.ones(2))
+ with pytest.raises(ValueError):
+ score_reconstruction(torch.ones(2), vector)
+
+
+def test_sampling_defaults_and_greedy():
+ defaults = NLASamplingConfig()
+ assert defaults.max_new_tokens == 512 and defaults.temperature == 1
+ assert defaults.do_sample is True and defaults.use_past_kv_cache is True
+ assert NLASamplingConfig(do_sample=False, temperature=0).do_sample is False
+ assert NLASamplingConfig(use_past_kv_cache=False).use_past_kv_cache is False
+
+
+@pytest.mark.parametrize(
+ "kwargs",
+ [
+ {"max_new_tokens": 0},
+ {"max_new_tokens": -1},
+ {"max_new_tokens": True},
+ {"max_new_tokens": 1.5},
+ {"temperature": 0},
+ {"temperature": -1},
+ {"temperature": float("nan")},
+ {"temperature": float("inf")},
+ {"temperature": True},
+ {"top_p": 0},
+ {"top_p": 1.1},
+ {"top_p": float("nan")},
+ {"top_p": True},
+ {"top_k": -1},
+ {"top_k": True},
+ {"do_sample": 1},
+ {"use_past_kv_cache": 0},
+ {"do_sample": False, "top_p": 0.9},
+ {"do_sample": False, "top_k": 10},
+ {"temperature": 10**400},
+ ],
+)
+def test_invalid_sampling(kwargs):
+ with pytest.raises((ValueError, TypeCheckError)):
+ NLASamplingConfig(**kwargs)
+
+
+def test_context_exact_boundaries():
+ validate_context_length(4, 4)
+ validate_context_length(1, 4, 3)
+ sampling = NLASamplingConfig(max_new_tokens=3)
+ sampling.validate_prompt(1, 4)
+ with pytest.raises(ValueError, match="exceeds"):
+ sampling.validate_prompt(2, 4)
+
+
+@pytest.mark.parametrize(
+ "args",
+ [
+ (0, 4, 0),
+ (1, 0, 0),
+ (True, 4, 0),
+ (1, False, 0),
+ (1, 4, -1),
+ (1, 4, True),
+ (5, 4, 0),
+ (1, 4, 4),
+ ],
+)
+def test_invalid_lengths(args):
+ with pytest.raises(ValueError):
+ validate_context_length(*args)
diff --git a/tests/unit/tools/test_nla_facade.py b/tests/unit/tools/test_nla_facade.py
new file mode 100644
index 000000000..aeee9fb55
--- /dev/null
+++ b/tests/unit/tools/test_nla_facade.py
@@ -0,0 +1,90 @@
+"""Private facade lifecycle and import contracts without model loading."""
+
+import pytest
+import torch
+from jaxtyping import TypeCheckError
+
+from transformer_lens.tools.analysis import NaturalLanguageAutoencoder
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig
+
+
+@pytest.fixture
+def nla():
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 4,
+ 1,
+ 3,
+ )
+ return NaturalLanguageAutoencoder(pair, "synthetic", local_files_only=True)
+
+
+def test_unknown_registry_fails_without_loading():
+ with pytest.raises(ValueError, match="No verified"):
+ NaturalLanguageAutoencoder.from_pretrained("unregistered", local_files_only=True)
+
+
+def test_context_closes_on_error(nla):
+ with pytest.raises(ValueError, match="example"):
+ with nla:
+ raise ValueError("example")
+ assert nla._closed and nla._av._model is None and nla._ar._model is None
+ nla.close()
+ with pytest.raises(RuntimeError, match="closed"):
+ nla.reconstruct("text")
+
+
+def test_unload_is_idempotent_and_validated(nla):
+ nla.unload("av")
+ nla.unload("av")
+ nla.unload("ar")
+ with pytest.raises((ValueError, TypeCheckError)):
+ nla.unload("invalid")
+ assert not nla._closed
+
+
+def test_operations_reject_concurrency(nla):
+ assert nla._lock.acquire(blocking=False)
+ try:
+ with pytest.raises(RuntimeError, match="concurrent"):
+ nla.load_target()
+ with pytest.raises(RuntimeError, match="during"):
+ nla.close()
+ finally:
+ nla._lock.release()
+
+
+def test_score_and_invalid_input(nla):
+ assert nla.score(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])).normalized_mse == 2
+ with pytest.raises(ValueError, match="nonempty"):
+ nla.reconstruct("")
+ assert nla._ar._model is None
+
+
+def test_in_memory_tokenizer_bos_does_not_require_loading_path():
+ from tokenizers import Tokenizer
+ from tokenizers.models import WordLevel
+ from transformers import PreTrainedTokenizerFast
+
+ from transformer_lens.utilities.tokenize_utils import get_tokenizer_with_bos
+
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=Tokenizer(WordLevel({"UNK": 0, "BOS": 1}, unk_token="UNK")),
+ unk_token="UNK",
+ bos_token="BOS",
+ )
+ assert get_tokenizer_with_bos(tokenizer) is tokenizer
+
+
+def test_public_support_export():
+ import transformer_lens.tools.analysis as analysis
+ from transformer_lens.tools.analysis.nla import (
+ NaturalLanguageAutoencoder,
+ NLAActivation,
+ )
+
+ assert analysis.NaturalLanguageAutoencoder is NaturalLanguageAutoencoder
+ assert analysis.NLAActivation is NLAActivation
diff --git a/tests/unit/tools/test_nla_reconstructor.py b/tests/unit/tools/test_nla_reconstructor.py
new file mode 100644
index 000000000..bd8b2f2ab
--- /dev/null
+++ b/tests/unit/tools/test_nla_reconstructor.py
@@ -0,0 +1,113 @@
+"""Local prompt, loading-key, and lifecycle contracts for reconstruction."""
+
+import pytest
+
+from transformer_lens.tools.analysis._nla_artifacts import (
+ ArtifactRef,
+ NLAMetadata,
+ NLAPairConfig,
+)
+from transformer_lens.tools.analysis._nla_reconstructor import (
+ ReconstructorRuntime,
+ prepare_critic_tokens,
+ validate_loading_info,
+)
+
+
+@pytest.fixture
+def metadata():
+ return NLAMetadata(
+ "ar",
+ 4,
+ 1,
+ "marker",
+ 7,
+ "{injection_char}",
+ "Summary: {explanation} END",
+ 2.0,
+ None,
+ None,
+ None,
+ (8, 9),
+ )
+
+
+class Tokenizer:
+ bos_token_id = 1
+
+ def __init__(self):
+ self.ids = [1, 3, 8, 9]
+ self.calls = []
+
+ def encode(self, text, *, add_special_tokens):
+ self.calls.append((text, add_special_tokens))
+ return self.ids
+
+
+def test_raw_template_and_authored_whitespace(metadata):
+ tokenizer = Tokenizer()
+ tokens = prepare_critic_tokens(tokenizer, metadata, " authored\n text ", 4)
+ assert tokenizer.calls == [("Summary: authored\n text END", True)]
+ assert tokens.tolist() == [[1, 3, 8, 9]]
+
+
+@pytest.mark.parametrize(
+ "ids,message",
+ [
+ ([3, 4, 8, 9], "BOS"),
+ ([1, 3, 9, 8], "suffix"),
+ ([1, 2, 3, 8, 9], "context"),
+ ([1, True, 8, 9], "integers"),
+ ],
+)
+def test_prompt_guards(metadata, ids, message):
+ tokenizer = Tokenizer()
+ tokenizer.ids = ids
+ with pytest.raises(ValueError, match=message):
+ prepare_critic_tokens(tokenizer, metadata, "text", 4)
+
+
+def test_lazy_runtime_and_cleanup():
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 4,
+ 1,
+ 3,
+ )
+ runtime = ReconstructorRuntime(pair, local_files_only=True)
+ assert runtime._model is None and runtime._head is None
+ with pytest.raises(ValueError, match="nonempty"):
+ runtime.reconstruct("")
+ runtime.unload()
+ runtime.unload()
+ assert runtime._lock.acquire(blocking=False)
+ try:
+ with pytest.raises(RuntimeError, match="concurrent"):
+ runtime.reconstruct("text")
+ with pytest.raises(RuntimeError, match="during"):
+ runtime.unload()
+ finally:
+ runtime._lock.release()
+
+
+def test_expected_missing_output_weights():
+ validate_loading_info({"missing_keys": ["model.norm.weight", "lm_head.weight"]})
+
+
+@pytest.mark.parametrize(
+ "info",
+ [
+ {},
+ {"missing_keys": ["model.norm.weight"]},
+ {"missing_keys": ["model.norm.weight", "lm_head.weight", "model.layers.0.weight"]},
+ {"missing_keys": ["model.norm.weight", "lm_head.weight"], "unexpected_keys": ["other"]},
+ {"missing_keys": ["model.norm.weight", "lm_head.weight"], "mismatched_keys": ["other"]},
+ {"missing_keys": ["model.norm.weight", "lm_head.weight"], "error_msgs": ["error"]},
+ ],
+)
+def test_loading_errors_are_not_suppressed(info):
+ with pytest.raises(ValueError):
+ validate_loading_info(info)
diff --git a/tests/unit/tools/test_nla_target.py b/tests/unit/tools/test_nla_target.py
new file mode 100644
index 000000000..ccf46e7a2
--- /dev/null
+++ b/tests/unit/tools/test_nla_target.py
@@ -0,0 +1,132 @@
+"""Local provenance and token-preparation checks without model loading."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig
+from transformer_lens.tools.analysis._nla_contracts import SourceProvenance
+from transformer_lens.tools.analysis._nla_target import (
+ prepare_tokens,
+ target_provenance,
+)
+
+
+@pytest.fixture
+def pair():
+ return NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 4,
+ 1,
+ 3,
+ )
+
+
+@pytest.fixture
+def model():
+ return SimpleNamespace(
+ original_model=SimpleNamespace(config=SimpleNamespace(vocab_size=10)),
+ cfg=SimpleNamespace(n_ctx=4, device="cpu", default_prepend_bos=True),
+ tokenizer=SimpleNamespace(pad_token_id=9, eos_token_id=8),
+ )
+
+
+def test_matching_external_loading_evidence(model, pair):
+ config = model.original_model.config
+ config._name_or_path = pair.target.repo_id
+ config._commit_hash = pair.target.revision
+ assert target_provenance(model, pair) is SourceProvenance.PINNED_LOAD
+
+
+def test_owned_loading_record(model, pair):
+ model._nla_target_artifact = pair.target
+ assert target_provenance(model, pair) is SourceProvenance.PINNED_LOAD
+
+
+def test_unknown_provenance_requires_attestation(model, pair):
+ with pytest.raises(ValueError, match="unknown"):
+ target_provenance(model, pair)
+ with pytest.warns(UserWarning, match="unverified"):
+ assert (
+ target_provenance(model, pair, allow_unverified_target=True)
+ is SourceProvenance.CALLER_ATTESTED
+ )
+
+
+@pytest.mark.parametrize(
+ "field,value", [("_name_or_path", "example/wrong"), ("_commit_hash", "d" * 40)]
+)
+def test_attestation_cannot_bypass_known_mismatch(model, pair, field, value):
+ setattr(model.original_model.config, field, value)
+ with pytest.raises(ValueError, match="does not match"):
+ target_provenance(model, pair, allow_unverified_target=True)
+
+
+def test_local_unknown_identity_can_be_attested(model, pair, tmp_path):
+ model.original_model.config._name_or_path = str(tmp_path)
+ with pytest.raises(ValueError, match="unknown"):
+ target_provenance(model, pair)
+ with pytest.warns(UserWarning, match="unverified"):
+ assert (
+ target_provenance(model, pair, allow_unverified_target=True)
+ is SourceProvenance.CALLER_ATTESTED
+ )
+
+
+def test_recorded_identity_mismatch(model, pair):
+ model._nla_target_artifact = pair.ar
+ with pytest.raises(ValueError, match="recorded"):
+ target_provenance(model, pair)
+
+
+@pytest.mark.parametrize(
+ "tokens", [torch.tensor([1, 2]), torch.tensor([[1, 2]], dtype=torch.int32)]
+)
+def test_explicit_tokens(model, tokens):
+ result, policy = prepare_tokens(model, tokens)
+ assert result.shape == (1, 2) and result.dtype == torch.int64
+ assert "caller" in policy
+
+
+@pytest.mark.parametrize(
+ "tokens",
+ [
+ torch.tensor([]),
+ torch.ones(2, 2, dtype=torch.int64),
+ torch.ones(1, 1, 2, dtype=torch.int64),
+ torch.tensor([-1]),
+ torch.tensor([10]),
+ torch.tensor([9]),
+ torch.ones(5, dtype=torch.int64),
+ ],
+)
+def test_invalid_tokens(model, tokens):
+ with pytest.raises(ValueError):
+ prepare_tokens(model, tokens)
+
+
+def test_explicit_bos_override_rejected(model):
+ with pytest.raises(ValueError, match="explicit"):
+ prepare_tokens(model, torch.tensor([1]), prepend_bos=True)
+
+
+def test_bos_policy_no_truncation_and_single_tokenization(model):
+ calls = []
+
+ def tokenize(text, *, prepend_bos, truncate):
+ calls.append((text, prepend_bos, truncate))
+ return torch.tensor([[1, 2]])
+
+ model.to_tokens = tokenize
+ tokens, policy = prepare_tokens(model, "hello", prepend_bos=False)
+ assert calls == [("hello", False, False)]
+ assert "prepend_bos=False" in policy and tokens.tolist() == [[1, 2]]
+
+
+def test_empty_prompt(model):
+ with pytest.raises(ValueError, match="nonempty"):
+ prepare_tokens(model, " ")
diff --git a/tests/unit/tools/test_nla_verbalizer.py b/tests/unit/tools/test_nla_verbalizer.py
new file mode 100644
index 000000000..e9ca97ab4
--- /dev/null
+++ b/tests/unit/tools/test_nla_verbalizer.py
@@ -0,0 +1,138 @@
+"""Exact synthetic checks of generation-aware embedding replacement."""
+
+import pytest
+import torch
+
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig
+from transformer_lens.tools.analysis._nla_verbalizer import (
+ VerbalizerRuntime,
+ embedding_overwrite,
+ validate_activation_provenance,
+)
+
+
+@pytest.mark.parametrize("cached", [True, False])
+def test_exact_prefill_locality(cached):
+ vector = torch.tensor([3.0, 4.0])
+ hook = embedding_overwrite(vector, 10.0, 1, 3, cached=cached)
+ value = torch.arange(6.0).reshape(1, 3, 2)
+ original = value.clone()
+ result = hook(value, None)
+ torch.testing.assert_close(result[0, 1], torch.tensor([6.0, 8.0]))
+ torch.testing.assert_close(result[0, 0], original[0, 0])
+ torch.testing.assert_close(result[0, 2], original[0, 2])
+ torch.testing.assert_close(value, original)
+ assert result.data_ptr() != value.data_ptr()
+
+
+def test_cached_one_token_prefill_and_decode():
+ hook = embedding_overwrite(torch.tensor([1.0, 0.0]), 5.0, 0, 1, cached=True)
+ value = torch.ones(1, 1, 2)
+ torch.testing.assert_close(hook(value, None)[0, 0], torch.tensor([5.0, 0.0]))
+ assert hook(value, None) is value
+
+
+def test_uncached_reapplies_at_original_position():
+ hook = embedding_overwrite(torch.tensor([1.0, 0.0]), 5.0, 1, 2, cached=False)
+ hook(torch.ones(1, 2, 2), None)
+ value = torch.ones(1, 4, 2)
+ result = hook(value, None)
+ torch.testing.assert_close(result[0, 1], torch.tensor([5.0, 0.0]))
+ torch.testing.assert_close(result[0, 2:], value[0, 2:])
+
+
+@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf")])
+def test_invalid_scale(scale):
+ with pytest.raises(ValueError):
+ embedding_overwrite(torch.ones(2), scale, 0, 1, cached=True)
+
+
+def test_wrong_prefill_shape():
+ hook = embedding_overwrite(torch.ones(2), 1.0, 0, 2, cached=True)
+ with pytest.raises(ValueError, match="prefill"):
+ hook(torch.ones(1, 1, 2), None)
+
+
+def test_runtime_is_lazy_and_rejects_concurrent_operations():
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 2,
+ 0,
+ 1,
+ )
+ runtime = VerbalizerRuntime(pair, "synthetic", local_files_only=True)
+ assert runtime._model is None
+ runtime.unload()
+ assert runtime._lock.acquire(blocking=False)
+ try:
+ with pytest.raises(RuntimeError, match="concurrent"):
+ runtime.verbalize(torch.ones(2))
+ with pytest.raises(RuntimeError, match="during"):
+ runtime.unload()
+ finally:
+ runtime._lock.release()
+ runtime.unload()
+
+
+@pytest.mark.parametrize("source", ["pinned_load", "caller_attested"])
+@pytest.mark.parametrize("mismatch", ["target", "revision", "layer"])
+def test_structured_provenance_mismatches_rejected(source, mismatch):
+ from dataclasses import replace
+
+ from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ NLAActivation,
+ SourceProvenance,
+ )
+
+ pair = NLAPairConfig(
+ ArtifactRef("example/target", "a" * 40),
+ ArtifactRef("example/av", "b" * 40),
+ ArtifactRef("example/ar", "c" * 40),
+ "qwen2",
+ 2,
+ 1,
+ 3,
+ )
+ provenance = ActivationProvenance(
+ "synthetic", SourceProvenance(source), pair.target, 1, 0, "tokens"
+ )
+ if mismatch == "target":
+ provenance = replace(provenance, target=pair.ar)
+ elif mismatch == "revision":
+ provenance = replace(provenance, target=ArtifactRef(pair.target.repo_id, "d" * 40))
+ else:
+ provenance = replace(provenance, layer=0)
+ activation = NLAActivation(torch.ones(2), provenance)
+ with pytest.raises(ValueError, match="match"):
+ validate_activation_provenance(activation, pair, "synthetic")
+ runtime = VerbalizerRuntime(pair, "synthetic", local_files_only=True)
+ with pytest.raises(ValueError, match="match"):
+ runtime.verbalize(activation)
+ assert runtime._model is None
+
+
+def test_reference_normalization_operation_order():
+ vector = torch.linspace(-1.0, 1.0, 3584)
+ expected = vector / (vector.float().norm().clamp_min(1e-12) / 150).to(vector.dtype)
+ actual = embedding_overwrite(vector, 150.0, 0, 1, cached=True)(torch.zeros(1, 1, 3584), None)[
+ 0, 0
+ ]
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("magnitude", [1e-300, 1e300])
+def test_extreme_input_fallback_is_scale_safe(magnitude):
+ vector = torch.tensor([magnitude, magnitude], dtype=torch.float64)
+ result = embedding_overwrite(vector, 5.0, 0, 1, cached=True)(torch.zeros(1, 1, 2), None)[0, 0]
+ assert torch.isfinite(result).all()
+ assert result.norm().item() == pytest.approx(5.0)
+
+
+def test_injection_representability():
+ hook = embedding_overwrite(torch.ones(2), 1e10, 0, 1, cached=True)
+ with pytest.raises(ValueError, match="representable"):
+ hook(torch.ones(1, 1, 2, dtype=torch.float16), None)
diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py
index 13ac09858..356db15b3 100644
--- a/transformer_lens/tools/analysis/__init__.py
+++ b/transformer_lens/tools/analysis/__init__.py
@@ -30,6 +30,8 @@
the ``logit_readout`` kernel that reads any ``d_model`` vector through the
real final norm, unembedding and adapter output transform, chunked over
the vocabulary with target ranks, top-k and vocabulary subsets.
+ - nla: Raw Qwen activation verbalization, reconstruction, and directional
+ scoring with pinned artifacts and lazy role loading.
- projection_kernel: Basis-invariant subspace overlap and TransformerBridge
attention-head OQ/OK/OV affinity.
- sparse_probing: Leakage-safe k-sparse binary probes over supplied
@@ -95,6 +97,17 @@
logit_lens,
logit_readout,
)
+from transformer_lens.tools.analysis.nla import (
+ ActivationProvenance,
+ ExplanationStatus,
+ NLAActivation,
+ NLAExplanation,
+ NLAReconstruction,
+ NLAReconstructionScore,
+ NLASamplingConfig,
+ NaturalLanguageAutoencoder,
+ SourceProvenance,
+)
from transformer_lens.tools.analysis.projection_kernel import (
AttentionHeadRef,
HeadAffinityPair,
@@ -134,6 +147,7 @@
__all__ = [
"ActivationProjection",
+ "ActivationProvenance",
"AttentionHeadRef",
"AttributionResult",
"BackwardLens",
@@ -145,6 +159,7 @@
"DegenerateDirectionError",
"DirectLogitAttribution",
"EdgeAttributionConfig",
+ "ExplanationStatus",
"FunctionSpec",
"HeadAffinityPair",
"HeadAffinityResult",
@@ -159,12 +174,19 @@
"LogitLensResult",
"LogitReadout",
"LogitSignature",
+ "NLAActivation",
+ "NLAExplanation",
+ "NLAReconstruction",
+ "NLAReconstructionScore",
+ "NLASamplingConfig",
+ "NaturalLanguageAutoencoder",
"Node",
"PatchResult",
"ProjectedFactor",
"ProjectionKernelResult",
"RandomSubspaceReference",
"RankReportRow",
+ "SourceProvenance",
"SparseProbeControl",
"SparseProbeConvergenceError",
"SparseProbeMetrics",
diff --git a/transformer_lens/tools/analysis/_nla.py b/transformer_lens/tools/analysis/_nla.py
new file mode 100644
index 000000000..980fc3383
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla.py
@@ -0,0 +1,219 @@
+"""Private orchestration of local NLA inference and owned role lifecycles."""
+
+from __future__ import annotations
+
+import threading
+from contextlib import contextmanager
+from typing import Iterator, Literal
+
+import torch
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla_artifacts import (
+ NLAPairConfig,
+ acquire_snapshot,
+ load_local_metadata,
+ load_registry,
+ load_snapshot_tokenizer,
+)
+from transformer_lens.tools.analysis._nla_contracts import (
+ NLAActivation,
+ NLAExplanation,
+ NLAReconstruction,
+ NLAReconstructionScore,
+ NLASamplingConfig,
+ score_reconstruction,
+)
+from transformer_lens.tools.analysis._nla_reconstructor import ReconstructorRuntime
+from transformer_lens.tools.analysis._nla_target import extract_activation, load_target
+from transformer_lens.tools.analysis._nla_verbalizer import VerbalizerRuntime
+
+
+class NaturalLanguageAutoencoder:
+ """Compose raw-target extraction, AV text, and AR reconstruction.
+
+ Role models are owned here; returned targets and detached results are owned
+ by the caller. Loading provenance does not establish live-weight integrity.
+ Closing releases role models and permanently closes this instance. Unloading
+ one role permits a later lazy reload. Operations on one instance are serial.
+ """
+
+ def __init__(
+ self,
+ pair: NLAPairConfig,
+ artifact_name: str,
+ *,
+ device: str = "cpu",
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+ ) -> None:
+ if not isinstance(artifact_name, str) or not artifact_name.strip():
+ raise ValueError("artifact_name must be nonempty text")
+ self.pair = pair
+ self.artifact_name = artifact_name
+ self.device = device
+ self.cache_dir = cache_dir
+ self.local_files_only = local_files_only
+ self._av = VerbalizerRuntime(
+ pair,
+ artifact_name,
+ device=device,
+ cache_dir=cache_dir,
+ local_files_only=local_files_only,
+ )
+ self._ar = ReconstructorRuntime(
+ pair, device=device, cache_dir=cache_dir, local_files_only=local_files_only
+ )
+ self._lock = threading.Lock()
+ self._closed = False
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ name: str,
+ *,
+ device: str = "cpu",
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+ ) -> NaturalLanguageAutoencoder:
+ """Resolve a supported pinned pair, validating metadata without model weights."""
+ registry = load_registry()
+ if name not in registry:
+ raise ValueError(f"No verified NLA artifact is registered for {name!r}")
+ pair = registry[name]
+ for role in ("av", "ar"):
+ ref = pair.av if role == "av" else pair.ar
+ directory = acquire_snapshot(
+ ref, cache_dir=cache_dir, local_files_only=local_files_only
+ )
+ tokenizer = load_snapshot_tokenizer(directory)
+ load_local_metadata(directory, pair, role, tokenizer)
+ return cls(
+ pair, name, device=device, cache_dir=cache_dir, local_files_only=local_files_only
+ )
+
+ @contextmanager
+ def _operation(self) -> Iterator[None]:
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("concurrent NLA operations are not supported")
+ try:
+ if self._closed:
+ raise RuntimeError("NLA instance is closed")
+ yield
+ finally:
+ self._lock.release()
+
+ def load_target(self) -> TransformerBridge:
+ """Return a pinned raw target without keeping a strong reference to it."""
+ with self._operation():
+ return load_target(
+ self.pair,
+ device=self.device,
+ cache_dir=self.cache_dir,
+ local_files_only=self.local_files_only,
+ )
+
+ def extract(
+ self,
+ target: TransformerBridge,
+ prompt: str | torch.Tensor,
+ *,
+ position: int = -1,
+ output_device: str = "cpu",
+ prepend_bos: bool | None = None,
+ allow_unverified_target: bool = False,
+ ) -> NLAActivation:
+ """Extract one detached raw activation at the artifact-fixed block."""
+ with self._operation():
+ return extract_activation(
+ target,
+ self.pair,
+ self.artifact_name,
+ prompt,
+ position=position,
+ output_device=output_device,
+ prepend_bos=prepend_bos,
+ allow_unverified_target=allow_unverified_target,
+ )
+
+ def verbalize(
+ self,
+ activation: NLAActivation | torch.Tensor,
+ *,
+ max_new_tokens: int = 512,
+ temperature: float | int = 1.0,
+ do_sample: bool = True,
+ use_past_kv_cache: bool = True,
+ top_p: float | int = 1.0,
+ top_k: int = 0,
+ ) -> NLAExplanation:
+ """Generate canonical inner explanation content and raw continuation data."""
+ sampling = NLASamplingConfig(
+ max_new_tokens, temperature, do_sample, use_past_kv_cache, top_p, top_k
+ )
+ with self._operation():
+ return self._av.verbalize(activation, sampling)
+
+ def explain(
+ self,
+ target: TransformerBridge,
+ prompt: str | torch.Tensor,
+ *,
+ position: int = -1,
+ prepend_bos: bool | None = None,
+ allow_unverified_target: bool = False,
+ sampling: NLASamplingConfig | None = None,
+ ) -> NLAExplanation:
+ """Extract then verbalize, leaving the caller-owned target resident."""
+ with self._operation():
+ activation = extract_activation(
+ target,
+ self.pair,
+ self.artifact_name,
+ prompt,
+ position=position,
+ prepend_bos=prepend_bos,
+ allow_unverified_target=allow_unverified_target,
+ )
+ return self._av.verbalize(activation, sampling)
+
+ def reconstruct(self, text: str, *, output_device: str = "cpu") -> NLAReconstruction:
+ """Run raw AR inference on nonempty untagged explanation content."""
+ with self._operation():
+ return self._ar.reconstruct(text, output_device=output_device)
+
+ @staticmethod
+ def score(
+ original: torch.Tensor | NLAActivation, reconstructed: torch.Tensor | NLAReconstruction
+ ) -> NLAReconstructionScore:
+ """Directional similarity only, not dataset FVE or explanation faithfulness."""
+ return score_reconstruction(original, reconstructed)
+
+ def unload(self, role: Literal["av", "ar"]) -> None:
+ """Release one owned role, permitting lazy reload on subsequent use."""
+ with self._operation():
+ if role == "av":
+ self._av.unload()
+ elif role == "ar":
+ self._ar.unload()
+ else:
+ raise ValueError("role must be av or ar")
+
+ def close(self) -> None:
+ """Release both owned roles; repeated closing is harmless."""
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("cannot close during an NLA operation")
+ try:
+ if not self._closed:
+ self._av.unload()
+ self._ar.unload()
+ self._closed = True
+ finally:
+ self._lock.release()
+
+ def __enter__(self) -> NaturalLanguageAutoencoder:
+ with self._operation():
+ return self
+
+ def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
+ self.close()
diff --git a/transformer_lens/tools/analysis/_nla_artifacts.py b/transformer_lens/tools/analysis/_nla_artifacts.py
new file mode 100644
index 000000000..25ee29d88
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla_artifacts.py
@@ -0,0 +1,373 @@
+"""Private validation of immutable natural-language autoencoder artifacts."""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+from dataclasses import dataclass
+from importlib.resources import files
+from pathlib import Path
+from string import Formatter
+from typing import Any, Literal, Mapping, Protocol, cast, runtime_checkable
+
+import torch
+import yaml
+from huggingface_hub import snapshot_download
+from safetensors import safe_open
+
+from transformer_lens.utilities.hf_utils import call_hf_with_retry
+
+Role = Literal["av", "ar"]
+
+
+@runtime_checkable
+class Tokenizer(Protocol):
+ """Tokenization operations needed to check the artifact contract."""
+
+ def encode(self, text: str, *, add_special_tokens: bool) -> list[int]:
+ ...
+
+ def apply_chat_template(
+ self,
+ conversation: list[dict[str, str]],
+ *,
+ tokenize: bool,
+ add_generation_prompt: bool,
+ ) -> object:
+ ...
+
+
+def _mapping(value: object, name: str) -> Mapping[str, Any]:
+ if not isinstance(value, dict) or not all(isinstance(key, str) for key in value):
+ raise ValueError(f"{name} must be a mapping with string keys")
+ return cast(Mapping[str, Any], value)
+
+
+def _integer(value: object, name: str, *, minimum: int = 0) -> int:
+ if type(value) is not int or value < minimum:
+ raise ValueError(f"{name} must be an integer >= {minimum}")
+ return value
+
+
+def _text(value: object, name: str) -> str:
+ if not isinstance(value, str) or not value.strip():
+ raise ValueError(f"{name} must be nonempty text")
+ return value
+
+
+def _positive(value: object, name: str) -> float:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{name} must be a finite positive number")
+ result = float(value)
+ if not math.isfinite(result) or result <= 0:
+ raise ValueError(f"{name} must be a finite positive number")
+ return result
+
+
+def _ids(value: object, name: str) -> tuple[int, ...]:
+ if isinstance(value, Mapping):
+ value = value.get("input_ids")
+ if not isinstance(value, (list, tuple)) or not value:
+ raise ValueError(f"{name} must contain token ids")
+ return tuple(_integer(item, name) for item in value)
+
+
+def _template(value: object, field: str, name: str) -> str:
+ text = _text(value, name)
+ try:
+ parts = list(Formatter().parse(text))
+ except ValueError as exc:
+ raise ValueError(f"{name} is not a valid template") from exc
+ fields = [key for _, key, _, _ in parts if key is not None]
+ if fields != [field] or any(spec or conversion for _, _, spec, conversion in parts):
+ raise ValueError(f"{name} must contain exactly one {{{field}}} field")
+ return text
+
+
+@dataclass(frozen=True)
+class ArtifactRef:
+ """One Hub repository at an immutable commit revision."""
+
+ repo_id: str
+ revision: str
+
+ def __post_init__(self) -> None:
+ if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*", self.repo_id):
+ raise ValueError("repo_id must be an owner/repository identifier")
+ if not re.fullmatch(r"[0-9a-f]{40}", self.revision):
+ raise ValueError("revision must be a lowercase immutable 40-character commit SHA")
+
+
+@dataclass(frozen=True)
+class NLAPairConfig:
+ """Pinned target and role identities, not proof of numerical compatibility."""
+
+ target: ArtifactRef
+ av: ArtifactRef
+ ar: ArtifactRef
+ model_type: str
+ d_model: int
+ layer: int
+ target_layers: int
+
+ def __post_init__(self) -> None:
+ _text(self.model_type, "model_type")
+ _integer(self.d_model, "d_model", minimum=1)
+ _integer(self.layer, "layer")
+ _integer(self.target_layers, "target_layers", minimum=1)
+ if self.layer >= self.target_layers:
+ raise ValueError("layer must be below target_layers")
+ if len({self.target.repo_id, self.av.repo_id, self.ar.repo_id}) != 3:
+ raise ValueError("target, av, and ar must identify distinct repositories")
+
+
+def parse_registry(value: object) -> dict[str, NLAPairConfig]:
+ """Parse supported entries without network access or model loading."""
+ registry = _mapping(value, "registry")
+ if _integer(registry.get("schema_version"), "registry schema_version") != 1:
+ raise ValueError("unsupported registry schema_version")
+ entries = _mapping(registry.get("models"), "registry models")
+ result = {}
+ for name, raw in entries.items():
+ _text(name, "registry name")
+ entry = _mapping(raw, "registry entry")
+ refs = {}
+ for role in ("target", "av", "ar"):
+ ref = _mapping(entry.get(role), role)
+ refs[role] = ArtifactRef(
+ _text(ref.get("repo_id"), "repo_id"), _text(ref.get("revision"), "revision")
+ )
+ result[name] = NLAPairConfig(
+ refs["target"],
+ refs["av"],
+ refs["ar"],
+ _text(entry.get("model_type"), "model_type"),
+ _integer(entry.get("d_model"), "d_model", minimum=1),
+ _integer(entry.get("layer"), "layer"),
+ _integer(entry.get("target_layers"), "target_layers", minimum=1),
+ )
+ return result
+
+
+def load_registry() -> dict[str, NLAPairConfig]:
+ """Read bundled supported pairs through package resources."""
+ resource = files("transformer_lens.tools.analysis").joinpath("nla_registry.json")
+ return parse_registry(json.loads(resource.read_text(encoding="utf-8")))
+
+
+@dataclass(frozen=True)
+class NLAMetadata:
+ """Validated role-specific metadata and canonical tokenization."""
+
+ role: Role
+ d_model: int
+ layer: int
+ injection_char: str
+ injection_token_id: int
+ actor_template: str
+ critic_template: str
+ mse_scale: float
+ injection_scale: float | None
+ actor_token_ids: tuple[int, ...] | None
+ injection_position: int | None
+ critic_suffix_ids: tuple[int, ...] | None
+
+
+def validate_metadata(
+ value: object,
+ model_config: object,
+ pair: NLAPairConfig,
+ role: Role,
+ tokenizer: Tokenizer,
+) -> NLAMetadata:
+ """Check schema-v2 sidecars against the pinned contract and live tokenizer."""
+ if role not in ("av", "ar"):
+ raise ValueError("role must be av or ar")
+ meta = _mapping(value, "nla_meta")
+ config = _mapping(model_config, "model config")
+ if meta.get("kind") != "nla_model":
+ raise ValueError("kind must be nla_model")
+ if _integer(meta.get("schema_version"), "schema_version") != 2:
+ raise ValueError("unsupported sidecar schema_version")
+ if meta.get("role") != role:
+ raise ValueError("sidecar role does not match requested role")
+ dimension = _integer(meta.get("d_model"), "d_model", minimum=1)
+ layer = _integer(meta.get("extraction_layer_index"), "extraction_layer_index")
+ if dimension != pair.d_model or layer != pair.layer:
+ raise ValueError("sidecar dimension/layer does not match the pinned pair")
+ if config.get("model_type") != pair.model_type:
+ raise ValueError("model_type does not match the pinned pair")
+ if _integer(config.get("hidden_size"), "hidden_size", minimum=1) != dimension:
+ raise ValueError("model hidden_size does not match sidecar")
+ depth = _integer(config.get("num_hidden_layers"), "num_hidden_layers", minimum=1)
+ if depth != (pair.target_layers if role == "av" else layer + 1):
+ raise ValueError("model depth does not match the artifact role")
+ critic = meta.get("critic")
+ if critic is not None:
+ critic_layer = _integer(
+ _mapping(critic, "critic").get("extraction_layer_index"), "critic layer"
+ )
+ if critic_layer != layer:
+ raise ValueError("critic layer does not match extraction_layer_index")
+ extraction = _mapping(meta.get("extraction"), "extraction")
+ if "injection_scale" not in extraction:
+ raise ValueError("extraction must record injection_scale, including null for ar")
+ mse_scale = _positive(extraction.get("mse_scale"), "mse_scale")
+ # Schema-v2 Qwen sidecars serialize sqrt(d) as a float64 value with 16 digits.
+ if not math.isclose(mse_scale, math.sqrt(dimension), rel_tol=1e-14, abs_tol=0):
+ raise ValueError("unsupported mse_scale; expected sqrt(d_model) normalization")
+ tokens = _mapping(meta.get("tokens"), "tokens")
+ if "critic_suffix_ids" not in tokens:
+ raise ValueError("tokens must record critic_suffix_ids, including null for av")
+ char = _text(tokens.get("injection_char"), "injection_char")
+ token_id = _integer(tokens.get("injection_token_id"), "injection_token_id")
+ left = _integer(tokens.get("injection_left_neighbor_id"), "left neighbor")
+ right = _integer(tokens.get("injection_right_neighbor_id"), "right neighbor")
+ live_id = _ids(tokenizer.encode(char, add_special_tokens=False), "live injection ids")
+ if live_id != (token_id,) or token_id == getattr(tokenizer, "unk_token_id", None):
+ raise ValueError("injection token encoding drift or UNK")
+ templates = _mapping(meta.get("prompt_templates"), "prompt_templates")
+ actor_template = _template(templates.get("av"), "injection_char", "av template")
+ critic_template = _template(templates.get("ar"), "explanation", "ar template")
+ scale = None
+ actor_ids = None
+ position = None
+ suffix = None
+ if role == "av":
+ scale = _positive(extraction.get("injection_scale"), "injection_scale")
+ actor_ids = _ids(
+ tokenizer.apply_chat_template(
+ [{"role": "user", "content": actor_template.format(injection_char=char)}],
+ tokenize=True,
+ add_generation_prompt=True,
+ ),
+ "actor token ids",
+ )
+ matches = [i for i, item in enumerate(actor_ids) if item == token_id]
+ if len(matches) != 1:
+ raise ValueError("actor prompt must contain exactly one injection token")
+ position = matches[0]
+ if not 0 < position < len(actor_ids) - 1:
+ raise ValueError("injection token must have both neighbors")
+ if actor_ids[position - 1] != left or actor_ids[position + 1] != right:
+ raise ValueError("actor prompt neighbor token drift")
+ if tokens.get("critic_suffix_ids") is not None:
+ raise ValueError("av critic_suffix_ids must be null")
+ else:
+ if extraction.get("injection_scale") is not None:
+ raise ValueError("ar injection_scale must be null")
+ suffix = _ids(tokens.get("critic_suffix_ids"), "critic_suffix_ids")
+ probe = _ids(
+ tokenizer.encode(critic_template.format(explanation="x"), add_special_tokens=True),
+ "critic probe ids",
+ )
+ if len(probe) < len(suffix) or probe[-len(suffix) :] != suffix:
+ raise ValueError("critic suffix token drift")
+ return NLAMetadata(
+ role,
+ dimension,
+ layer,
+ char,
+ token_id,
+ actor_template,
+ critic_template,
+ mse_scale,
+ scale,
+ actor_ids,
+ position,
+ suffix,
+ )
+
+
+def load_local_metadata(
+ directory: Path, pair: NLAPairConfig, role: Role, tokenizer: Tokenizer
+) -> NLAMetadata:
+ """Validate files from one local snapshot without acquiring any remote files."""
+ try:
+ meta = yaml.safe_load((directory / "nla_meta.yaml").read_text(encoding="utf-8"))
+ config = json.loads((directory / "config.json").read_text(encoding="utf-8"))
+ except (yaml.YAMLError, json.JSONDecodeError) as exc:
+ raise ValueError("invalid artifact metadata syntax") from exc
+ return validate_metadata(meta, config, pair, role, tokenizer)
+
+
+def validate_snapshot_path(directory: Path, ref: ArtifactRef) -> Path:
+ """Require the cache namespace and revision of the requested repository."""
+ expected_namespace = "models--" + ref.repo_id.replace("/", "--")
+ if (
+ directory.name != ref.revision
+ or directory.parent.name != "snapshots"
+ or directory.parent.parent.name != expected_namespace
+ ):
+ raise ValueError("Hub snapshot path does not match the pinned repository and revision")
+ return directory
+
+
+def acquire_snapshot(
+ ref: ArtifactRef,
+ *,
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+ include_weights: bool = False,
+) -> Path:
+ """Acquire one pinned snapshot, excluding large model weights by default."""
+ patterns = [
+ "nla_meta.yaml",
+ "config.json",
+ "tokenizer*",
+ "special_tokens_map.json",
+ "added_tokens.json",
+ "vocab.json",
+ "merges.txt",
+ "chat_template*",
+ "generation_config.json",
+ ]
+ if include_weights:
+ patterns.extend(["*.safetensors", "*.safetensors.index.json"])
+ directory = call_hf_with_retry(
+ snapshot_download,
+ repo_id=ref.repo_id,
+ revision=ref.revision,
+ cache_dir=cache_dir,
+ local_files_only=local_files_only,
+ allow_patterns=patterns,
+ )
+ return validate_snapshot_path(Path(directory), ref)
+
+
+def load_snapshot_tokenizer(directory: Path) -> Tokenizer:
+ """Load only local tokenizer assets, with remote code disabled."""
+ from transformers import AutoTokenizer
+
+ return cast(
+ Tokenizer,
+ AutoTokenizer.from_pretrained(
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ ),
+ )
+
+
+def validate_value_head(state: Mapping[str, torch.Tensor], d_model: int) -> None:
+ """Require the released bias-free Linear(d,d) state contract."""
+ _integer(d_model, "d_model", minimum=1)
+ if set(state) != {"weight"}:
+ raise ValueError("value head must contain only weight, with no bias")
+ weight = state["weight"]
+ if not isinstance(weight, torch.Tensor) or tuple(weight.shape) != (d_model, d_model):
+ raise ValueError("value head weight must have shape (d_model, d_model)")
+ if not weight.is_floating_point() or not torch.isfinite(weight).all().item():
+ raise ValueError("value head weight must be finite floating point")
+
+
+def validate_value_head_file(path: Path, d_model: int) -> None:
+ """Inspect safetensors keys and shape without materializing a full backbone."""
+ _integer(d_model, "d_model", minimum=1)
+ with safe_open(path, framework="pt", device="cpu") as handle:
+ if set(handle.keys()) != {"weight"}:
+ raise ValueError("value head must contain only weight, with no bias")
+ if handle.get_slice("weight").get_shape() != [d_model, d_model]:
+ raise ValueError("value head weight must have shape (d_model, d_model)")
+ validate_value_head({"weight": handle.get_tensor("weight")}, d_model)
diff --git a/transformer_lens/tools/analysis/_nla_contracts.py b/transformer_lens/tools/analysis/_nla_contracts.py
new file mode 100644
index 000000000..c320dc601
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla_contracts.py
@@ -0,0 +1,270 @@
+"""Private result, text, and numerical contracts for NLA inference."""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+
+from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, _integer
+
+
+class SourceProvenance(str, Enum):
+ """Loading evidence, not continuous verification of live model weights."""
+
+ PINNED_LOAD = "pinned_load"
+ CALLER_ATTESTED = "caller_attested"
+ UNKNOWN_VECTOR = "unknown_vector"
+
+
+@dataclass(frozen=True)
+class ActivationProvenance:
+ """Source identity and extraction conventions without owning a model."""
+
+ artifact_name: str
+ source: SourceProvenance
+ target: ArtifactRef | None = None
+ layer: int | None = None
+ position: int | None = None
+ tokenization_policy: str | None = None
+ basis: str = "raw_hf"
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.artifact_name, str) or not self.artifact_name.strip():
+ raise ValueError("artifact_name must be nonempty text")
+ if not isinstance(self.source, SourceProvenance):
+ raise ValueError("source must be a SourceProvenance value")
+ if self.basis != "raw_hf":
+ raise ValueError("only the raw_hf residual basis is supported")
+ for name, value in (("layer", self.layer), ("position", self.position)):
+ if value is not None:
+ _integer(value, name)
+ if self.tokenization_policy is not None and (
+ not isinstance(self.tokenization_policy, str) or not self.tokenization_policy.strip()
+ ):
+ raise ValueError("tokenization_policy must be nonempty text")
+ if self.source is SourceProvenance.PINNED_LOAD and (
+ self.target is None
+ or self.layer is None
+ or self.position is None
+ or self.tokenization_policy is None
+ ):
+ raise ValueError(
+ "pinned loading provenance requires identity and extraction conventions"
+ )
+
+ @property
+ def source_verified(self) -> bool:
+ """Whether immutable loading provenance was recorded, not weight integrity."""
+ return self.source is SourceProvenance.PINNED_LOAD
+
+
+def detached_vector(vector: torch.Tensor) -> torch.Tensor:
+ """Validate and copy one real activation without retaining its graph or storage."""
+ if vector.ndim != 1 or vector.numel() == 0 or not vector.is_floating_point():
+ raise ValueError("activation must be a nonempty one-dimensional floating-point tensor")
+ if not torch.isfinite(vector).all().item():
+ raise ValueError("activation must contain only finite values")
+ if not torch.any(vector != 0).item():
+ raise ValueError("activation must be nonzero")
+ return vector.detach().clone()
+
+
+@dataclass(frozen=True)
+class NLAActivation:
+ """Owned detached activation vector and source loading evidence."""
+
+ vector: torch.Tensor
+ provenance: ActivationProvenance
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "vector", detached_vector(self.vector))
+
+
+@dataclass(frozen=True)
+class NLAReconstruction:
+ """Owned detached AR vector and immutable reconstructor identity."""
+
+ vector: torch.Tensor
+ artifact: ArtifactRef
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "vector", detached_vector(self.vector))
+
+
+class ExplanationStatus(str, Enum):
+ """Observable outcomes of canonical explanation-region parsing."""
+
+ SUCCESS = "success"
+ MISSING = "missing"
+ MULTIPLE = "multiple"
+ NESTED = "nested"
+ INCOMPLETE = "incomplete"
+ EMPTY = "empty"
+ MALFORMED = "malformed"
+
+
+@dataclass(frozen=True)
+class NLAExplanation:
+ """Parsed inner text and untouched continuation data, including failures."""
+
+ text: str | None
+ raw_text: str
+ token_ids: tuple[int, ...]
+ parse_status: ExplanationStatus
+ provenance: ActivationProvenance
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.raw_text, str):
+ raise ValueError("raw_text must be text")
+ if not isinstance(self.parse_status, ExplanationStatus):
+ raise ValueError("parse_status must be an ExplanationStatus value")
+ ids = tuple(_integer(item, "continuation token id") for item in self.token_ids)
+ object.__setattr__(self, "token_ids", ids)
+ if self.parse_status is ExplanationStatus.SUCCESS:
+ text = validate_reconstruction_text(self.text)
+ if text != text.strip():
+ raise ValueError("parsed explanation text must have stripped outer whitespace")
+ elif self.text is not None:
+ raise ValueError("failed explanations must have text=None")
+
+
+# Recognize broken wrappers too, rather than falling back to the full response.
+_WRAPPER = re.compile(r"<\s*/?\s*explanation\b", re.IGNORECASE)
+_TAG = re.compile(r"")
+
+
+def validate_reconstruction_text(text: object) -> str:
+ """Accept nonempty untagged content without guessing its semantic origin."""
+ if not isinstance(text, str) or not text.strip():
+ raise ValueError("reconstruction text must be nonempty explanation content")
+ if _WRAPPER.search(text):
+ raise ValueError("reconstruction text must not contain explanation wrappers")
+ return text
+
+
+def parse_explanation(
+ raw_text: str, token_ids: tuple[int, ...], provenance: ActivationProvenance
+) -> NLAExplanation:
+ """Parse one exact region, preserving raw continuation on every outcome."""
+ tags = list(_TAG.finditer(raw_text))
+ wrappers = list(_WRAPPER.finditer(raw_text))
+ text = None
+ if not wrappers:
+ status = ExplanationStatus.MISSING
+ elif len(wrappers) != len(tags):
+ status = ExplanationStatus.MALFORMED
+ else:
+ depth = 0
+ regions = 0
+ malformed_order = False
+ nested = False
+ for tag in tags:
+ if tag.group() == "":
+ nested |= depth > 0
+ depth += 1
+ elif depth == 0:
+ malformed_order = True
+ else:
+ depth -= 1
+ if depth == 0:
+ regions += 1
+ if nested:
+ status = ExplanationStatus.NESTED
+ elif malformed_order or depth:
+ status = ExplanationStatus.INCOMPLETE
+ elif regions != 1:
+ status = ExplanationStatus.MULTIPLE
+ else:
+ text = raw_text[tags[0].end() : tags[1].start()].strip()
+ status = ExplanationStatus.SUCCESS if text else ExplanationStatus.EMPTY
+ if not text:
+ text = None
+ return NLAExplanation(text, raw_text, token_ids, status, provenance)
+
+
+@dataclass(frozen=True)
+class NLAReconstructionScore:
+ """Directional diagnostics, not raw MSE, dataset FVE, or text faithfulness."""
+
+ cosine: float
+ normalized_mse: float
+
+
+def score_reconstruction(
+ original: torch.Tensor | NLAActivation, reconstructed: torch.Tensor | NLAReconstruction
+) -> NLAReconstructionScore:
+ """Compute 2*(1-cosine) using detached, scale-safe CPU float64 accumulation."""
+ a = detached_vector(original.vector if isinstance(original, NLAActivation) else original)
+ b = detached_vector(
+ reconstructed.vector if isinstance(reconstructed, NLAReconstruction) else reconstructed
+ )
+ if a.shape != b.shape:
+ raise ValueError("original and reconstructed dimensions must match")
+ a = a.to(device="cpu", dtype=torch.float64)
+ b = b.to(device="cpu", dtype=torch.float64)
+ # Rescale before taking norms to avoid overflow and underflow at finite extremes.
+ a = a / a.abs().max()
+ b = b / b.abs().max()
+ cosine = float(torch.dot(a / a.norm(), b / b.norm()).item())
+ # Clamp only roundoff-sized overshoot; larger violations indicate a defect.
+ if not math.isfinite(cosine) or abs(cosine) > 1 + 1e-12:
+ raise ValueError("cosine outside its numerical range")
+ cosine = max(-1.0, min(1.0, cosine))
+ return NLAReconstructionScore(cosine, 2 * (1 - cosine))
+
+
+def _number(value: object, name: str) -> float:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{name} must be a finite number")
+ try:
+ result = float(value)
+ except OverflowError as exc:
+ raise ValueError(f"{name} must be a finite number") from exc
+ if not math.isfinite(result):
+ raise ValueError(f"{name} must be a finite number")
+ return result
+
+
+@dataclass(frozen=True)
+class NLASamplingConfig:
+ """Bounded AV generation policy with explicit sampling versus greedy mode."""
+
+ max_new_tokens: int = 512
+ temperature: float | int = 1.0
+ do_sample: bool = True
+ use_past_kv_cache: bool = True
+ top_p: float | int = 1.0
+ top_k: int = 0
+
+ def __post_init__(self) -> None:
+ _integer(self.max_new_tokens, "max_new_tokens", minimum=1)
+ _integer(self.top_k, "top_k")
+ if type(self.do_sample) is not bool or type(self.use_past_kv_cache) is not bool:
+ raise ValueError("sampling and caching flags must be bool")
+ temperature = _number(self.temperature, "temperature")
+ top_p = _number(self.top_p, "top_p")
+ if temperature < 0 or (self.do_sample and temperature == 0):
+ raise ValueError("temperature must be positive for sampling, nonnegative for greedy")
+ if not 0 < top_p <= 1:
+ raise ValueError("top_p must be in (0, 1]")
+ if not self.do_sample and (top_p != 1 or self.top_k != 0):
+ raise ValueError("top_p and top_k filtering require sampling")
+
+ def validate_prompt(self, prompt_length: int, context_length: int) -> None:
+ """Reject prompt plus generation budgets exceeding the model context."""
+ validate_context_length(prompt_length, context_length, self.max_new_tokens)
+
+
+def validate_context_length(
+ token_count: int, context_length: int, generation_budget: int = 0
+) -> None:
+ """Check exact tokenized lengths without truncating prompts or budgets."""
+ _integer(token_count, "token_count", minimum=1)
+ _integer(context_length, "context_length", minimum=1)
+ _integer(generation_budget, "generation_budget")
+ if token_count + generation_budget > context_length:
+ raise ValueError("tokenized prompt plus generation budget exceeds model context length")
diff --git a/transformer_lens/tools/analysis/_nla_reconstructor.py b/transformer_lens/tools/analysis/_nla_reconstructor.py
new file mode 100644
index 000000000..3f357f4b5
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla_reconstructor.py
@@ -0,0 +1,215 @@
+"""Private truncated Qwen reconstructor execution with an independent value head."""
+
+from __future__ import annotations
+
+import json
+import threading
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoModelForCausalLM
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla_artifacts import (
+ NLAMetadata,
+ NLAPairConfig,
+ acquire_snapshot,
+ load_local_metadata,
+ load_snapshot_tokenizer,
+ validate_value_head,
+ validate_value_head_file,
+)
+from transformer_lens.tools.analysis._nla_contracts import (
+ NLAReconstruction,
+ validate_context_length,
+ validate_reconstruction_text,
+)
+from transformer_lens.utilities.hf_utils import call_hf_with_retry
+
+_EXPECTED_MISSING = frozenset({"model.norm.weight", "lm_head.weight"})
+
+
+def validate_loading_info(info: dict[str, Any]) -> None:
+ """Accept only the released Qwen AR's omitted final output weights."""
+ if set(info.get("missing_keys", ())) != _EXPECTED_MISSING:
+ raise ValueError("AR missing keys must be exactly final norm and LM head weights")
+ for key in ("unexpected_keys", "mismatched_keys", "error_msgs"):
+ if info.get(key):
+ raise ValueError(f"AR loading reported {key}")
+
+
+def prepare_critic_tokens(
+ tokenizer: Any, metadata: NLAMetadata, text: str, context: int
+) -> torch.Tensor:
+ """Render one raw prompt with special tokens, checking BOS and suffix."""
+ content = validate_reconstruction_text(text)
+ if metadata.role != "ar" or not metadata.critic_suffix_ids:
+ raise ValueError("reconstruction requires validated AR metadata")
+ prompt = metadata.critic_template.format(explanation=content)
+ ids = tokenizer.encode(prompt, add_special_tokens=True)
+ validate_context_length(len(ids), context)
+ if any(type(item) is not int or item < 0 for item in ids):
+ raise ValueError("critic token ids must be nonnegative integers")
+ suffix = metadata.critic_suffix_ids
+ if tuple(ids[-len(suffix) :]) != suffix:
+ raise ValueError("critic suffix token drift")
+ bos = tokenizer.bos_token_id
+ if bos is not None and ids[0] != bos:
+ raise ValueError("critic special-token policy must preserve BOS")
+ return torch.tensor([ids], dtype=torch.int64)
+
+
+def load_critic_model(
+ directory: Path, pair: NLAPairConfig, tokenizer: Any, device: str
+) -> TransformerBridge:
+ """Load an owned Qwen AR for raw backbone execution without its LM head."""
+ if pair.model_type != "qwen2":
+ raise ValueError("AR loading currently supports qwen2 only")
+ config = json.loads((directory / "config.json").read_text(encoding="utf-8"))
+ if config.get("tie_word_embeddings") is not False:
+ raise ValueError("AR loading requires an independent omitted LM head")
+ loaded: Any = call_hf_with_retry(
+ AutoModelForCausalLM.from_pretrained,
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ output_loading_info=True,
+ )
+ original, info = loaded
+ validate_loading_info(info)
+ if (
+ original.config.num_hidden_layers != pair.layer + 1
+ or original.config.hidden_size != pair.d_model
+ ):
+ raise ValueError("AR backbone dimension/depth mismatch")
+ model = TransformerBridge.boot_transformers(
+ pair.ar.repo_id,
+ hf_model=original,
+ tokenizer=tokenizer,
+ device=device,
+ dtype=torch.float32,
+ trust_remote_code=False,
+ )
+ # Bridge setup inspects output weights; bypass normalization only after wrapping.
+ original.model.norm = torch.nn.Identity()
+ model.eval()
+ return model
+
+
+def reconstruct_vector(
+ model: TransformerBridge,
+ head: torch.nn.Linear,
+ tokens: torch.Tensor,
+ metadata: NLAMetadata,
+ pair: NLAPairConfig,
+ *,
+ output_device: str = "cpu",
+) -> NLAReconstruction:
+ """Read the raw last-block final token, never a normalized output vector."""
+ if metadata.role != "ar" or metadata.layer != pair.layer or metadata.d_model != pair.d_model:
+ raise ValueError("AR metadata does not match the artifact")
+ if model.compatibility_mode or getattr(model, "_weights_processed", False):
+ raise ValueError("AR requires raw unprocessed Bridge weights")
+ if model.training or model.original_model.training or head.training:
+ raise ValueError("AR and value head must be in eval mode")
+ if model.cfg.n_layers != pair.layer + 1 or model.cfg.d_model != pair.d_model:
+ raise ValueError("AR model dimensions/depth do not match the artifact")
+ if head.bias is not None:
+ raise ValueError("AR value head must be bias-free")
+ if tokens.ndim != 2 or tokens.shape[0] != 1 or tokens.dtype != torch.int64:
+ raise ValueError("AR requires one int64 token sequence")
+ validate_context_length(tokens.shape[1], model.cfg.n_ctx)
+ validate_value_head(head.state_dict(), pair.d_model)
+ hook = f"blocks.{pair.layer}.hook_out"
+ captured: list[torch.Tensor] = []
+
+ def capture(value: torch.Tensor, hook: Any) -> torch.Tensor:
+ captured.append(value)
+ return value
+
+ original: Any = model.original_model
+ backbone = original.base_model
+ with torch.inference_mode(), model.hooks(fwd_hooks=[(hook, capture)]):
+ backbone(input_ids=tokens.to(model.cfg.device), use_cache=False)
+ if len(captured) != 1:
+ raise ValueError("AR block output must fire exactly once")
+ residual = captured[0]
+ if residual.shape != (1, tokens.shape[1], pair.d_model):
+ raise ValueError("AR block output must be a single residual stream")
+ head.to(device=residual.device, dtype=residual.dtype)
+ vector = head(residual[0, -1]).detach().to(output_device).clone()
+ del captured, residual
+ return NLAReconstruction(vector, pair.ar)
+
+
+class ReconstructorRuntime:
+ """Lazy owned AR/head lifecycle; concurrent operations are rejected."""
+
+ def __init__(
+ self,
+ pair: NLAPairConfig,
+ *,
+ device: str = "cpu",
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+ ) -> None:
+ self.pair = pair
+ self.device = device
+ self.cache_dir = cache_dir
+ self.local_files_only = local_files_only
+ self._model: TransformerBridge | None = None
+ self._head: torch.nn.Linear | None = None
+ self._lock = threading.Lock()
+
+ def reconstruct(self, text: str, *, output_device: str = "cpu") -> NLAReconstruction:
+ """Validate prompt and small assets before acquiring full backbone weights."""
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("concurrent reconstructor operations are not supported")
+ try:
+ validate_reconstruction_text(text)
+ directory = acquire_snapshot(
+ self.pair.ar, cache_dir=self.cache_dir, local_files_only=self.local_files_only
+ )
+ tokenizer = load_snapshot_tokenizer(directory)
+ metadata = load_local_metadata(directory, self.pair, "ar", tokenizer)
+ config = json.loads((directory / "config.json").read_text(encoding="utf-8"))
+ tokens = prepare_critic_tokens(
+ tokenizer, metadata, text, config["max_position_embeddings"]
+ )
+ if self._model is None:
+ snapshot = acquire_snapshot(
+ self.pair.ar,
+ cache_dir=self.cache_dir,
+ local_files_only=self.local_files_only,
+ include_weights=True,
+ )
+ if snapshot != directory:
+ raise ValueError("AR weights must use the metadata snapshot")
+ path = directory / "value_head.safetensors"
+ validate_value_head_file(path, self.pair.d_model)
+ head = torch.nn.Linear(self.pair.d_model, self.pair.d_model, bias=False)
+ head.load_state_dict(load_file(str(path), device="cpu"), strict=True)
+ head.eval()
+ model = load_critic_model(directory, self.pair, tokenizer, self.device)
+ self._model, self._head = model, head
+ if self._head is None:
+ raise RuntimeError("AR value head is unavailable")
+ return reconstruct_vector(
+ self._model, self._head, tokens, metadata, self.pair, output_device=output_device
+ )
+ finally:
+ self._lock.release()
+
+ def unload(self) -> None:
+ """Drop owned backbone/head references without freeing caller tensors."""
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("cannot unload during a reconstructor operation")
+ try:
+ self._model = None
+ self._head = None
+ finally:
+ self._lock.release()
diff --git a/transformer_lens/tools/analysis/_nla_target.py b/transformer_lens/tools/analysis/_nla_target.py
new file mode 100644
index 000000000..d9584d60d
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla_target.py
@@ -0,0 +1,229 @@
+"""Private raw-target loading and single-token NLA activation extraction."""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla_artifacts import (
+ NLAPairConfig,
+ _integer,
+ acquire_snapshot,
+ load_snapshot_tokenizer,
+)
+from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ NLAActivation,
+ SourceProvenance,
+ validate_context_length,
+)
+
+
+def validate_target_structure(model: Any, pair: NLAPairConfig) -> None:
+ """Require raw single-stream causal Bridge weights on one fp32 device."""
+ if not isinstance(model, TransformerBridge):
+ raise TypeError("NLA extraction requires a TransformerBridge")
+ if getattr(model, "compatibility_mode", False) or getattr(model, "_weights_processed", False):
+ raise ValueError("NLA extraction requires the raw unprocessed residual basis")
+ original = model.original_model
+ config = original.config
+ if getattr(config, "model_type", None) != pair.model_type:
+ raise ValueError("target architecture does not match the artifact")
+ if model.cfg.d_model != pair.d_model or model.cfg.n_layers != pair.target_layers:
+ raise ValueError("target dimensions/depth do not match the artifact")
+ if not model.adapter.supports_generation or getattr(config, "is_encoder_decoder", False):
+ raise ValueError("NLA extraction requires a causal decoder-only model")
+ if getattr(model.cfg, "attention_dir", "causal") != "causal":
+ raise ValueError("NLA extraction requires causal attention")
+ if getattr(model.cfg, "total_ut_steps", 1) != 1:
+ raise ValueError("NLA extraction requires non-recurrent block execution")
+ if getattr(original, "is_quantized", False) or getattr(config, "quantization_config", None):
+ raise ValueError("quantized targets are not supported")
+ if getattr(config, "_attn_implementation", None) != "eager":
+ raise ValueError("NLA target extraction requires eager attention")
+ if model.training or original.training:
+ raise ValueError("NLA target extraction requires eval mode")
+ parameters = list(original.parameters())
+ if not parameters or any(
+ p.is_floating_point() and p.dtype != torch.float32 for p in parameters
+ ):
+ raise ValueError("NLA target extraction currently requires float32 weights")
+ if len({p.device for p in parameters}) != 1 or any(p.device.type == "meta" for p in parameters):
+ raise ValueError("NLA target extraction requires materialized single-device weights")
+ mapping = model.adapter.get_component_mapping()
+ block = mapping.get("blocks")
+ if not getattr(block, "hook_out_is_single_residual_stream", False):
+ raise ValueError("target blocks must expose a single residual stream")
+ if f"blocks.{pair.layer}.hook_out" not in model.hook_dict:
+ raise ValueError("artifact block output hook is unavailable")
+
+
+def target_provenance(
+ model: Any, pair: NLAPairConfig, *, allow_unverified_target: bool = False
+) -> SourceProvenance:
+ """Check available loading identity, without claiming live-weight integrity."""
+ if type(allow_unverified_target) is not bool:
+ raise ValueError("allow_unverified_target must be bool")
+ config = model.original_model.config
+ identity = getattr(config, "_name_or_path", None)
+ revision = getattr(config, "_commit_hash", None)
+ recorded = getattr(model, "_nla_target_artifact", None)
+ local_identity = isinstance(identity, str) and (
+ Path(identity).is_absolute()
+ or identity.startswith(("./", "../"))
+ or Path(identity).exists()
+ )
+ if identity and not local_identity and identity != pair.target.repo_id:
+ raise ValueError("target repository identity does not match the pinned artifact")
+ if revision and revision != pair.target.revision:
+ raise ValueError("target revision does not match the pinned artifact")
+ if recorded is not None and recorded != pair.target:
+ raise ValueError("recorded target loading identity does not match the artifact")
+ if recorded == pair.target or (
+ identity == pair.target.repo_id and revision == pair.target.revision
+ ):
+ return SourceProvenance.PINNED_LOAD
+ if not allow_unverified_target:
+ raise ValueError("target loading provenance is unknown; explicit attestation is required")
+ warnings.warn(
+ "Target loading provenance is unverified; results are outside verified artifact support",
+ UserWarning,
+ stacklevel=2,
+ )
+ return SourceProvenance.CALLER_ATTESTED
+
+
+def load_target(
+ pair: NLAPairConfig,
+ *,
+ device: str = "cpu",
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+) -> TransformerBridge:
+ """Load pinned fp32 raw weights with eager attention and caller ownership."""
+ if cache_dir is not None or local_files_only:
+ from transformers import AutoModelForCausalLM
+
+ directory = acquire_snapshot(
+ pair.target,
+ cache_dir=cache_dir,
+ local_files_only=local_files_only,
+ include_weights=True,
+ )
+ tokenizer = load_snapshot_tokenizer(directory)
+ original = AutoModelForCausalLM.from_pretrained(
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ )
+ model = TransformerBridge.boot_transformers(
+ pair.target.repo_id,
+ hf_model=original,
+ tokenizer=tokenizer,
+ device=device,
+ dtype=torch.float32,
+ trust_remote_code=False,
+ )
+ model.eval()
+ validate_target_structure(model, pair)
+ model.__dict__["_nla_target_artifact"] = pair.target
+ return model
+ model = TransformerBridge.boot_transformers(
+ pair.target.repo_id,
+ revision=pair.target.revision,
+ device=device,
+ dtype=torch.float32,
+ trust_remote_code=False,
+ )
+ model.eval()
+ validate_target_structure(model, pair)
+ # Immutable loading evidence contains no reference back to the model.
+ model.__dict__["_nla_target_artifact"] = pair.target
+ return model
+
+
+def prepare_tokens(
+ model: Any, prompt: str | torch.Tensor, *, prepend_bos: bool | None = None
+) -> tuple[torch.Tensor, str]:
+ """Tokenize exactly once or accept one explicit unpadded integer sequence."""
+ if prepend_bos is not None and type(prepend_bos) is not bool:
+ raise ValueError("prepend_bos must be bool or None")
+ if isinstance(prompt, str):
+ if not prompt.strip():
+ raise ValueError("prompt must be nonempty text")
+ tokens = model.to_tokens(prompt, prepend_bos=prepend_bos, truncate=False)
+ effective_bos = model.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
+ policy = f"Bridge string tokenization; prepend_bos={effective_bos}; no chat template"
+ else:
+ if prepend_bos is not None:
+ raise ValueError("prepend_bos cannot modify explicit token input")
+ if prompt.dtype not in (torch.int32, torch.int64):
+ raise ValueError("explicit tokens must be int32 or int64")
+ if prompt.ndim == 1:
+ tokens = prompt.unsqueeze(0)
+ elif prompt.ndim == 2 and prompt.shape[0] == 1:
+ tokens = prompt
+ else:
+ raise ValueError("only one unpadded token sequence is supported")
+ policy = "Explicit tokens; BOS supplied by caller"
+ validate_context_length(tokens.shape[-1], model.cfg.n_ctx)
+ if tokens.ndim != 2 or tokens.shape[0] != 1:
+ raise ValueError("only one token sequence is supported")
+ vocab_size = model.original_model.config.vocab_size
+ if torch.any(tokens < 0).item() or torch.any(tokens >= vocab_size).item():
+ raise ValueError("token ids must be within the model vocabulary")
+ tokenizer = model.tokenizer
+ pad = getattr(tokenizer, "pad_token_id", None)
+ eos = getattr(tokenizer, "eos_token_id", None)
+ # When pad aliases EOS, a token id alone cannot distinguish padding from content.
+ if pad is not None and pad != eos and torch.any(tokens == pad).item():
+ raise ValueError("padding tokens are not supported")
+ return tokens.to(device=model.cfg.device, dtype=torch.int64), policy
+
+
+def extract_activation(
+ model: TransformerBridge,
+ pair: NLAPairConfig,
+ artifact_name: str,
+ prompt: str | torch.Tensor,
+ *,
+ position: int = -1,
+ output_device: str = "cpu",
+ prepend_bos: bool | None = None,
+ allow_unverified_target: bool = False,
+) -> NLAActivation:
+ """Copy one raw block-output vector, releasing the full temporary cache."""
+ validate_target_structure(model, pair)
+ source = target_provenance(model, pair, allow_unverified_target=allow_unverified_target)
+ tokens, policy = prepare_tokens(model, prompt, prepend_bos=prepend_bos)
+ if type(position) is not int:
+ raise ValueError("position must be an integer")
+ resolved = position if position >= 0 else tokens.shape[1] + position
+ if not 0 <= resolved < tokens.shape[1]:
+ raise ValueError("position is outside the original prompt")
+ _integer(pair.layer, "layer")
+ hook = f"blocks.{pair.layer}.hook_out"
+ with torch.no_grad():
+ _, cache = model.run_with_cache(
+ tokens, names_filter=hook, return_cache_object=False, return_type=None
+ )
+ residual = cache[hook]
+ if residual.shape != (1, tokens.shape[1], pair.d_model):
+ raise ValueError("block output is not the expected single residual stream")
+ vector = residual[0, resolved].detach().to(output_device).clone()
+ del cache, residual
+ provenance = ActivationProvenance(
+ artifact_name,
+ source,
+ pair.target if source is SourceProvenance.PINNED_LOAD else None,
+ pair.layer,
+ resolved,
+ policy,
+ )
+ return NLAActivation(vector, provenance)
diff --git a/transformer_lens/tools/analysis/_nla_verbalizer.py b/transformer_lens/tools/analysis/_nla_verbalizer.py
new file mode 100644
index 000000000..d4042a41e
--- /dev/null
+++ b/transformer_lens/tools/analysis/_nla_verbalizer.py
@@ -0,0 +1,248 @@
+"""Private generation-aware activation verbalizer execution."""
+
+from __future__ import annotations
+
+import json
+import threading
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+from transformers import AutoModelForCausalLM
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._nla_artifacts import (
+ NLAMetadata,
+ NLAPairConfig,
+ acquire_snapshot,
+ load_local_metadata,
+ load_snapshot_tokenizer,
+)
+from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ NLAActivation,
+ NLAExplanation,
+ NLASamplingConfig,
+ SourceProvenance,
+ detached_vector,
+ parse_explanation,
+)
+from transformer_lens.utilities.hf_utils import call_hf_with_retry
+
+
+def embedding_overwrite(
+ vector: torch.Tensor, scale: float, position: int, prompt_length: int, *, cached: bool
+) -> Callable[[torch.Tensor, Any], torch.Tensor]:
+ """Construct per-generation state, including one-token prefill handling."""
+ source = detached_vector(vector).to(device="cpu", dtype=torch.float64)
+ if not torch.isfinite(torch.tensor(scale, dtype=torch.float64)).item() or scale <= 0:
+ raise ValueError("injection scale must be finite and positive")
+ if (
+ type(position) is not int
+ or type(prompt_length) is not int
+ or not 0 <= position < prompt_length
+ ):
+ raise ValueError("injection position must be within the original prompt")
+ if type(cached) is not bool:
+ raise ValueError("cached must be bool")
+ reference_source = source.float()
+ reference_norm = reference_source.norm()
+ if torch.isfinite(reference_norm).item() and reference_norm.item() >= 1e-12:
+ # Match the released client's fp32 norm and division order on its input domain.
+ injected = reference_source / (reference_norm / scale)
+ else:
+ unit = source / source.abs().max()
+ injected = unit / unit.norm() * scale
+ first_call = True
+
+ def overwrite(value: torch.Tensor, hook: Any) -> torch.Tensor:
+ nonlocal first_call
+ if cached and not first_call:
+ return value
+ if value.ndim != 3 or value.shape[0] != 1 or value.shape[2] != source.numel():
+ raise ValueError("embedding hook must expose one [batch, position, d_model] stream")
+ if value.shape[1] < prompt_length or (first_call and value.shape[1] != prompt_length):
+ raise ValueError("embedding prefill length does not match the canonical prompt")
+ replacement = injected.to(device=value.device, dtype=value.dtype)
+ if not torch.isfinite(replacement).all().item():
+ raise ValueError("injection scale is not representable in embedding dtype")
+ output = value.clone()
+ output[0, position] = replacement
+ first_call = False
+ return output
+
+ return overwrite
+
+
+def generate_explanation(
+ model: TransformerBridge,
+ metadata: NLAMetadata,
+ activation: NLAActivation,
+ sampling: NLASamplingConfig,
+) -> NLAExplanation:
+ """Execute already validated canonical tokens with scoped embedding injection."""
+ if metadata.role != "av" or metadata.actor_token_ids is None:
+ raise ValueError("verbalization requires validated actor metadata")
+ if metadata.injection_position is None or metadata.injection_scale is None:
+ raise ValueError("actor metadata is missing its injection contract")
+ if activation.vector.numel() != metadata.d_model or model.cfg.d_model != metadata.d_model:
+ raise ValueError("activation or model dimension does not match actor metadata")
+ if model.compatibility_mode or getattr(model, "_weights_processed", False):
+ raise ValueError("actor must use raw unprocessed Bridge weights")
+ if model.training or model.original_model.training:
+ raise ValueError("actor must be in eval mode")
+ sampling.validate_prompt(len(metadata.actor_token_ids), model.cfg.n_ctx)
+ config: Any = model.original_model.config
+ if getattr(config, "_attn_implementation", None) != "eager":
+ raise ValueError("actor requires eager attention")
+ if getattr(model.original_model, "is_quantized", False):
+ raise ValueError("quantized actors are not supported")
+ if any(
+ p.is_floating_point() and p.dtype != torch.float32
+ for p in model.original_model.parameters()
+ ):
+ raise ValueError("actor currently requires float32 weights")
+ if sampling.top_k > config.vocab_size:
+ raise ValueError("top_k exceeds actor vocabulary size")
+ if model.tokenizer is None:
+ raise ValueError("actor tokenizer is unavailable")
+ tokens = torch.tensor([metadata.actor_token_ids], dtype=torch.int64, device=model.cfg.device)
+ hook = embedding_overwrite(
+ activation.vector,
+ metadata.injection_scale,
+ metadata.injection_position,
+ tokens.shape[1],
+ cached=sampling.use_past_kv_cache,
+ )
+ with torch.inference_mode(), model.hooks(fwd_hooks=[("hook_embed", hook)]):
+ output = model.generate(
+ tokens,
+ return_type="tokens",
+ max_new_tokens=sampling.max_new_tokens,
+ temperature=float(sampling.temperature),
+ do_sample=sampling.do_sample,
+ top_p=float(sampling.top_p) if sampling.do_sample else None,
+ top_k=sampling.top_k or None,
+ use_past_kv_cache=sampling.use_past_kv_cache,
+ verbose=False,
+ )
+ if not isinstance(output, torch.Tensor) or output.ndim != 2 or output.shape[0] != 1:
+ raise ValueError("actor generation must return one token sequence")
+ if output.shape[1] < tokens.shape[1] or not torch.equal(output[:, : tokens.shape[1]], tokens):
+ raise ValueError("actor generation did not preserve the original prompt tokens")
+ ids = tuple(output[0, tokens.shape[1] :].detach().cpu().tolist())
+ raw = model.tokenizer.decode(list(ids), skip_special_tokens=True)
+ return parse_explanation(raw, ids, activation.provenance)
+
+
+def validate_activation_provenance(
+ activation: NLAActivation, pair: NLAPairConfig, artifact_name: str
+) -> None:
+ """Reject known source incompatibilities even for caller-attested inputs."""
+ provenance = activation.provenance
+ if provenance.artifact_name != artifact_name:
+ raise ValueError("activation provenance belongs to a different artifact")
+ if provenance.target is not None and provenance.target != pair.target:
+ raise ValueError("activation target identity/revision does not match the artifact")
+ if provenance.layer is not None and provenance.layer != pair.layer:
+ raise ValueError("activation layer does not match the artifact")
+ if provenance.basis != "raw_hf":
+ raise ValueError("activation must use the raw residual basis")
+
+
+class VerbalizerRuntime:
+ """Lazy fp32 actor loading with operation-local hooks and non-reentrant ownership."""
+
+ def __init__(
+ self,
+ pair: NLAPairConfig,
+ artifact_name: str,
+ *,
+ device: str = "cpu",
+ cache_dir: str | None = None,
+ local_files_only: bool = False,
+ ) -> None:
+ self.pair = pair
+ self.artifact_name = artifact_name
+ self.device = device
+ self.cache_dir = cache_dir
+ self.local_files_only = local_files_only
+ self._model: TransformerBridge | None = None
+ self._lock = threading.Lock()
+
+ def _metadata(self) -> tuple[Path, NLAMetadata, Any]:
+ directory = acquire_snapshot(
+ self.pair.av,
+ cache_dir=self.cache_dir,
+ local_files_only=self.local_files_only,
+ )
+ tokenizer = load_snapshot_tokenizer(directory)
+ metadata = load_local_metadata(directory, self.pair, "av", tokenizer)
+ return directory, metadata, tokenizer
+
+ def _load(self, directory: Path, tokenizer: Any) -> TransformerBridge:
+ snapshot = acquire_snapshot(
+ self.pair.av,
+ cache_dir=self.cache_dir,
+ local_files_only=self.local_files_only,
+ include_weights=True,
+ )
+ if snapshot != directory:
+ raise ValueError("actor weights must use the metadata snapshot")
+ original = call_hf_with_retry(
+ AutoModelForCausalLM.from_pretrained,
+ str(directory),
+ local_files_only=True,
+ trust_remote_code=False,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ )
+ model = TransformerBridge.boot_transformers(
+ self.pair.av.repo_id,
+ hf_model=original,
+ tokenizer=tokenizer,
+ device=self.device,
+ dtype=torch.float32,
+ trust_remote_code=False,
+ )
+ model.eval()
+ return model
+
+ def verbalize(
+ self, activation: NLAActivation | torch.Tensor, sampling: NLASamplingConfig | None = None
+ ) -> NLAExplanation:
+ """Validate small assets and budget before materializing full actor weights."""
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("concurrent verbalizer operations are not supported")
+ try:
+ sampling = sampling if sampling is not None else NLASamplingConfig()
+ if isinstance(activation, torch.Tensor):
+ activation = NLAActivation(
+ activation,
+ ActivationProvenance(self.artifact_name, SourceProvenance.UNKNOWN_VECTOR),
+ )
+ validate_activation_provenance(activation, self.pair, self.artifact_name)
+ directory, metadata, tokenizer = self._metadata()
+ detached_vector(activation.vector)
+ if activation.vector.numel() != metadata.d_model:
+ raise ValueError("activation dimension does not match the artifact")
+ config = json.loads((directory / "config.json").read_text(encoding="utf-8"))
+ sampling.validate_prompt(
+ len(metadata.actor_token_ids or ()), config["max_position_embeddings"]
+ )
+ if sampling.top_k > config["vocab_size"]:
+ raise ValueError("top_k exceeds actor vocabulary size")
+ if self._model is None:
+ self._model = self._load(directory, tokenizer)
+ return generate_explanation(self._model, metadata, activation, sampling)
+ finally:
+ self._lock.release()
+
+ def unload(self) -> None:
+ """Release only actor references owned by this runtime, idempotently."""
+ if not self._lock.acquire(blocking=False):
+ raise RuntimeError("cannot unload during a verbalizer operation")
+ try:
+ self._model = None
+ finally:
+ self._lock.release()
diff --git a/transformer_lens/tools/analysis/nla.py b/transformer_lens/tools/analysis/nla.py
new file mode 100644
index 000000000..373b4ad79
--- /dev/null
+++ b/transformer_lens/tools/analysis/nla.py
@@ -0,0 +1,30 @@
+"""Bridge-native Qwen natural-language autoencoder inference.
+
+Use raw fp32 models with eager attention. Reconstruction similarity measures
+vector preservation, not human-readable faithfulness. Models load lazily; the
+returned target is caller-owned and AV/AR resources are context-owned.
+"""
+
+from transformer_lens.tools.analysis._nla import NaturalLanguageAutoencoder
+from transformer_lens.tools.analysis._nla_contracts import (
+ ActivationProvenance,
+ ExplanationStatus,
+ NLAActivation,
+ NLAExplanation,
+ NLAReconstruction,
+ NLAReconstructionScore,
+ NLASamplingConfig,
+ SourceProvenance,
+)
+
+__all__ = [
+ "ActivationProvenance",
+ "ExplanationStatus",
+ "NLAActivation",
+ "NLAExplanation",
+ "NLAReconstruction",
+ "NLAReconstructionScore",
+ "NLASamplingConfig",
+ "NaturalLanguageAutoencoder",
+ "SourceProvenance",
+]
diff --git a/transformer_lens/tools/analysis/nla_registry.json b/transformer_lens/tools/analysis/nla_registry.json
new file mode 100644
index 000000000..da90ad8e0
--- /dev/null
+++ b/transformer_lens/tools/analysis/nla_registry.json
@@ -0,0 +1,23 @@
+{
+ "schema_version": 1,
+ "models": {
+ "qwen2.5-7b-instruct-L20": {
+ "target": {
+ "repo_id": "Qwen/Qwen2.5-7B-Instruct",
+ "revision": "a09a35458c702b33eeacc393d103063234e8bc28"
+ },
+ "av": {
+ "repo_id": "kitft/nla-qwen2.5-7b-L20-av",
+ "revision": "b88469162777ae6553bc14208eb0cb579336f8f4"
+ },
+ "ar": {
+ "repo_id": "kitft/nla-qwen2.5-7b-L20-ar",
+ "revision": "e2c9e57eac213d37a31612087f645ab6332c1bb6"
+ },
+ "model_type": "qwen2",
+ "d_model": 3584,
+ "layer": 20,
+ "target_layers": 28
+ }
+ }
+}
diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py
index 66b945995..8ec5e6ee0 100644
--- a/transformer_lens/utilities/tokenize_utils.py
+++ b/transformer_lens/utilities/tokenize_utils.py
@@ -152,7 +152,9 @@ def get_tokenizer_with_bos(tokenizer: PreTrainedTokenizerBase) -> PreTrainedToke
return tokenizer
init_kwargs = deepcopy(tokenizer.init_kwargs)
- pretrained_model_name_or_path = init_kwargs.pop("name_or_path")
+ pretrained_model_name_or_path = init_kwargs.pop("name_or_path", tokenizer.name_or_path)
+ if not pretrained_model_name_or_path:
+ return tokenizer
add_bos_token = init_kwargs.pop("add_bos_token", None)
if add_bos_token is None:
add_bos_token = getattr(tokenizer, "add_bos_token", False)
diff --git a/uv.lock b/uv.lock
index df22232fd..4e51ca0f8 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2,32 +2,32 @@ version = 1
revision = 3
requires-python = ">=3.10, <4.0"
resolution-markers = [
+ "python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
conflicts = [[
{ package = "transformer-lens", extra = "lit" },
@@ -1208,27 +1208,27 @@ name = "contourpy"
version = "1.3.3"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-16-transformer-lens-lit') or (extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -1825,10 +1825,10 @@ name = "etils"
version = "1.14.0"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*'",
- "python_full_version == '3.11.*'",
"python_full_version >= '3.14'",
"python_full_version == '3.13.*'",
+ "python_full_version == '3.12.*'",
+ "python_full_version == '3.11.*'",
]
sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" }
wheels = [
@@ -2571,14 +2571,14 @@ name = "grpcio"
version = "1.80.0"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version < '3.11' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version < '3.11' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version < '3.11' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "typing-extensions", marker = "(python_full_version < '3.14' and sys_platform == 'linux' and extra == 'extra-16-transformer-lens-vllm') or (python_full_version < '3.14' and extra == 'extra-16-transformer-lens-lit') or (extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -2707,10 +2707,10 @@ name = "grpcio-status"
version = "1.62.3"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
+ "python_full_version == '3.13.*'",
"python_full_version == '3.12.*'",
"python_full_version == '3.11.*'",
"python_full_version < '3.11'",
- "python_full_version == '3.13.*'",
]
dependencies = [
{ name = "googleapis-common-protos", marker = "python_full_version < '3.14'" },
@@ -3172,27 +3172,27 @@ name = "ipython"
version = "9.10.0"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "colorama", marker = "(python_full_version >= '3.11' and sys_platform == 'win32') or (python_full_version < '3.11' and extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm') or (sys_platform != 'win32' and extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -3286,27 +3286,27 @@ name = "jaxtyping"
version = "0.3.11"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "wadler-lindig", marker = "python_full_version >= '3.11' or (extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -5193,27 +5193,27 @@ name = "networkx"
version = "3.5"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
sdist = { url = "https://files.pythonhosted.org/packages/6c/4f/ccdb8ad3a38e583f214547fd2f7ff1fc160c43a75af88e6aec213404b96a/networkx-3.5.tar.gz", hash = "sha256:d4c6f9cf81f52d69230866796b82afbccdec3db7ae4fbd1b65ea750feed50037", size = 2471065, upload-time = "2025-05-29T11:35:07.804Z" }
wheels = [
@@ -5333,11 +5333,11 @@ name = "numpy"
version = "1.26.4"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
+ "python_full_version >= '3.14'",
+ "python_full_version == '3.13.*'",
"python_full_version == '3.12.*'",
"python_full_version == '3.11.*'",
"python_full_version < '3.11'",
- "python_full_version >= '3.14'",
- "python_full_version == '3.13.*'",
]
sdist = { url = "https://files.pythonhosted.org/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129, upload-time = "2024-02-06T00:26:44.495Z" }
wheels = [
@@ -5440,23 +5440,23 @@ name = "numpy"
version = "2.3.5"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
sdist = { url = "https://files.pythonhosted.org/packages/76/65/21b3bc86aac7b8f2862db1e808f1ea22b028e30a225a34a5ede9bf8678f2/numpy-2.3.5.tar.gz", hash = "sha256:784db1dcdab56bf0517743e746dfb0f885fc68d948aba86eeec2cba234bdf1c0", size = 20584950, upload-time = "2025-11-16T22:52:42.067Z" }
wheels = [
@@ -8143,10 +8143,10 @@ name = "scikit-image"
version = "0.26.0"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*'",
- "python_full_version == '3.11.*'",
"python_full_version >= '3.14'",
"python_full_version == '3.13.*'",
+ "python_full_version == '3.12.*'",
+ "python_full_version == '3.11.*'",
]
dependencies = [
{ name = "imageio", marker = "python_full_version >= '3.11'" },
@@ -8267,27 +8267,27 @@ name = "scikit-learn"
version = "1.8.0"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "joblib", marker = "python_full_version >= '3.11' or (extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -8405,27 +8405,27 @@ name = "scipy"
version = "1.17.1"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.13.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.14' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.13.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.12.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.12.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform == 'darwin' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm'",
+ "python_full_version == '3.11.*' and extra == 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
"python_full_version == '3.11.*' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
- "python_full_version >= '3.13' and extra != 'extra-16-transformer-lens-lit' and extra != 'extra-16-transformer-lens-vllm'",
]
dependencies = [
{ name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-16-transformer-lens-lit') or (extra == 'extra-16-transformer-lens-lit' and extra == 'extra-16-transformer-lens-vllm')" },
@@ -9099,10 +9099,10 @@ name = "tifffile"
version = "2026.3.3"
source = { registry = "https://pypi.org/simple" }
resolution-markers = [
- "python_full_version == '3.12.*'",
- "python_full_version == '3.11.*'",
"python_full_version >= '3.14'",
"python_full_version == '3.13.*'",
+ "python_full_version == '3.12.*'",
+ "python_full_version == '3.11.*'",
]
dependencies = [
{ name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
@@ -9517,7 +9517,9 @@ dependencies = [
{ name = "packaging" },
{ name = "pandas" },
{ name = "protobuf" },
+ { name = "pyyaml" },
{ name = "rich" },
+ { name = "safetensors" },
{ name = "sentencepiece" },
{ name = "torch" },
{ name = "tqdm" },
@@ -9562,6 +9564,7 @@ dev = [
{ name = "pytest-doctestplus" },
{ name = "pytest-rerunfailures" },
{ name = "pytest-xdist" },
+ { name = "types-pyyaml" },
]
docs = [
{ name = "furo" },
@@ -9615,7 +9618,9 @@ requires-dist = [
{ name = "packaging", specifier = ">=23.0" },
{ name = "pandas", specifier = ">=1.1.5" },
{ name = "protobuf", specifier = ">=3.20.0" },
+ { name = "pyyaml", specifier = ">=6.0" },
{ name = "rich", specifier = ">=12.6.0" },
+ { name = "safetensors", specifier = ">=0.4.3" },
{ name = "sentencepiece" },
{ name = "torch", specifier = ">=2.6" },
{ name = "tqdm", specifier = ">=4.64.1" },
@@ -9647,6 +9652,7 @@ dev = [
{ name = "pytest-doctestplus", specifier = ">=1.0.0" },
{ name = "pytest-rerunfailures", specifier = ">=12.0" },
{ name = "pytest-xdist", specifier = ">=3.8.0" },
+ { name = "types-pyyaml", specifier = ">=6.0" },
]
docs = [
{ name = "furo", specifier = ">=2023.3.27,<2024.0.0" },
@@ -9786,6 +9792,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/72/52/43e70a8e57fefb172c22a21000b03ebcc15e47e97f5cb8495b9c2832efb4/types_python_dateutil-2.9.0.20250708-py3-none-any.whl", hash = "sha256:4d6d0cc1cc4d24a2dc3816024e502564094497b713f7befda4d5bc7a8e3fd21f", size = 17724, upload-time = "2025-07-08T03:14:02.593Z" },
]
+[[package]]
+name = "types-pyyaml"
+version = "6.0.12.20260906"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/90/6e/abec85b9013db5b934b0280a6dd104904d84f7bcbaab2e2f3def87ac7463/types_pyyaml-6.0.12.20260906.tar.gz", hash = "sha256:f59c1cc05010b833d2d72287bbaa72610106b28d42d89a907313117faba85212", size = 18649, upload-time = "2026-09-06T06:35:35.362Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/15/c0/fc0644b7ddcfb969e95845837143cb5173ddd6e06ee4ba5fc493cd9329b7/types_pyyaml-6.0.12.20260906-py3-none-any.whl", hash = "sha256:bca893ff0d51df5c9053137d5d0e6ccd36e939a196356f1d5c16372422f5137b", size = 21282, upload-time = "2026-09-06T06:35:34.372Z" },
+]
+
[[package]]
name = "typing-extensions"
version = "4.14.1"