refactor: 优化参数传递，清理导入样式

tests/conftest.py

@@ -1,11 +1,12 @@
-import os
 import json
-import numpy as np
-import tempfile
+import os
 import shutil
-import torch
+import tempfile
+
+import numpy as np
 import pytest
 import safetensors.torch as st
+import torch
 from tokenizers import pre_tokenizers
 from torch.utils.data import Dataset
 

tests/data/test_checkpoint.py

@@ -1,9 +1,10 @@
-import torch
 import tempfile
-import torch.distributed as dist
 
+import torch
+import torch.distributed as dist
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from astrai.data.serialization import Checkpoint
 from astrai.parallel.setup import get_rank, spawn_parallel_fn
 

tests/data/test_dataset.py

@@ -1,8 +1,8 @@
-import torch
 import numpy as np
+import torch
 
-from astrai.data.serialization import save_h5
 from astrai.data.dataset import *
+from astrai.data.serialization import save_h5
 
 
 def test_dataset_loader_random_paths(base_test_env):

tests/data/test_sampler.py

@@ -1,5 +1,5 @@
-from astrai.trainer import *
 from astrai.data import *
+from astrai.trainer import *
 
 
 def test_random_sampler_consistency(random_dataset):

tests/inference/conftest.py

@@ -1,8 +1,10 @@
 """Shared fixtures for inference tests."""
 
-import pytest
 from unittest.mock import MagicMock, patch
+
+import pytest
 from fastapi.testclient import TestClient
+
 from astrai.inference.server import app
 
 

tests/inference/test_server.py

@@ -1,9 +1,6 @@
 """Unit tests for the inference HTTP server."""
 
 import pytest
-from unittest.mock import MagicMock, patch
-from fastapi.testclient import TestClient
-from astrai.inference.server import app
 
 
 def test_health_no_model(client, monkeypatch):

tests/module/test_module.py

@@ -1,10 +1,12 @@
 import os
+
 import torch
-from astrai.trainer import *
+
 from astrai.config import *
-from astrai.model import *
 from astrai.data import *
 from astrai.inference.generator import EmbeddingEncoderCore, GeneratorCore
+from astrai.model import *
+from astrai.trainer import *
 
 
 def test_model_parameter(test_env):

tests/module/test_tie_weight.py

@@ -1,11 +1,13 @@
-import os
 import json
-import torch
-import pytest
+import os
 import tempfile
+
+import pytest
 import safetensors.torch as st
-from astrai.model.transformer import Transformer
+import torch
+
 from astrai.config.model_config import ModelConfig
+from astrai.model.transformer import Transformer
 
 
 @pytest.fixture

tests/trainer/conftest.py

@@ -1,6 +1,9 @@
+import pytest
 import torch
 from torch.utils.data import Dataset
-import pytest
+
+from astrai.config import TrainConfig
+from astrai.trainer.schedule import SchedulerFactory
 
 
 class TrainerDataset(Dataset):
