diff --git a/tests/model/test_config.py b/tests/model/test_config.py index 93f6f4a..5789660 100644 --- a/tests/model/test_config.py +++ b/tests/model/test_config.py @@ -1,9 +1,9 @@ import os import tempfile import pytest -import torch.nn as nn from aiia import AIIAConfig + def test_aiia_config_initialization(): config = AIIAConfig() assert config.model_type == "AIIA" @@ -12,7 +12,7 @@ def test_aiia_config_initialization(): assert config.hidden_size == 512 assert config.num_hidden_layers == 12 assert config.num_channels == 3 - assert config.learning_rate == 5e-5 + def test_aiia_config_custom_initialization(): config = AIIAConfig( @@ -21,8 +21,7 @@ def test_aiia_config_custom_initialization(): activation_function="ReLU", hidden_size=1024, num_hidden_layers=8, - num_channels=1, - learning_rate=1e-4 + num_channels=1 ) assert config.model_type == "CustomModel" assert config.kernel_size == 5 @@ -30,19 +29,20 @@ def test_aiia_config_custom_initialization(): assert config.hidden_size == 1024 assert config.num_hidden_layers == 8 assert config.num_channels == 1 - assert config.learning_rate == 1e-4 + def test_aiia_config_invalid_activation_function(): with pytest.raises(ValueError): AIIAConfig(activation_function="InvalidFunction") + def test_aiia_config_to_dict(): config = AIIAConfig() config_dict = config.to_dict() assert isinstance(config_dict, dict) - assert config_dict["model_type"] == "AIIA" assert config_dict["kernel_size"] == 3 + def test_aiia_config_save_pretrained_and_from_pretrained(): with tempfile.TemporaryDirectory() as tmpdir: config = AIIAConfig(model_type="TempModel") @@ -54,6 +54,7 @@ def test_aiia_config_save_pretrained_and_from_pretrained(): assert loaded_config.kernel_size == 3 assert loaded_config.activation_function == "GELU" + def test_aiia_config_save_pretrained_and_load_with_custom_attributes(): with tempfile.TemporaryDirectory() as tmpdir: config = AIIAConfig(model_type="TempModel", custom_attr="value") @@ -64,6 +65,7 @@ def test_aiia_config_save_pretrained_and_load_with_custom_attributes(): assert loaded_config.model_type == "TempModel" assert loaded_config.custom_attr == "value" + def test_aiia_config_save_pretrained_and_load_with_nested_attributes(): with tempfile.TemporaryDirectory() as tmpdir: config = AIIAConfig(model_type="TempModel", nested={"key": "value"}) @@ -72,4 +74,4 @@ def test_aiia_config_save_pretrained_and_load_with_nested_attributes(): loaded_config = AIIAConfig.from_pretrained(save_pretrained_path) assert loaded_config.model_type == "TempModel" - assert loaded_config.nested == {"key": "value"} + assert loaded_config.nested == {"key": "value"} \ No newline at end of file