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

astrai/model/__init__.py

@@ -1,9 +1,9 @@
 from astrai.model.module import (
+    GQA,
+    MLP,
+    DecoderBlock,
     Linear,
     RMSNorm,
-    MLP,
-    GQA,
-    DecoderBlock,
 )
 from astrai.model.transformer import Transformer
 

astrai/model/module.py

@@ -1,9 +1,9 @@
+from typing import Optional, Tuple
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 from torch import Tensor
-from typing import Optional, Tuple
 
 
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:

astrai/model/transformer.py

@@ -1,12 +1,13 @@
+from typing import Any, Mapping, Optional, Tuple
+
 import torch
 import torch.nn as nn
-
 from torch import Tensor
-from typing import Any, Mapping, Optional, Tuple
+
 from astrai.config.model_config import ModelConfig
 from astrai.model.module import (
-    Embedding,
     DecoderBlock,
+    Embedding,
     Linear,
     RMSNorm,
     RotaryEmbedding,