File size: 1,656 Bytes
2bc8e46 | 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 | """Tests for configuration loading."""
from __future__ import annotations
import yaml
from obliteratus.config import StudyConfig
SAMPLE_CONFIG = {
"model": {
"name": "gpt2",
"task": "causal_lm",
"dtype": "float32",
"device": "cpu",
},
"dataset": {
"name": "wikitext",
"subset": "wikitext-2-raw-v1",
"split": "test",
"text_column": "text",
"max_samples": 50,
},
"strategies": [
{"name": "layer_removal", "params": {}},
{"name": "ffn_ablation", "params": {}},
],
"metrics": ["perplexity"],
"batch_size": 4,
"max_length": 256,
"output_dir": "results/test",
}
class TestStudyConfig:
def test_from_dict(self):
config = StudyConfig.from_dict(SAMPLE_CONFIG)
assert config.model.name == "gpt2"
assert config.model.task == "causal_lm"
assert config.dataset.name == "wikitext"
assert len(config.strategies) == 2
assert config.strategies[0].name == "layer_removal"
def test_from_yaml(self, tmp_path):
yaml_path = tmp_path / "test_config.yaml"
yaml_path.write_text(yaml.dump(SAMPLE_CONFIG))
config = StudyConfig.from_yaml(yaml_path)
assert config.model.name == "gpt2"
assert config.batch_size == 4
def test_roundtrip(self):
config = StudyConfig.from_dict(SAMPLE_CONFIG)
d = config.to_dict()
config2 = StudyConfig.from_dict(d)
assert config2.model.name == config.model.name
assert config2.dataset.name == config.dataset.name
assert len(config2.strategies) == len(config.strategies)
|