Spaces:
Running on Zero
Running on Zero
File size: 2,432 Bytes
9838759 | 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 | from __future__ import annotations
import csv
import sys
import types
from pathlib import Path
def _stub_transformers() -> None:
if 'transformers' in sys.modules:
return
stub = types.ModuleType('transformers')
stub.AutoModelForCausalLM = type('AutoModelForCausalLM', (), {})
stub.AutoTokenizer = type('AutoTokenizer', (), {})
sys.modules['transformers'] = stub
def test_run_all_forwards_activation_collection_tuning(monkeypatch, tmp_path: Path) -> None:
from experiments import run_all
calls: list[list[str]] = []
def fake_run(command, **_kwargs):
calls.append(list(command))
return types.SimpleNamespace(returncode=0)
monkeypatch.setattr(run_all.subprocess, 'run', fake_run)
missing = tmp_path / 'not-created'
run_all.run(
'experiments.collect_activations',
outputs=[missing],
resume=True,
extra_args=['--batch-size', '8', '--max-length', '192'],
)
assert calls
assert calls[0][-4:] == ['--batch-size', '8', '--max-length', '192']
def test_causal_checkpoint_helpers_roundtrip(tmp_path: Path) -> None:
_stub_transformers()
from experiments import run_causal
output = tmp_path / 'causal_results.csv'
rows = [
{'task_id': 'task-1', 'concept': 'mathematics', 'condition': 'sae_feature'},
{'task_id': 'task-1', 'concept': 'mathematics', 'condition': 'random_norm_matched'},
]
run_causal._write_rows_atomic(output, rows)
loaded = run_causal._load_checkpoint_rows(output)
assert loaded == rows
assert run_causal._completion_marker(output).name == 'causal_results.csv.complete'
assert not output.with_suffix('.csv.tmp').exists()
def test_feature_set_checkpoint_helpers_roundtrip(tmp_path: Path) -> None:
_stub_transformers()
from experiments import run_feature_sets
output = tmp_path / 'feature_set_results.csv'
rows = [
{'task_id': 'task-1', 'set_size': '1', 'condition': 'sae_feature_set'},
{'task_id': 'task-1', 'set_size': '1', 'condition': 'random_norm_matched'},
]
run_feature_sets._write_rows_atomic(output, rows)
with output.open(newline='', encoding='utf-8') as handle:
saved = list(csv.DictReader(handle))
assert saved == rows
assert run_feature_sets._load_checkpoint_rows(output) == rows
assert run_feature_sets._completion_marker(output).name == 'feature_set_results.csv.complete'
|