Coverage for transformer_lens/tools/analysis/jacobian_lens_causal_swap_benchmark_cli.py: 51%

1"""Artifact-generation entry point for the causal coordinate-swap benchmark. 

2 

3Kept separate from 

4``transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark`` because the analysis 

5package's ``__init__`` imports that module: running it as a script would execute it twice and 

6``runpy`` would warn. This module is not imported by the package, so ``python -m`` runs it 

7once. 

8 

9Reproduce the frozen artifact with (no ``HF_TOKEN`` needed -- GPT-2 is not gated):: 

10 

11 uv run python -m transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark_cli 

12""" 

13 

14from __future__ import annotations 

15 

16import argparse 

17import json 

18from pathlib import Path 

19from typing import List, Optional, Sequence, Tuple 

20 

21from transformer_lens.tools.analysis.jacobian_lens import DEFAULT_K, JacobianLens 

22from transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark import ( 

23 BenchmarkCorpus, 

24 FunctionSpec, 

25 build_protocol_manifest, 

26 corpus_definition, 

27 run_causal_swap_benchmark, 

28 serialize_artifact, 

29 success_rate_ci, 

30) 

31 

32_COUNTRY_CORPUS = BenchmarkCorpus( 

33 name="countries", 

34 concepts=("France", "Canada", "China", "Egypt"), 

35 functions=( 

36 FunctionSpec( 

37 name="capital", 

38 template="The capital of {arg} is the city of", 

39 answers={"France": "Paris", "Canada": "Ottawa", "China": "Beijing", "Egypt": "Cairo"}, 

40 ), 

41 FunctionSpec( 

42 name="language", 

43 template="Most people in {arg} speak", 

44 answers={ 

45 "France": "French", 

46 "Canada": "English", 

47 "China": "Chinese", 

48 "Egypt": "Arabic", 

49 }, 

50 ), 

51 FunctionSpec( 

52 name="continent", 

53 template="{arg} is a country on the continent of", 

54 answers={"France": "Europe", "Canada": "North", "China": "Asia", "Egypt": "Africa"}, 

55 ), 

56 FunctionSpec( 

57 name="currency", 

58 template="The single-word name for the currency now used in {arg} is the", 

59 answers={"France": "Euro", "Canada": "Dollar", "China": "Yuan", "Egypt": "Pound"}, 

60 ), 

61 ), 

62) 

63 

64_GPT2_LENS_REPO = "neuronpedia/jacobian-lens" 

65_GPT2_LENS_FILE = "gpt2-small/jlens/Salesforce-wikitext/gpt2_jacobian_lens.pt" 

66_GPT2_LENS_REVISION = "a4114d7752d11eb546e6cf372213d7e75526d3a1" 

67# The corpus mirrors the country config ``Jacobian_Lens_Demo.ipynb`` pins, so it reuses that 

68# notebook's repo/path/revision triple rather than inventing a second provenance scheme. 

69_CORPUS_REPO = "anthropics/jacobian-lens" 

70_CORPUS_PATH = "data/experiments/flexible-generalization.json" 

71_CORPUS_REVISION = "581d398613e5602a5af361e1c34d3a92ea82ba8e" 

72_DEFAULT_CONTROL_SEEDS = (0, 1, 2, 3, 4) 

73_DEFAULT_ARTIFACT_PATH = ( 

74 Path(__file__).resolve().parents[3] 

75 / "demos" 

76 / "data" 

77 / "jacobian_lens_causal_swap_benchmark_gpt2.json" 

78) 

79 

80 

81def _parse_seed_list(value: str) -> Tuple[int, ...]: 

82 """Parses a comma-separated seed list for ``--control-seeds``.""" 

83 seeds = tuple(int(part) for part in value.split(",") if part.strip()) 

84 if not seeds: 

85 raise ValueError("at least one control seed is required") 

86 return seeds 

87 

88 

89def _build_parser() -> argparse.ArgumentParser: 

90 """Builds the CLI argument parser.""" 

91 parser = argparse.ArgumentParser( 

92 description="Generates the frozen GPT-2 causal-swap benchmark artifact." 

93 ) 

94 parser.add_argument("--output", type=Path, default=_DEFAULT_ARTIFACT_PATH) 

95 parser.add_argument("--control-tolerance", type=float, default=0.1) 

