From 7e4896e492bc69815647763d104b34bdb1af14eb Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 09:17:50 +0530 Subject: [PATCH 01/10] Add validated NLA artifact configuration --- pyproject.toml | 3 + tests/unit/tools/test_nla.py | 380 ++++++++++++++++++ .../tools/analysis/_nla_artifacts.py | 375 +++++++++++++++++ .../tools/analysis/nla_registry.json | 4 + uv.lock | 233 ++++++----- 5 files changed, 886 insertions(+), 109 deletions(-) create mode 100644 tests/unit/tools/test_nla.py create mode 100644 transformer_lens/tools/analysis/_nla_artifacts.py create mode 100644 transformer_lens/tools/analysis/nla_registry.json diff --git a/pyproject.toml b/pyproject.toml index fc7142c950..20a1c923f2 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/unit/tools/test_nla.py b/tests/unit/tools/test_nla.py new file mode 100644 index 0000000000..6a2290269c --- /dev/null +++ b/tests/unit/tools/test_nla.py @@ -0,0 +1,380 @@ +"""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): + assert load_registry() == {} + 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/transformer_lens/tools/analysis/_nla_artifacts.py b/transformer_lens/tools/analysis/_nla_artifacts.py new file mode 100644 index 0000000000..595c61c0e7 --- /dev/null +++ b/transformer_lens/tools/analysis/_nla_artifacts.py @@ -0,0 +1,375 @@ +"""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.""" + + unk_token_id: int | None + + 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 == tokenizer.unk_token_id: + 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_registry.json b/transformer_lens/tools/analysis/nla_registry.json new file mode 100644 index 0000000000..c9890248e7 --- /dev/null +++ b/transformer_lens/tools/analysis/nla_registry.json @@ -0,0 +1,4 @@ +{ + "schema_version": 1, + "models": {} +} diff --git a/uv.lock b/uv.lock index df22232fd4..4e51ca0f8e 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" From e90be3819ca1d70f3697eea231993be0b4cd02cf Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 09:26:02 +0530 Subject: [PATCH 02/10] Add NLA result and reconstruction score contracts --- tests/unit/tools/test_nla_contracts.py | 283 ++++++++++++++++++ .../tools/analysis/_nla_contracts.py | 270 +++++++++++++++++ 2 files changed, 553 insertions(+) create mode 100644 tests/unit/tools/test_nla_contracts.py create mode 100644 transformer_lens/tools/analysis/_nla_contracts.py diff --git a/tests/unit/tools/test_nla_contracts.py b/tests/unit/tools/test_nla_contracts.py new file mode 100644 index 0000000000..2e4d1c9707 --- /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/transformer_lens/tools/analysis/_nla_contracts.py b/transformer_lens/tools/analysis/_nla_contracts.py new file mode 100644 index 0000000000..c320dc6017 --- /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") From 09e473eac4a8fc9142500a09e82d3a5bf9ceb179 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 09:32:08 +0530 Subject: [PATCH 03/10] Add pinned raw-target NLA activation extraction --- .../integration/test_nla_bridge_execution.py | 122 +++++++++++ tests/unit/tools/test_nla_target.py | 132 ++++++++++++ .../tools/analysis/_nla_target.py | 189 ++++++++++++++++++ 3 files changed, 443 insertions(+) create mode 100644 tests/integration/test_nla_bridge_execution.py create mode 100644 tests/unit/tools/test_nla_target.py create mode 100644 transformer_lens/tools/analysis/_nla_target.py diff --git a/tests/integration/test_nla_bridge_execution.py b/tests/integration/test_nla_bridge_execution.py new file mode 100644 index 0000000000..1df3d63810 --- /dev/null +++ b/tests/integration/test_nla_bridge_execution.py @@ -0,0 +1,122 @@ +"""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, NLAPairConfig +from transformer_lens.tools.analysis._nla_target import ( + extract_activation, + load_target, + validate_target_structure, +) + + +@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"] + + +def test_non_bridge_rejected(target_pair): + with pytest.raises(TypeError, match="TransformerBridge"): + validate_target_structure(object(), target_pair) diff --git a/tests/unit/tools/test_nla_target.py b/tests/unit/tools/test_nla_target.py new file mode 100644 index 0000000000..ccf46e7a24 --- /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/transformer_lens/tools/analysis/_nla_target.py b/transformer_lens/tools/analysis/_nla_target.py new file mode 100644 index 0000000000..ddc138959e --- /dev/null +++ b/transformer_lens/tools/analysis/_nla_target.py @@ -0,0 +1,189 @@ +"""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 +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") -> TransformerBridge: + """Load pinned fp32 raw weights with eager attention and caller ownership.""" + 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) From ef74a352a82aff09a983bdd65ef7eaa3a2f42a46 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 09:45:54 +0530 Subject: [PATCH 04/10] Add generation-aware NLA verbalizer inference --- .../integration/test_nla_bridge_execution.py | 112 ++++++++- tests/unit/tools/test_nla_verbalizer.py | 82 +++++++ .../tools/analysis/_nla_verbalizer.py | 228 ++++++++++++++++++ 3 files changed, 421 insertions(+), 1 deletion(-) create mode 100644 tests/unit/tools/test_nla_verbalizer.py create mode 100644 transformer_lens/tools/analysis/_nla_verbalizer.py diff --git a/tests/integration/test_nla_bridge_execution.py b/tests/integration/test_nla_bridge_execution.py index 1df3d63810..0663b2b42b 100644 --- a/tests/integration/test_nla_bridge_execution.py +++ b/tests/integration/test_nla_bridge_execution.py @@ -6,12 +6,23 @@ import pytest import torch -from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig +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") @@ -117,6 +128,105 @@ def observe(value, hook): 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_non_bridge_rejected(target_pair): with pytest.raises(TypeError, match="TransformerBridge"): validate_target_structure(object(), target_pair) diff --git a/tests/unit/tools/test_nla_verbalizer.py b/tests/unit/tools/test_nla_verbalizer.py new file mode 100644 index 0000000000..d839b35ecb --- /dev/null +++ b/tests/unit/tools/test_nla_verbalizer.py @@ -0,0 +1,82 @@ +"""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, +) + + +@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() + + +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/_nla_verbalizer.py b/transformer_lens/tools/analysis/_nla_verbalizer.py new file mode 100644 index 0000000000..b5b14acf7a --- /dev/null +++ b/transformer_lens/tools/analysis/_nla_verbalizer.py @@ -0,0 +1,228 @@ +"""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") + 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) + + +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), + ) + if activation.provenance.artifact_name != self.artifact_name: + raise ValueError("activation provenance belongs to a different artifact") + 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() From 5825e5469c3ed90bd8ee895e5f0fefd40d8c5a14 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 10:03:37 +0530 Subject: [PATCH 05/10] Add truncated NLA reconstruction with released value head --- .../integration/test_nla_bridge_execution.py | 100 ++++++++ tests/unit/tools/test_nla_reconstructor.py | 113 +++++++++ .../tools/analysis/_nla_reconstructor.py | 215 ++++++++++++++++++ 3 files changed, 428 insertions(+) create mode 100644 tests/unit/tools/test_nla_reconstructor.py create mode 100644 transformer_lens/tools/analysis/_nla_reconstructor.py diff --git a/tests/integration/test_nla_bridge_execution.py b/tests/integration/test_nla_bridge_execution.py index 0663b2b42b..b592b2010a 100644 --- a/tests/integration/test_nla_bridge_execution.py +++ b/tests/integration/test_nla_bridge_execution.py @@ -227,6 +227,106 @@ def test_actor_budget_rejected(target, actor_metadata): ) +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/unit/tools/test_nla_reconstructor.py b/tests/unit/tools/test_nla_reconstructor.py new file mode 100644 index 0000000000..bd8b2f2ab1 --- /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/transformer_lens/tools/analysis/_nla_reconstructor.py b/transformer_lens/tools/analysis/_nla_reconstructor.py new file mode 100644 index 0000000000..3f357f4b56 --- /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() From 30e654e12a25b6383c3164ce63c1ba5964af2915 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 10:16:12 +0530 Subject: [PATCH 06/10] Wire private NLA inference and owned-resource lifecycle --- .../integration/test_nla_facade_execution.py | 181 +++++++++++++++ tests/unit/tools/test_nla_facade.py | 85 +++++++ transformer_lens/tools/analysis/_nla.py | 219 ++++++++++++++++++ .../tools/analysis/_nla_artifacts.py | 4 +- .../tools/analysis/_nla_target.py | 44 +++- transformer_lens/utilities/tokenize_utils.py | 4 +- 6 files changed, 531 insertions(+), 6 deletions(-) create mode 100644 tests/integration/test_nla_facade_execution.py create mode 100644 tests/unit/tools/test_nla_facade.py create mode 100644 transformer_lens/tools/analysis/_nla.py diff --git a/tests/integration/test_nla_facade_execution.py b/tests/integration/test_nla_facade_execution.py new file mode 100644 index 0000000000..ec15073174 --- /dev/null +++ b/tests/integration/test_nla_facade_execution.py @@ -0,0 +1,181 @@ +"""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"] * 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_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/unit/tools/test_nla_facade.py b/tests/unit/tools/test_nla_facade.py new file mode 100644 index 0000000000..f2bf80b949 --- /dev/null +++ b/tests/unit/tools/test_nla_facade.py @@ -0,0 +1,85 @@ +"""Private facade lifecycle and import contracts without model loading.""" + +import pytest +import torch +from jaxtyping import TypeCheckError + +from transformer_lens.tools.analysis._nla 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_no_public_support_export(): + import transformer_lens.tools.analysis as analysis + + assert not hasattr(analysis, "NaturalLanguageAutoencoder") diff --git a/transformer_lens/tools/analysis/_nla.py b/transformer_lens/tools/analysis/_nla.py new file mode 100644 index 0000000000..980fc33839 --- /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 index 595c61c0e7..25ee29d882 100644 --- a/transformer_lens/tools/analysis/_nla_artifacts.py +++ b/transformer_lens/tools/analysis/_nla_artifacts.py @@ -25,8 +25,6 @@ class Tokenizer(Protocol): """Tokenization operations needed to check the artifact contract.""" - unk_token_id: int | None - def encode(self, text: str, *, add_special_tokens: bool) -> list[int]: ... @@ -227,7 +225,7 @@ def validate_metadata( 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 == tokenizer.unk_token_id: + 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") diff --git a/transformer_lens/tools/analysis/_nla_target.py b/transformer_lens/tools/analysis/_nla_target.py index ddc138959e..d9584d60dd 100644 --- a/transformer_lens/tools/analysis/_nla_target.py +++ b/transformer_lens/tools/analysis/_nla_target.py @@ -9,7 +9,12 @@ import torch from transformer_lens.model_bridge import TransformerBridge -from transformer_lens.tools.analysis._nla_artifacts import NLAPairConfig, _integer +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, @@ -92,8 +97,43 @@ def target_provenance( return SourceProvenance.CALLER_ATTESTED -def load_target(pair: NLAPairConfig, *, device: str = "cpu") -> TransformerBridge: +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, diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py index 66b9459957..8ec5e6ee07 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) From 7991ff5fbbcfd58ee30d8efe4f4da2d517560bd4 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 10:26:18 +0530 Subject: [PATCH 07/10] Add deterministic Qwen NLA artifact parity tests --- tests/integration/test_nla_oracle_parity.py | 337 ++++++++++++++++++++ 1 file changed, 337 insertions(+) create mode 100644 tests/integration/test_nla_oracle_parity.py diff --git a/tests/integration/test_nla_oracle_parity.py b/tests/integration/test_nla_oracle_parity.py new file mode 100644 index 0000000000..89673dea2d --- /dev/null +++ b/tests/integration/test_nla_oracle_parity.py @@ -0,0 +1,337 @@ +"""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_contracts import NLASamplingConfig +from transformer_lens.tools.analysis._nla_reconstructor import ( + ReconstructorRuntime, + 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().double() + source = source / source.abs().max() + source = source / torch.linalg.vector_norm(source) * scale + 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(): + require_storage((QWEN_PAIR.target, QWEN_PAIR.av, QWEN_PAIR.ar)) + activation = _normal_activation() + actor = VerbalizerRuntime(QWEN_PAIR, "qwen-parity", device="cpu") + try: + explanation = actor.verbalize(activation, NLASamplingConfig(max_new_tokens=512)) + finally: + actor.unload() + del actor + gc.collect() + assert ( + explanation.text is not None + ), f"AV parse failure: {explanation.parse_status}; {explanation.raw_text!r}" + critic = ReconstructorRuntime(QWEN_PAIR, device="cpu") + try: + result = critic.reconstruct(explanation.text) + finally: + critic.unload() + assert result.vector.shape == activation.shape + assert torch.isfinite(result.vector).all() + assert torch.any(result.vector != 0) From 9401eda9f2e25f5cf0cd7c132e80a6a67a3ffd43 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 10:39:55 +0530 Subject: [PATCH 08/10] Document experimental NLA inference and verify wheel resources --- .../content/natural_language_autoencoders.md | 210 ++++++++++++++++++ docs/source/index.md | 1 + tests/integration/test_nla_packaging.py | 89 ++++++++ 3 files changed, 300 insertions(+) create mode 100644 docs/source/content/natural_language_autoencoders.md create mode 100644 tests/integration/test_nla_packaging.py diff --git a/docs/source/content/natural_language_autoencoders.md b/docs/source/content/natural_language_autoencoders.md new file mode 100644 index 0000000000..0ff4ad665f --- /dev/null +++ b/docs/source/content/natural_language_autoencoders.md @@ -0,0 +1,210 @@ +# Natural language autoencoder inference + +## Availability and validation boundary + +Natural language autoencoder (NLA) execution is currently private experimental +infrastructure. There is no supported public `NaturalLanguageAutoencoder` export, +and the bundled support registry contains no verified pair. Calling its private +`from_pretrained` resolver fails for unregistered names rather than silently +loading an unverified artifact. Private interfaces may change. + +Local synthetic checkpoints exercise the inference mechanics. This does not +establish numerical compatibility with the released Qwen checkpoints. The slow +released-artifact checks must pass before any pair is advertised as supported. +The implementation currently uses raw fp32, eager-attention, single-device +execution; this is an execution restriction, not a universal precision/device +compatibility claim. + +## 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. + +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. + +## Experimental execution examples + +The following examples are for contributors with an explicitly constructed, +immutable `NLAPairConfig` whose artifacts they are validating. They deliberately +use the private constructor rather than claiming a working public registry name. +Replace `pair` with that explicit configuration; do not copy test fixture identities +into a real analysis or interpret these examples as verified model support. + +### 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._nla import NaturalLanguageAutoencoder + +with NaturalLanguageAutoencoder(pair, "experimental-pair", 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._nla import NaturalLanguageAutoencoder + +with NaturalLanguageAutoencoder(pair, "experimental-pair", 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. Released-checkpoint peak RAM/VRAM +and timing have not been measured; no accelerator fit 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 in a temporary directory, reads its resources, +and imports its artifact helper from an isolated extracted-wheel package with +network calls forbidden. This verifies 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 e81b4cf4cf..fae4cc375f 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/tests/integration/test_nla_packaging.py b/tests/integration/test_nla_packaging.py new file mode 100644 index 0000000000..41a560f639 --- /dev/null +++ b/tests/integration/test_nla_packaging.py @@ -0,0 +1,89 @@ +"""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", +) + + +def test_built_wheel_nla_resources(tmp_path): + repository = Path(__file__).resolve().parents[2] + output = tmp_path / "dist" + subprocess.run( + ["uv", "build", "--wheel", "--offline", "--out-dir", str(output)], + cwd=repository, + 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 importlib.util +import json +import pathlib +import socket +import sys +import types + +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 +# Empty package shells isolate resource imports from unrelated package bootstrapping. +for name, relative in ( + ("transformer_lens", "transformer_lens"), + ("transformer_lens.tools", "transformer_lens/tools"), + ("transformer_lens.tools.analysis", "transformer_lens/tools/analysis"), + ("transformer_lens.utilities", "transformer_lens/utilities"), +): + module = types.ModuleType(name) + module.__path__ = [str(root / relative)] + module.__spec__ = importlib.util.spec_from_file_location( + name, root / relative / "__init__.py", + submodule_search_locations=module.__path__, + ) + sys.modules[name] = module +from transformer_lens.tools.analysis._nla_artifacts import load_registry +from transformer_lens.tools.analysis._nla_contracts import NLASamplingConfig +assert isinstance(load_registry(), dict) +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 From da20fa3944ddb3a9792df5fafee3e905b23cb333 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sun, 11 Oct 2026 11:48:46 +0530 Subject: [PATCH 09/10] Fix NLA provenance validation and reference coverage --- .../content/natural_language_autoencoders.md | 14 +++-- .../integration/test_nla_facade_execution.py | 56 ++++++++++++++++++- tests/integration/test_nla_oracle_parity.py | 6 +- tests/integration/test_nla_packaging.py | 27 +++------ tests/unit/tools/test_nla_verbalizer.py | 56 +++++++++++++++++++ .../tools/analysis/_nla_verbalizer.py | 28 ++++++++-- 6 files changed, 157 insertions(+), 30 deletions(-) diff --git a/docs/source/content/natural_language_autoencoders.md b/docs/source/content/natural_language_autoencoders.md index 0ff4ad665f..ff9a5a8d59 100644 --- a/docs/source/content/natural_language_autoencoders.md +++ b/docs/source/content/natural_language_autoencoders.md @@ -36,7 +36,10 @@ 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. +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} `. @@ -196,9 +199,12 @@ uv build uv run pytest tests/integration/test_nla_packaging.py ``` -The packaging test builds a wheel in a temporary directory, reads its resources, -and imports its artifact helper from an isolated extracted-wheel package with -network calls forbidden. This verifies resource delivery, not model compatibility. +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 diff --git a/tests/integration/test_nla_facade_execution.py b/tests/integration/test_nla_facade_execution.py index ec15073174..4b11a856b3 100644 --- a/tests/integration/test_nla_facade_execution.py +++ b/tests/integration/test_nla_facade_execution.py @@ -25,7 +25,7 @@ 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"] * 10, + ["left MARK right Summary x END objects in space"] * 10, BpeTrainer( vocab_size=300, special_tokens=["UNK", "MARK"], initial_alphabet=ByteLevel.alphabet() ), @@ -162,6 +162,60 @@ def test_facade_local_round_trip_and_reload(local_artifacts): 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) diff --git a/tests/integration/test_nla_oracle_parity.py b/tests/integration/test_nla_oracle_parity.py index 89673dea2d..0afc518352 100644 --- a/tests/integration/test_nla_oracle_parity.py +++ b/tests/integration/test_nla_oracle_parity.py @@ -79,9 +79,9 @@ 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().double() - source = source / source.abs().max() - source = source / torch.linalg.vector_norm(source) * scale + 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) diff --git a/tests/integration/test_nla_packaging.py b/tests/integration/test_nla_packaging.py index 41a560f639..56ef236979 100644 --- a/tests/integration/test_nla_packaging.py +++ b/tests/integration/test_nla_packaging.py @@ -21,9 +21,12 @@ 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", "--offline", "--out-dir", str(output)], + ["uv", "build", "--wheel", "--out-dir", str(output)], cwd=repository, + env=build_environment, check=True, capture_output=True, text=True, @@ -39,12 +42,9 @@ def test_built_wheel_nla_resources(tmp_path): assert isinstance(registry["models"], dict) wheel.extractall(extracted) script = r""" -import importlib.util -import json import pathlib import socket import sys -import types root = pathlib.Path(sys.argv[1]) @@ -53,20 +53,11 @@ def forbid_network(*args, **kwargs): socket.socket.connect = forbid_network socket.create_connection = forbid_network -# Empty package shells isolate resource imports from unrelated package bootstrapping. -for name, relative in ( - ("transformer_lens", "transformer_lens"), - ("transformer_lens.tools", "transformer_lens/tools"), - ("transformer_lens.tools.analysis", "transformer_lens/tools/analysis"), - ("transformer_lens.utilities", "transformer_lens/utilities"), -): - module = types.ModuleType(name) - module.__path__ = [str(root / relative)] - module.__spec__ = importlib.util.spec_from_file_location( - name, root / relative / "__init__.py", - submodule_search_locations=module.__path__, - ) - sys.modules[name] = module +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 isinstance(load_registry(), dict) diff --git a/tests/unit/tools/test_nla_verbalizer.py b/tests/unit/tools/test_nla_verbalizer.py index d839b35ecb..e9ca97ab46 100644 --- a/tests/unit/tools/test_nla_verbalizer.py +++ b/tests/unit/tools/test_nla_verbalizer.py @@ -7,6 +7,7 @@ from transformer_lens.tools.analysis._nla_verbalizer import ( VerbalizerRuntime, embedding_overwrite, + validate_activation_provenance, ) @@ -76,6 +77,61 @@ def test_runtime_is_lazy_and_rejects_concurrent_operations(): 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"): diff --git a/transformer_lens/tools/analysis/_nla_verbalizer.py b/transformer_lens/tools/analysis/_nla_verbalizer.py index b5b14acf7a..d4042a41ea 100644 --- a/transformer_lens/tools/analysis/_nla_verbalizer.py +++ b/transformer_lens/tools/analysis/_nla_verbalizer.py @@ -45,8 +45,14 @@ def embedding_overwrite( raise ValueError("injection position must be within the original prompt") if type(cached) is not bool: raise ValueError("cached must be bool") - unit = source / source.abs().max() - injected = unit / unit.norm() * scale + 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: @@ -129,6 +135,21 @@ def generate_explanation( 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.""" @@ -200,8 +221,7 @@ def verbalize( activation, ActivationProvenance(self.artifact_name, SourceProvenance.UNKNOWN_VECTOR), ) - if activation.provenance.artifact_name != self.artifact_name: - raise ValueError("activation provenance belongs to a different artifact") + 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: From 599124cefde26086492346933a144217b3a9b9c9 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 12 Oct 2026 06:50:20 +0530 Subject: [PATCH 10/10] Expose verified Qwen NLA inference and public workflow --- .../content/natural_language_autoencoders.md | 55 +++++++++++-------- tests/integration/test_nla_oracle_parity.py | 46 +++++++++------- tests/integration/test_nla_packaging.py | 5 +- tests/unit/tools/test_nla.py | 4 +- tests/unit/tools/test_nla_facade.py | 11 +++- transformer_lens/tools/analysis/__init__.py | 22 ++++++++ transformer_lens/tools/analysis/nla.py | 30 ++++++++++ .../tools/analysis/nla_registry.json | 21 ++++++- 8 files changed, 144 insertions(+), 50 deletions(-) create mode 100644 transformer_lens/tools/analysis/nla.py diff --git a/docs/source/content/natural_language_autoencoders.md b/docs/source/content/natural_language_autoencoders.md index ff9a5a8d59..47bda6119c 100644 --- a/docs/source/content/natural_language_autoencoders.md +++ b/docs/source/content/natural_language_autoencoders.md @@ -2,18 +2,17 @@ ## Availability and validation boundary -Natural language autoencoder (NLA) execution is currently private experimental -infrastructure. There is no supported public `NaturalLanguageAutoencoder` export, -and the bundled support registry contains no verified pair. Calling its private -`from_pretrained` resolver fails for unregistered names rather than silently -loading an unverified artifact. Private interfaces may change. - -Local synthetic checkpoints exercise the inference mechanics. This does not -establish numerical compatibility with the released Qwen checkpoints. The slow -released-artifact checks must pass before any pair is advertised as supported. -The implementation currently uses raw fp32, eager-attention, single-device -execution; this is an execution restriction, not a universal precision/device -compatibility claim. +`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 @@ -79,13 +78,11 @@ 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. -## Experimental execution examples +## Execution examples -The following examples are for contributors with an explicitly constructed, -immutable `NLAPairConfig` whose artifacts they are validating. They deliberately -use the private constructor rather than claiming a working public registry name. -Replace `pair` with that explicit configuration; do not copy test fixture identities -into a real analysis or interpret these examples as verified model support. +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 @@ -93,9 +90,11 @@ 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._nla import NaturalLanguageAutoencoder +from transformer_lens.tools.analysis import NaturalLanguageAutoencoder -with NaturalLanguageAutoencoder(pair, "experimental-pair", device="cpu") as nla: +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.") @@ -115,9 +114,11 @@ 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._nla import NaturalLanguageAutoencoder +from transformer_lens.tools.analysis import NaturalLanguageAutoencoder -with NaturalLanguageAutoencoder(pair, "experimental-pair", device="cpu") as nla: +with NaturalLanguageAutoencoder.from_pretrained( + "qwen2.5-7b-instruct-L20", device="cpu" +) as nla: target = nla.load_target() try: activation = nla.extract( @@ -169,8 +170,14 @@ The workflow involves three separate 7B-class checkpoints. A pinned-artifact fil 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. Released-checkpoint peak RAM/VRAM -and timing have not been measured; no accelerator fit claim is made. +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 diff --git a/tests/integration/test_nla_oracle_parity.py b/tests/integration/test_nla_oracle_parity.py index 0afc518352..0cc4f96bf1 100644 --- a/tests/integration/test_nla_oracle_parity.py +++ b/tests/integration/test_nla_oracle_parity.py @@ -36,9 +36,7 @@ load_snapshot_tokenizer, validate_value_head_file, ) -from transformer_lens.tools.analysis._nla_contracts import NLASamplingConfig from transformer_lens.tools.analysis._nla_reconstructor import ( - ReconstructorRuntime, load_critic_model, prepare_critic_tokens, reconstruct_vector, @@ -315,23 +313,31 @@ def test_released_qwen_critic_vector_parity(): @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)) - activation = _normal_activation() - actor = VerbalizerRuntime(QWEN_PAIR, "qwen-parity", device="cpu") - try: - explanation = actor.verbalize(activation, NLASamplingConfig(max_new_tokens=512)) - finally: - actor.unload() - del actor + 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() - assert ( - explanation.text is not None - ), f"AV parse failure: {explanation.parse_status}; {explanation.raw_text!r}" - critic = ReconstructorRuntime(QWEN_PAIR, device="cpu") - try: - result = critic.reconstruct(explanation.text) - finally: - critic.unload() - assert result.vector.shape == activation.shape - assert torch.isfinite(result.vector).all() - assert torch.any(result.vector != 0) + 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 index 56ef236979..c01f6e93dd 100644 --- a/tests/integration/test_nla_packaging.py +++ b/tests/integration/test_nla_packaging.py @@ -15,6 +15,7 @@ "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", ) @@ -60,7 +61,9 @@ def forbid_network(*args, **kwargs): 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 isinstance(load_registry(), dict) +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 diff --git a/tests/unit/tools/test_nla.py b/tests/unit/tools/test_nla.py index 6a2290269c..b5b2d9e7a8 100644 --- a/tests/unit/tools/test_nla.py +++ b/tests/unit/tools/test_nla.py @@ -279,7 +279,9 @@ def registry_entry(): def test_registry_resource_and_parser(pair): - assert load_registry() == {} + 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 } diff --git a/tests/unit/tools/test_nla_facade.py b/tests/unit/tools/test_nla_facade.py index f2bf80b949..aeee9fb55f 100644 --- a/tests/unit/tools/test_nla_facade.py +++ b/tests/unit/tools/test_nla_facade.py @@ -4,7 +4,7 @@ import torch from jaxtyping import TypeCheckError -from transformer_lens.tools.analysis._nla import NaturalLanguageAutoencoder +from transformer_lens.tools.analysis import NaturalLanguageAutoencoder from transformer_lens.tools.analysis._nla_artifacts import ArtifactRef, NLAPairConfig @@ -79,7 +79,12 @@ def test_in_memory_tokenizer_bos_does_not_require_loading_path(): assert get_tokenizer_with_bos(tokenizer) is tokenizer -def test_no_public_support_export(): +def test_public_support_export(): import transformer_lens.tools.analysis as analysis + from transformer_lens.tools.analysis.nla import ( + NaturalLanguageAutoencoder, + NLAActivation, + ) - assert not hasattr(analysis, "NaturalLanguageAutoencoder") + assert analysis.NaturalLanguageAutoencoder is NaturalLanguageAutoencoder + assert analysis.NLAActivation is NLAActivation diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 13ac09858e..356db15b33 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 0000000000..373b4ad79c --- /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 index c9890248e7..da90ad8e03 100644 --- a/transformer_lens/tools/analysis/nla_registry.json +++ b/transformer_lens/tools/analysis/nla_registry.json @@ -1,4 +1,23 @@ { "schema_version": 1, - "models": {} + "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 + } + } }