Coverage for transformer_lens/model_bridge/sources/inspect/wire.py: 89%

1"""Serialization chokepoint for the Inspect activation wire format. 

2 

3Activations ride in ``ModelOutput.metadata["activations"]`` as a flat 

4``{"<layer>:<kind>": {"data": <b64>, "dtype": str, "shape": [...]}}`` map (keys 

5are :func:`hooks.wire_key`). For vllm-lens interop, decode also understands its 

6*documented* nested ``{"residual_stream": {layer: ...}}`` shape (mapped to 

7``resid_post``) — unverified against a live vllm-lens provider. 

8Numpy-only (no torch) so both the torch-using provider and the torch-free driver 

9import it; the single place to patch on format drift. 

10""" 

11from __future__ import annotations 

12 

13import base64 

14from typing import Any, Iterable, Mapping 

15 

16import numpy as np 

17 

18_RESIDUAL = "residual_stream" # vllm-lens's nested key (interop decode only) 

19 

20 

21def encode_array(arr: np.ndarray) -> dict[str, Any]: 

22 """numpy array → ``{"data": b64, "dtype": str, "shape": [...]}``.""" 

23 contiguous = np.ascontiguousarray(arr) 

24 return { 

25 "data": base64.b64encode(contiguous.tobytes()).decode("ascii"), 

26 "dtype": str(contiguous.dtype), 

27 "shape": list(contiguous.shape), 

28 } 

29 

30 

31def decode_array(entry: Any) -> np.ndarray: 

32 """Inverse of :func:`encode_array`; passes an already-decoded ndarray through.""" 

33 if isinstance(entry, np.ndarray): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 return entry 

35 raw = base64.b64decode(entry["data"]) 

36 # frombuffer is read-only; copy so the downstream torch tensor is writable. 

37 return np.frombuffer(raw, dtype=np.dtype(entry["dtype"])).reshape(entry["shape"]).copy() 

38 

39 

40def encode_activations(captured: Mapping[str, np.ndarray]) -> dict[str, Any]: 

41 """``{wire_key: array}`` → the ``metadata["activations"]`` payload.""" 

42 return {key: encode_array(arr) for key, arr in captured.items()} 

43 

44 

45def decode_activations( 

46 metadata: Mapping[str, Any] | None, wire_keys: Iterable[str] 

47) -> dict[str, np.ndarray]: 

48 """Pull the requested ``<layer>:<kind>`` keys out of ``metadata["activations"]``, 

49 falling back to the nested ``residual_stream`` for ``resid_post``. Missing keys are 

50 skipped — the caller decides.""" 

51 activations = (metadata or {}).get("activations") or {} 

52 residual = activations.get(_RESIDUAL, {}) 

53 out: dict[str, np.ndarray] = {} 

54 for key in wire_keys: 

55 if key in activations: 

56 out[key] = decode_array(activations[key]) 

57 continue 

58 layer, _, kind = key.partition(":") 

59 if kind == "resid_post": 59 ↛ 54line 59 didn't jump to line 54 because the condition on line 59 was always true

60 entry = residual.get(layer, residual.get(_safe_int(layer))) 

61 if entry is not None: 

62 out[key] = decode_array(entry) 

63 return out 

64 

65 

66def _safe_int(value: str) -> Any: 

67 try: 

68 return int(value) 

69 except ValueError: 

70 return value