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

astrai/data/__init__.py

@@ -1,16 +1,15 @@
 from astrai.data.dataset import (
     BaseDataset,
-    SEQDataset,
+    DatasetFactory,
+    DatasetLoader,
     DPODataset,
-    SFTDataset,
     GRPODataset,
     MultiSegmentFetcher,
-    DatasetLoader,
-    DatasetFactory,
+    SEQDataset,
+    SFTDataset,
 )
-
-from astrai.data.tokenizer import BpeTokenizer
 from astrai.data.sampler import ResumableDistributedSampler
+from astrai.data.tokenizer import BpeTokenizer
 
 __all__ = [
     # Base classes

astrai/data/dataset.py

@@ -1,13 +1,14 @@
 """Dataset implementations with factory pattern for training."""
 
-import torch
 import bisect
-
 from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+import torch
 from torch import Tensor
 from torch.utils.data import Dataset
+
 from astrai.data.serialization import load_h5
-from typing import List, Dict, Optional, Union
 
 
 class BaseSegmentFetcher:

astrai/data/sampler.py

@@ -1,8 +1,8 @@
+from typing import Optional
+
 import torch
 import torch.distributed as dist
-
 from torch.utils.data import Dataset, Sampler
-from typing import Optional
 
 
 class ResumableDistributedSampler(Sampler[int]):

astrai/data/serialization.py

@@ -1,13 +1,14 @@
-import os
-import h5py
-import torch
 import json
-import safetensors.torch as st
-import torch.distributed as dist
-
+import os
 from pathlib import Path
-from torch import Tensor
 from typing import Any, Dict, List
+
+import h5py
+import safetensors.torch as st
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
 from astrai.parallel.setup import get_rank
 
 

astrai/data/tokenizer.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from tokenizers import Tokenizer, decoders, processors, normalizers, pre_tokenizers
+from typing import List, Union
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
 from tokenizers.models import BPE
 from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
-from typing import List, Union
 
 
 class BaseTokenizer(ABC):