mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""Tests for configuration loading."""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import yaml
|
|
import pytest
|
|
|
|
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)
|
|
|
|
def test_model_quantization_and_label_count_roundtrip(self):
|
|
raw = {
|
|
**SAMPLE_CONFIG,
|
|
"model": {
|
|
**SAMPLE_CONFIG["model"],
|
|
"num_labels": 7,
|
|
"quantization": "4bit",
|
|
},
|
|
}
|
|
config = StudyConfig.from_dict(raw)
|
|
assert config.model.quantization == "4bit"
|
|
assert config.model.num_labels == 7
|
|
assert StudyConfig.from_dict(config.to_dict()).model == config.model
|
|
|
|
def test_remote_config_roundtrip_preserves_security_and_execution_settings(self):
|
|
raw = {
|
|
**SAMPLE_CONFIG,
|
|
"remote": {
|
|
"host": "compute.example",
|
|
"user": "runner",
|
|
"port": 2222,
|
|
"ssh_key": "/keys/id",
|
|
"known_hosts_file": "/keys/known_hosts",
|
|
"remote_dir": "/srv/obliteratus",
|
|
"install_timeout": 120,
|
|
"python": "/opt/python",
|
|
"sync_results": False,
|
|
"gpus": "02, 0",
|
|
"install_source": "obliteratus==0.1.2",
|
|
},
|
|
}
|
|
config = StudyConfig.from_dict(raw)
|
|
assert config.remote is not None
|
|
assert config.remote.gpus == "2,0"
|
|
assert StudyConfig.from_dict(config.to_dict()) == config
|
|
|
|
@pytest.mark.parametrize(
|
|
"remote",
|
|
[
|
|
{"host": ""},
|
|
{"host": "host", "port": 0},
|
|
{"host": "host", "remote_dir": "relative"},
|
|
{"host": "host", "gpus": "0; injected"},
|
|
{"host": "host", "install_timeout": 0},
|
|
],
|
|
)
|
|
def test_remote_config_rejects_invalid_public_values(self, remote):
|
|
with pytest.raises(ValueError):
|
|
StudyConfig.from_dict({**SAMPLE_CONFIG, "remote": remote})
|