File size: 2,658 Bytes
e0265b9 | 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 | from __future__ import annotations
from pathlib import Path
import json
import pytest
from adam.assets import AssetRegistry
from adam.commands import CommandValidationError, TrainingCommand
class FakeConfig:
def __init__(self, values: dict) -> None:
self.values = values
def get(self, key: str, default=None):
return self.values.get(key, default)
def test_asset_registry_prefers_exact_friendly_name(tmp_path: Path) -> None:
exact = tmp_path / "Mario"
similar = tmp_path / "Mario 2"
exact.mkdir()
similar.mkdir()
registry = AssetRegistry(tmp_path)
registry.register(kind="dataset", name="Mario", path=str(exact))
registry.register(kind="dataset", name="Mario 2", path=str(similar))
assert [item.name for item in registry.find("dataset", "Mario")] == ["Mario"]
def test_training_command_rejects_uncontrolled_fields() -> None:
with pytest.raises(CommandValidationError, match="Unsupported command fields"):
TrainingCommand.from_dict(
{
"action": "train",
"trainer": "ddpm",
"dataset": "dataset",
"model_name": "model",
"epochs": 10,
"shell_command": "unsafe",
}
)
def test_training_command_accepts_only_safe_ddpm_options() -> None:
command = TrainingCommand.from_dict({
"action": "train", "trainer": "ddpm", "dataset": "dataset", "model_name": "model", "epochs": 10,
"training_options": {"resolution": 256, "batch_size": 2, "learning_rate": 0.0001},
})
assert command.training_options["resolution"] == 256
def test_flow_discovery_recovers_dataset_from_adam_job_history(tmp_path: Path) -> None:
flow_root = tmp_path / "Flow"
dataset = tmp_path / "Dataset"
model = flow_root / "output_flow_models" / "Model"
dataset.mkdir()
(model / "unet").mkdir(parents=True)
(model / "unet" / "config.json").write_text("{}", encoding="utf-8")
(model / "flow_model_info.json").write_text(
'{"model_type":"rectified_flow","name":"Friendly Flow","resolution":128}',
encoding="utf-8",
)
(tmp_path / "data").mkdir()
(tmp_path / "data" / "jobs.json").write_text(json.dumps({"jobs": [{"plan": {"steps": [{
"tool_id": "flow_trainer", "arguments": {
"output_dir": str(model), "dataset_dir": str(dataset),
},
}]}}]}), encoding="utf-8")
registry = AssetRegistry(tmp_path)
registry.discover(FakeConfig({"tool_folders": {"flow_trainer": str(flow_root)}}))
flow = registry.find("model", "Friendly Flow", trainer="flow")[0]
assert flow.dataset_id
|