96 parser.add_argument( 

97 "--control-seeds", 

98 type=_parse_seed_list, 

99 default=_DEFAULT_CONTROL_SEEDS, 

100 help="comma-separated seeds to draw the control arm under (default: %(default)s)", 

101 ) 

102 parser.add_argument("--alpha", type=float, default=1.0) 

103 parser.add_argument("--k", type=int, default=DEFAULT_K) 

104 return parser 

105 

106 

107def main(argv: Optional[Sequence[str]] = None) -> None: 

108 """Generates the frozen GPT-2 causal-swap benchmark artifact. 

109 

110 Loads the published GPT-2-small lens (the same artifact 

111 ``tests/integration/test_jacobian_lens.py`` uses), runs the reused country corpus over 

112 every one of the lens's fitted source layers, and writes the versioned, fingerprinted JSON 

113 artifact consumed by ``demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb``. 

114 """ 

115 import torch 

116 

117 from transformer_lens.model_bridge import TransformerBridge 

118 

119 args = _build_parser().parse_args(argv) 

120 

121 device = "cuda" if torch.cuda.is_available() else "cpu" 

122 model = TransformerBridge.boot_transformers("gpt2", dtype=torch.float32, device=device) 

123 lens = JacobianLens.from_pretrained( 

124 _GPT2_LENS_REPO, 

125 filename=_GPT2_LENS_FILE, 

126 revision=_GPT2_LENS_REVISION, 

127 model=model, 

128 ) 

129 layers = list(lens.source_layers) 

130 

131 trials, excluded = run_causal_swap_benchmark( 

132 lens, 

133 model, 

134 _COUNTRY_CORPUS, 

135 layers, 

136 control_tolerance=args.control_tolerance, 

137 control_seeds=args.control_seeds, 

138 alpha=args.alpha, 

139 k=args.k, 

140 ) 

141 ok_trials = [trial for trial in trials if trial.status == "ok"] 

142 if not ok_trials: 

143 raise RuntimeError( 

144 "no trial survived baseline filtering, the active-support check, and " 

145 "control-token selection" 

146 ) 

147 

148 real_successes: List[bool] = [] 

149 control_successes: List[bool] = [] 

150 for trial in ok_trials: 

151 assert trial.real_target_metrics is not None 

152 assert trial.control_target_metrics is not None 

153 real_successes.append(trial.real_target_metrics.target_is_top1) 

154 control_successes.append(trial.control_target_metrics.target_is_top1) 

155 

156 manifest = build_protocol_manifest( 

157 model_id="gpt2", 

158 model_revision="n/a", 

159 lens_repo=_GPT2_LENS_REPO, 

160 lens_file=_GPT2_LENS_FILE, 

161 lens_revision=_GPT2_LENS_REVISION, 

162 corpus_name=_COUNTRY_CORPUS.name, 

163 corpus=corpus_definition(_COUNTRY_CORPUS), 

164 corpus_repo=_CORPUS_REPO, 

165 corpus_path=_CORPUS_PATH, 

166 corpus_revision=_CORPUS_REVISION, 

167 layers_swept=layers, 

168 alpha=args.alpha, 

169 k=args.k, 

170 control_tolerance=args.control_tolerance, 

171 control_seeds=list(args.control_seeds), 

172 success_definition="target token id equals deterministic argmax token id", 

173 baseline_definition="source answer token id equals deterministic argmax token id", 

174 rank_definition="1 + count(logits strictly greater than target logit)", 

175 ) 

176 real_ci = success_rate_ci(real_successes) 

177 control_ci = success_rate_ci(control_successes) 

178 artifact = serialize_artifact(manifest, trials, excluded, real_ci, control_ci) 

179 

180 args.output.parent.mkdir(parents=True, exist_ok=True) 

181 args.output.write_text(json.dumps(artifact, indent=2, sort_keys=True) + "\n") 

182 status_counts = artifact["trial_status_counts"] 

183 n_prompts = artifact["n_independent_prompts"] 

184 print( 

185 f"wrote {len(trials)} trials ({len(ok_trials)} ok over {n_prompts} independent " 

186 f"prompt(s) at layers {artifact['layers_executed']}, " 

187 f"{status_counts.get('skipped_source_inactive', 0)} source-inactive, " 

188 f"{status_counts.get('skipped_no_control_token', 0)} no-control-token, " 

189 f"{len(excluded)} excluded prompts) to {args.output}" 

190 ) 

191 

192 

193if __name__ == "__main__": 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 main()