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

scripts/demo/download.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+
 from huggingface_hub import snapshot_download
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]

scripts/demo/generate_ar.py

@@ -1,7 +1,9 @@
-import torch
 from pathlib import Path
+
+import torch
+
 from astrai.config.param_config import ModelParameter
-from astrai.inference.generator import GeneratorFactory, GenerationRequest
+from astrai.inference.generator import GenerationRequest, GeneratorFactory
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")

scripts/demo/generate_batch.py

@@ -1,7 +1,9 @@
-import torch
 from pathlib import Path
+
+import torch
+
 from astrai.config.param_config import ModelParameter
-from astrai.inference.generator import GeneratorFactory, GenerationRequest
+from astrai.inference.generator import GenerationRequest, GeneratorFactory
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")

scripts/demo/stream_chat.py

@@ -1,7 +1,9 @@
-import torch
 from pathlib import Path
+
+import torch
+
 from astrai.config.param_config import ModelParameter
-from astrai.inference.generator import GeneratorFactory, GenerationRequest
+from astrai.inference.generator import GenerationRequest, GeneratorFactory
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")

scripts/tools/benchmark.py

@@ -1,6 +1,8 @@
-import torch
-from typing import Dict, Any
 from dataclasses import dataclass
+from typing import Any, Dict
+
+import torch
+
 from astrai.model.transformer import ModelConfig, Transformer
 
 

scripts/tools/generate.py

@@ -1,6 +1,7 @@
-import torch
-import json
 import argparse
+import json
+
+import torch
 
 from astrai.config.param_config import ModelParameter
 from astrai.inference.generator import BatchGenerator, GenerationRequest

scripts/tools/perplexity.py

@@ -1,11 +1,12 @@
+import argparse
 import json
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import argparse
 import tqdm
-
 from torch import Tensor
+
 from astrai.config.param_config import ModelParameter
 
 

scripts/tools/server.py

@@ -1,5 +1,6 @@
 import argparse
 from pathlib import Path
+
 from astrai.inference.server import run_server
 
 

scripts/tools/train.py

@@ -1,15 +1,16 @@
-import os
 import argparse
+import os
+from functools import partial
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from functools import partial
+from astrai.config import ModelParameter, TrainConfig
 from astrai.data import DatasetLoader
-from astrai.config import ModelParameter, TrainConfig, CosineScheduleConfig
-from astrai.trainer import Trainer, SchedulerFactory
 from astrai.parallel import get_rank
+from astrai.trainer import SchedulerFactory, Trainer
 
 
 def parse_args() -> argparse.Namespace:
@@ -158,7 +159,7 @@ def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
 def create_scheduler(
     optimizer: optim.Optimizer, **kwargs
 ) -> optim.lr_scheduler.LRScheduler:
-    return SchedulerFactory.load(optimizer, **kwargs)
+    return SchedulerFactory.create(optimizer, **kwargs)
 
 
 def prepare_checkpoint(model: nn.Module) -> dict:
@@ -211,11 +212,6 @@ def train(
         stride=stride,
     )
 
-    schedule_config = CosineScheduleConfig(
-        warmup_steps=warmup_steps,
-        total_steps=len(dataset) * n_epoch // (batch_size * nprocs),
-    )
-
     optimizer_fn = partial(
         create_optimizer,
         **{
@@ -224,7 +220,16 @@ def train(
             "weight_decay": adamw_weight_decay,
         },
     )
-    scheduler_fn = partial(create_scheduler, **{"schedule_config": schedule_config})
+
+    toltal_steps = len(dataset) * n_epoch // (batch_size * nprocs)
+    scheduler_fn = partial(
+        create_scheduler,
+        **{
+            "scheduler": "cosine",
+            "warmup_steps": warmup_steps,
+            "lr_decay_steps": toltal_steps - warmup_steps,
+        },
+    )
 
     train_config = TrainConfig(
         model=model,