refactor: 从data 模块分离tokenizer

tests/conftest.py

@@ -11,7 +11,7 @@ from tokenizers import pre_tokenizers
 from torch.utils.data import Dataset
 
 from astrai.config.model_config import ModelConfig
-from astrai.data.tokenizer import BpeTokenizer, BpeTrainer
+from astrai.tokenizer import BpeTokenizer, BpeTrainer
 from astrai.model.transformer import Transformer
 
 

tests/data/test_checkpoint.py

@@ -5,7 +5,7 @@ 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.serialization import Checkpoint
 from astrai.parallel.setup import get_rank, spawn_parallel_fn
 
 

tests/data/test_dataset.py

@@ -2,7 +2,7 @@ import numpy as np
 import torch
 
 from astrai.data.dataset import DatasetFactory
-from astrai.data.serialization import save_h5
+from astrai.serialization import save_h5
 
 
 def test_dataset_loader_random_paths(base_test_env):

tests/trainer/test_early_stopping.py

@@ -4,7 +4,7 @@ import numpy as np
 import torch
 
 from astrai.config.train_config import TrainConfig
-from astrai.data.serialization import Checkpoint
+from astrai.serialization import Checkpoint
 from astrai.trainer.schedule import SchedulerFactory
 from astrai.trainer.trainer import Trainer