@@ -54,13 +57,11 @@ def create_train_config(
     Returns:
         TrainConfig instance configured for testing
     """
-    from astrai.config import TrainConfig
-    from astrai.config.schedule_config import CosineScheduleConfig
-    from astrai.trainer.schedule import SchedulerFactory
 
-    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
     optimizer_fn = lambda m: torch.optim.AdamW(m.parameters(), lr=0.001)
-    scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
+    scheduler_fn = lambda optim: SchedulerFactory.create(
+        optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+    )
 
     return TrainConfig(
         strategy=strategy,

tests/trainer/test_callbacks.py

@@ -6,10 +6,10 @@ from astrai.trainer import *
 
 def test_callback_integration(base_test_env, random_dataset):
     """Test that all callbacks are properly integrated"""
-    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-
     optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
-    scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
+    scheduler_fn = lambda optim: SchedulerFactory.create(
+        optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+    )
 
     train_config = TrainConfig(
         model=base_test_env["model"],

tests/trainer/test_early_stopping.py

@@ -1,18 +1,20 @@
 import os
-import torch
+
 import numpy as np
+import torch
+
 from astrai.config import *
-from astrai.trainer import *
 from astrai.data.serialization import Checkpoint
+from astrai.trainer import *
 
 
 def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
     """Simulate early stopping behavior"""
 
-    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-
     optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
-    scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
+    scheduler_fn = lambda optim: SchedulerFactory.create(
+        optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+    )
 
     train_config = TrainConfig(
         strategy="seq",

tests/trainer/test_train_strategy.py

@@ -1,10 +1,9 @@
-import torch
 import numpy as np
-import pytest
+import torch
 
 from astrai.config import *
-from astrai.trainer.schedule import *
 from astrai.data.dataset import *
+from astrai.trainer.schedule import *
 
 
 def test_schedule_factory_random_configs():
@@ -16,41 +15,57 @@ def test_schedule_factory_random_configs():
 
     # Test multiple random configurations
     for _ in range(5):  # Test 5 random configurations
-        schedule_configs = [
-            CosineScheduleConfig(
-                warmup_steps=np.random.randint(50, 200),
-                total_steps=np.random.randint(1000, 5000),
-                min_rate=np.random.uniform(0.01, 0.1),
-            ),
-            SGDRScheduleConfig(
-                warmup_steps=np.random.randint(50, 200),
-                cycle_length=np.random.randint(500, 2000),
-                t_mult=np.random.randint(1, 3),
-                min_rate=np.random.uniform(0.01, 0.1),
-            ),
-        ]
-
-        for config in schedule_configs:
-            # Validate configuration
-            config.validate()
-
-            # Create scheduler using factory
-            scheduler = SchedulerFactory.load(optimizer, config)
-
-            # Verify scheduler type
-            if isinstance(config, CosineScheduleConfig):
+        # Test multiple random configurations
+        cosine_params = {
+            "schedule_type": "cosine",
+            "warmup_steps": np.random.randint(50, 200),
+            "total_steps": np.random.randint(1000, 5000),
+            "min_rate": np.random.uniform(0.01, 0.1),
+        }
+        sgdr_params = {
+            "schedule_type": "sgdr",
+            "warmup_steps": np.random.randint(50, 200),
+            "cycle_length": np.random.randint(500, 2000),
+            "t_mult": np.random.randint(1, 3),
+            "min_rate": np.random.uniform(0.01, 0.1),
+        }
+        for params in [cosine_params, sgdr_params]:
+            schedule_type = params["schedule_type"]
+            # Convert parameters for scheduler constructor
+            if schedule_type == "cosine":
+                warmup_steps = params["warmup_steps"]
+                total_steps = params["total_steps"]
+                min_rate = params["min_rate"]
+                lr_decay_steps = total_steps - warmup_steps
+                scheduler = SchedulerFactory.create(
+                    optimizer,
+                    schedule_type,
+                    warmup_steps=warmup_steps,
+                    lr_decay_steps=lr_decay_steps,
+                    min_rate=min_rate,
+                )
                 assert isinstance(scheduler, CosineScheduler)
-                assert scheduler.warmup_steps == config.warmup_steps
-                assert (
-                    scheduler.lr_decay_steps == config.total_steps - config.warmup_steps
+                assert scheduler.warmup_steps == warmup_steps
+                assert scheduler.lr_decay_steps == lr_decay_steps
+                assert scheduler.min_rate == min_rate
+            elif schedule_type == "sgdr":
+                warmup_steps = params["warmup_steps"]
+                cycle_length = params["cycle_length"]
+                t_mult = params["t_mult"]
+                min_rate = params["min_rate"]
+                scheduler = SchedulerFactory.create(
+                    optimizer,
+                    schedule_type,
+                    warmup_steps=warmup_steps,
+                    cycle_length=cycle_length,
+                    t_mult=t_mult,
+                    min_rate=min_rate,
                 )
-                assert scheduler.min_rate == config.min_rate
-            elif isinstance(config, SGDRScheduleConfig):
                 assert isinstance(scheduler, SGDRScheduler)
-                assert scheduler.warmup_steps == config.warmup_steps
-                assert scheduler.cycle_length == config.cycle_length
-                assert scheduler.t_mult == config.t_mult
-                assert scheduler.min_rate == config.min_rate
+                assert scheduler.warmup_steps == warmup_steps
+                assert scheduler.cycle_length == cycle_length
+                assert scheduler.t_mult == t_mult
+                assert scheduler.min_rate == min_rate
 
             # Test scheduler state dict functionality
             state_dict = scheduler.state_dict()
@@ -76,16 +91,25 @@ def test_schedule_factory_edge_cases():
     # Test edge cases for CosineScheduleConfig
     edge_cases = [
         # Minimal warmup and steps
-        CosineScheduleConfig(warmup_steps=1, total_steps=10, min_rate=0.01),
+        {"warmup_steps": 1, "total_steps": 10, "min_rate": 0.01},
         # Large values
-        CosineScheduleConfig(warmup_steps=1000, total_steps=10000, min_rate=0.5),
+        {"warmup_steps": 1000, "total_steps": 10000, "min_rate": 0.5},
         # Zero min_rate (edge case)
-        CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=0.0),
+        {"warmup_steps": 100, "total_steps": 1000, "min_rate": 0.0},
     ]
 
-    for config in edge_cases:
-        config.validate()
-        scheduler = SchedulerFactory.load(optimizer, config)
+    for params in edge_cases:
+        warmup_steps = params["warmup_steps"]
+        total_steps = params["total_steps"]
+        min_rate = params["min_rate"]
+        lr_decay_steps = total_steps - warmup_steps
+        scheduler = SchedulerFactory.create(
+            optimizer,
+            "cosine",
+            warmup_steps=warmup_steps,
+            lr_decay_steps=lr_decay_steps,
+            min_rate=min_rate,
+        )
         assert scheduler is not None
 
         # Test multiple steps
@@ -93,34 +117,24 @@ def test_schedule_factory_edge_cases():
             scheduler.step()
 
 
-def test_schedule_factory_invalid_configs():
-    """Test scheduler factory with invalid configurations"""
-
-    # Test invalid configurations that should raise errors
-    invalid_configs = [
-        # Negative warmup steps
-        {"warmup_steps": -10, "total_steps": 1000, "min_rate": 0.1},
-        # Total steps less than warmup steps
-        {"warmup_steps": 500, "total_steps": 400, "min_rate": 0.1},
-        # Invalid min_rate
-        {"warmup_steps": 100, "total_steps": 1000, "min_rate": -0.1},
-        {"warmup_steps": 100, "total_steps": 1000, "min_rate": 1.1},
-    ]
-
-    for kwargs in invalid_configs:
-        with pytest.raises(ValueError):
-            config = CosineScheduleConfig(**kwargs)
-            config.validate()
-
-
 def test_schedule_factory_state_persistence():
     """Test scheduler state persistence (save/load)"""
 
     model = torch.nn.Linear(10, 2)
     optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
 
-    config = CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=0.1)
-    scheduler = SchedulerFactory.load(optimizer, config)
+    # Create scheduler directly with parameters
+    warmup_steps = 100
+    total_steps = 1000
+    min_rate = 0.1
+    lr_decay_steps = total_steps - warmup_steps
+    scheduler = SchedulerFactory.create(
+        optimizer,
+        "cosine",
+        warmup_steps=warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        min_rate=min_rate,
+    )
 
     # Take a few steps
     for _ in range(5):
@@ -129,8 +143,14 @@ def test_schedule_factory_state_persistence():
     # Save state
     state_dict = scheduler.state_dict()
 
-    # Create new scheduler and load state
-    new_scheduler = SchedulerFactory.load(optimizer, config)
+    # Create new scheduler with same parameters
+    new_scheduler = SchedulerFactory.create(
+        optimizer,
+        "cosine",
+        warmup_steps=warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        min_rate=min_rate,
+    )
     new_scheduler.load_state_dict(state_dict)
 
     # Verify states match

tests/trainer/test_trainer.py

@@ -1,5 +1,3 @@
-import torch
-
 from astrai.data.dataset import *
 from astrai.trainer import Trainer