File size: 4,401 Bytes
9d24374
 
3a2b2e4
9d24374
 
 
 
3a2b2e4
 
9d24374
 
 
3a2b2e4
 
 
 
 
 
 
9838759
 
 
 
 
 
 
 
 
 
 
 
3a2b2e4
 
 
9838759
 
 
 
 
 
 
3a2b2e4
 
 
9838759
9d24374
 
 
 
 
3a2b2e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9838759
 
 
 
3a2b2e4
 
 
 
 
 
 
 
 
 
 
b784950
3a2b2e4
 
b784950
3a2b2e4
b784950
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3a2b2e4
9838759
3a2b2e4
 
9838759
 
 
 
3a2b2e4
9838759
3a2b2e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9d24374
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

from featurelens.config import SETTINGS

ROOT = Path(__file__).resolve().parents[1]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Run the full FeatureLens offline study.')
    parser.add_argument(
        '--resume',
        action='store_true',
        help='Skip stages whose expected outputs already exist.',
    )
    parser.add_argument(
        '--activation-batch-size',
        type=int,
        default=16,
        help='Batch size used only by experiments.collect_activations.',
    )
    parser.add_argument(
        '--activation-max-length',
        type=int,
        default=192,
        help='Maximum prompt length used only by experiments.collect_activations.',
    )
    return parser.parse_args()


def run(
    module: str,
    *,
    outputs: list[Path],
    resume: bool,
    extra_args: list[str] | None = None,
) -> None:
    if resume and outputs and all(path.exists() for path in outputs):
        print(f'\nSKIP {module}: expected outputs already exist.', flush=True)
        return
    command = [sys.executable, '-m', module, *(extra_args or [])]
    print('\n$', ' '.join(command), flush=True)
    subprocess.run(command, cwd=ROOT, check=True)


def main() -> None:
    args = parse_args()
    artifact_dir = ROOT / 'artifacts'
    activation_dir = artifact_dir / 'activations'

    run(
        'experiments.build_dataset',
        outputs=[ROOT / 'data' / 'prompts.jsonl', ROOT / 'data' / 'causal_tasks.jsonl'],
        resume=args.resume,
    )
    run(
        'experiments.collect_activations',
        outputs=[
            activation_dir / 'metadata.json',
            *[activation_dir / f'features_layer{layer}.npz' for layer in SETTINGS.layers],
            *[activation_dir / f'features_final_layer{layer}.npz' for layer in SETTINGS.layers],
        ],
        resume=args.resume,
        extra_args=[
            '--batch-size', str(args.activation_batch_size),
            '--max-length', str(args.activation_max_length),
        ],
    )
    run(
        'experiments.evaluate_features',
        outputs=[
            artifact_dir / 'feature_catalog.csv',
            artifact_dir / 'layer_metrics.csv',
            artifact_dir / 'stability.csv',
            artifact_dir / 'split.json',
        ],
        resume=args.resume,
    )
    final_output = artifact_dir / 'causal_results_final_token.csv'
    run(
        'experiments.run_causal',
        outputs=[final_output, final_output.with_suffix(final_output.suffix + '.complete')],
        resume=args.resume,
        extra_args=[
            '--position-policy', 'final_token',
            '--output', str(final_output),
            *(['--resume'] if args.resume else []),
        ],
    )
    max_active_output = artifact_dir / 'causal_results_max_active.csv'
    run(
        'experiments.run_causal',
        outputs=[
            max_active_output,
            max_active_output.with_suffix(max_active_output.suffix + '.complete'),
        ],
        resume=args.resume,
        extra_args=[
            '--position-policy', 'max_feature_activation',
            '--output', str(max_active_output),
            *(['--resume'] if args.resume else []),
        ],
    )
    feature_set_output = artifact_dir / 'feature_set_results.csv'
    run(
        'experiments.run_feature_sets',
        outputs=[
            feature_set_output,
            feature_set_output.with_suffix(feature_set_output.suffix + '.complete'),
        ],
        resume=args.resume,
        extra_args=['--resume'] if args.resume else None,
    )
    run(
        'experiments.analyze_stability',
        outputs=[artifact_dir / 'selection_stability.csv'],
        resume=args.resume,
    )
    run(
        'experiments.analyze_study',
        outputs=[artifact_dir / 'study_feature_summary.csv', artifact_dir / 'study_summary.json'],
        resume=args.resume,
    )
    run(
        'experiments.make_report',
        outputs=[artifact_dir / 'summary.json', artifact_dir / 'report.md'],
        resume=args.resume,
    )
    subprocess.run(
        [sys.executable, '-m', 'scripts.validate_artifacts'],
        cwd=ROOT,
        check=True,
    )
    print('\nFeatureLens experiment pipeline complete. See artifacts/report.md')


if __name__ == '__main__':
    main()