refactor(khaosz): 重构项目结构

tests/conftest.py

@@ -9,9 +9,11 @@ import pytest
 import matplotlib
 from torch.utils.data import Dataset
 
-from khaosz.core import *
-from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+from khaosz.config.model_config import TransformerConfig
+from khaosz.data.data_util import build_attention_mask, build_loss_mask
+from khaosz.data.tokenizer import BpeTokenizer
+from khaosz.model.transformer import Transformer
+
 
 matplotlib.use("Agg")
 

tests/test_callbacks.py

@@ -1,8 +1,7 @@
 import torch
 
-from khaosz.core import *
+from khaosz.config import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
 
 def test_callback_integration(base_test_env, random_dataset):
     """Test that all callbacks are properly integrated"""

tests/test_dataset_loader.py

@@ -3,9 +3,9 @@ import torch
 import pickle
 import numpy as np
 
-from khaosz.core import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+from khaosz.data.data_util import *
+
 
 def test_dataset_loader_random_paths(base_test_env):
     """Test dataset loader with multiple random paths"""

tests/test_early_stopping.py

@@ -1,8 +1,7 @@
 import torch
-
-from khaosz.core import *
+from khaosz.config import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+
 
 def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
     """Simulate early stopping behavior"""

tests/test_module.py

@@ -5,8 +5,11 @@ import shutil
 import pytest
 import tempfile
 import safetensors.torch as st
-from khaosz.core import *
-from khaosz.core.generator import EmbeddingEncoderCore, GeneratorCore
+from khaosz.trainer import *
+from khaosz.config import *
+from khaosz.model import *
+from khaosz.data import *
+from khaosz.inference.generator import EmbeddingEncoderCore, GeneratorCore
 from tokenizers import pre_tokenizers
 
 @pytest.fixture

tests/test_sampler.py

@@ -1,14 +1,13 @@
-from khaosz.core import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+from khaosz.data.data_util import *
 
 def test_random_sampler_consistency(random_dataset):
     """Test RandomSampler produces consistent results with same seed"""
     dataset = random_dataset
     
     # Create two samplers with same seed
-    sampler1 = RandomSampler(dataset, seed=42)
-    sampler2 = RandomSampler(dataset, seed=42)
+    sampler1 = ResumeableRandomSampler(dataset, seed=42)
+    sampler2 = ResumeableRandomSampler(dataset, seed=42)
     
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
@@ -20,8 +19,8 @@ def test_random_sampler_different_seeds(random_dataset):
     dataset = random_dataset
     
     # Create two samplers with different seeds
-    sampler1 = RandomSampler(dataset, seed=42)
-    sampler2 = RandomSampler(dataset, seed=123)
+    sampler1 = ResumeableRandomSampler(dataset, seed=42)
+    sampler2 = ResumeableRandomSampler(dataset, seed=123)
     
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
@@ -29,38 +28,13 @@ def test_random_sampler_different_seeds(random_dataset):
     # Very high probability they should be different
     assert indices1 != indices2
 
-def test_sampler_state_persistence(random_dataset):
-    """Test that sampler state is correctly saved and loaded"""
-    dataset = random_dataset
-    n = len(dataset)
-    
-    # Create sampler and get some indices
-    sampler = RandomSampler(dataset, seed=42)
-    iter1 = iter(sampler)
-    indices1 = [next(iter1) for _ in range(min(10, n))]
-    
-    # Save state
-    state_dict = sampler.state_dict()
-    
-    # Get more indices
-    indices2 = [next(iter1) for _ in range(min(10, n - len(indices1)))]
-    
-    # Create new sampler and load state
-    sampler2 = RandomSampler(dataset, seed=42)
-    sampler2.load_state_dict(state_dict)
-    
-    # Check that new sampler produces same sequence from saved point
-    iter2 = iter(sampler2)
-    indices3 = [next(iter2) for _ in range(min(10, n - len(indices1)))]
-    
-    assert indices2 == indices3
 
 def test_sampler_across_epochs(random_dataset):
     """Test sampler behavior across multiple epochs"""
     dataset = random_dataset
     n = len(dataset)
     
-    sampler = RandomSampler(dataset, seed=42)
+    sampler = ResumeableRandomSampler(dataset, seed=42)
     
     # Get indices for first epoch
     epoch1_indices = list(iter(sampler))

tests/test_train_config.py

@@ -1,9 +1,10 @@
 import torch
 import numpy as np
 
-from khaosz.core import *
+
+from khaosz.config import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+from khaosz.data.data_util import *
 
 def test_different_batch_sizes(base_test_env, random_dataset):
     """Test training with different batch sizes"""

tests/test_train_strategy.py

@@ -1,9 +1,9 @@
 import torch
 import numpy as np
 
-from khaosz.core import *
+from khaosz.config import *
 from khaosz.trainer import *
-from khaosz.trainer.data_util import *
+from khaosz.data.data_util import *
 
 def test_multi_turn_training(base_test_env, multi_turn_dataset):
     """Test training with multi-turn conversation data"""