Initial commit

2025-09-27 12:02:22 +08:00
commit a4443765ee
33 changed files with 3896 additions and 0 deletions
@@ -0,0 +1,103 @@
+import os
+import json
+import torch
+import shutil
+import pytest
+import tempfile
+import safetensors.torch as st
+from khaosz.core import *
+from khaosz.core.generator import EmbeddingEncoderCore, GeneratorCore
+from tokenizers import pre_tokenizers
+
+@pytest.fixture
+def test_env():
+    test_dir = tempfile.mkdtemp()
+    config_path = os.path.join(test_dir, "config.json")
+    tokenizer_path = os.path.join(test_dir, "tokenizer.json")
+    model_path = os.path.join(test_dir, "model.safetensors")
+    
+    config = {
+        "vocab_size": 1000,
+        "n_dim": 128,
+        "n_head": 4,
+        "n_kvhead": 2,
+        "d_ffn": 256,
+        "m_len": 64,
+        "n_layer": 2,
+        "norm_eps": 1e-5
+    }
+    with open(config_path, 'w') as f:
+        json.dump(config, f)
+    
+    tokenizer = BpeTokenizer()
+    sp_token_iter = iter(pre_tokenizers.ByteLevel.alphabet())
+    tokenizer.train_from_iterator(sp_token_iter, config["vocab_size"], 1)
+    tokenizer.save(tokenizer_path)
+    
+    transformer_config = TransformerConfig().load(config_path)
+    model = Transformer(transformer_config)
+    st.save_file(model.state_dict(), model_path)
+    
+    yield {
+        "test_dir": test_dir,
+        "model": model,
+        "tokenizer": tokenizer,
+        "transformer_config": transformer_config,
+    }
+
+    shutil.rmtree(test_dir)
+
+# parameter loader
+def test_parameter_loader(test_env):
+    loaded_param = ParameterLoader.load(test_env["test_dir"])
+    assert loaded_param.model is not None
+    assert loaded_param.tokenizer is not None
+    assert loaded_param.config == test_env["transformer_config"]
+
+def test_model_parameter(test_env):
+    save_dir = os.path.join(test_env["test_dir"], "save")
+    model_param = ModelParameter(test_env["model"],test_env["tokenizer"] , test_env["transformer_config"])
+    model_param.save(save_dir)
+    
+    assert os.path.exists(os.path.join(save_dir, "model.safetensors"))
+    assert os.path.exists(os.path.join(save_dir, "tokenizer.json"))
+    assert os.path.exists(os.path.join(save_dir, "config.json"))
+
+# transformer
+def test_transformer(test_env):
+    model = test_env["model"]
+    input_ids = torch.randint(0, test_env["transformer_config"].vocab_size, 
+                              (4, test_env["transformer_config"].m_len))
+    output_logits = model(input_ids)["logits"]
+    target_shape = (4, test_env["transformer_config"].m_len, test_env["transformer_config"].vocab_size)
+    assert output_logits.shape == target_shape
+    
+# generator
+def test_embedding_encoder_core(test_env):
+    parameter = ModelParameter(
+        test_env["model"],
+        test_env["tokenizer"],
+        test_env["transformer_config"]
+    )
+    encoder = EmbeddingEncoderCore(parameter)
+    
+    single_emb = encoder.encode("测试文本")
+    assert isinstance(single_emb, torch.Tensor)
+    assert single_emb.shape[-1] == test_env["transformer_config"].n_dim
+    
+
+    batch_emb = encoder.encode(["测试1", "测试2"])
+    assert isinstance(batch_emb, list)
+    assert len(batch_emb) == 2
+    
+def test_generator_core(test_env):
+    parameter = ModelParameter(
+        test_env["model"],
+        test_env["tokenizer"],
+        test_env["transformer_config"]
+    )
+    generator = GeneratorCore(parameter)
+    logits, incr = generator.compute_logits(torch.randint(0, test_env["transformer_config"].vocab_size, (4, 10)))
+    
+    assert logits.shape == (4, test_env["transformer_config"].vocab_size)
+    assert incr == 10
@@ -0,0 +1,203 @@
+import os
+import json
+import torch
+import shutil
+import pytest
+import pickle
+import tempfile
+import matplotlib
+
+from torch.utils.data import Dataset
+from khaosz.core import *
+from khaosz.trainer import *
+
+# to avoid _tkinter.TclError
+matplotlib.use('Agg')
+
+
+@pytest.fixture
+def test_env():
+    test_dir = tempfile.mkdtemp()
+    config_path = os.path.join(test_dir, "config.json")
+    
+    config = {
+        "vocab_size": 1000,
+        "n_dim": 128,
+        "n_head": 4,
+        "n_kvhead": 2,
+        "d_ffn": 256,
+        "m_len": 64,
+        "n_layer": 2,
+        "norm_eps": 1e-5
+    }
+    
+    with open(config_path, 'w') as f:
+        json.dump(config, f)
+    
+    transformer_config = TransformerConfig().load(config_path)
+    model = Transformer(transformer_config)
+    tokenizer = BpeTokenizer()
+    
+    class DummyDataset(Dataset):
+        def __init__(self, length=10):
+            self.length = length
+            
+        def __len__(self):
+            return self.length
+            
+        def __getitem__(self, idx):
+            return (
+                torch.randint(0, 1000, (64,)),
+                torch.randint(0, 1000, (64,))
+            )
+    
+    dataset = DummyDataset()
+    
+    yield {
+        "test_dir": test_dir,
+        "config_path": config_path,
+        "transformer_config": transformer_config,
+        "model": model,
+        "tokenizer": tokenizer,
+        "dataset": dataset
+    }
+    
+    shutil.rmtree(test_dir)
+
+def test_dataset_loader(test_env):
+    test_dir = test_env["test_dir"]
+    pkl_path = os.path.join(test_dir, "test_data.pkl")
+    
+    dummy_data = {"sequence": torch.randint(0, 1000, (64,))}
+    with open(pkl_path, "wb") as f:
+        pickle.dump(dummy_data, f)
+    
+    loaded_dataset = DatasetLoader.load(train_type="seq", load_path=pkl_path, max_len=64, device="cpu")
+    assert loaded_dataset is not None
+
+def test_training_config(test_env):
+    optimizer = torch.optim.AdamW(test_env["model"].parameters())
+    train_config = TrainConfig(
+        train_type="seq",
+        dataset=test_env["dataset"],
+        optimizer=optimizer,
+        ckpt_dir=test_env["test_dir"],
+        n_epoch=1,
+        batch_size=2,
+        n_iter_ckpt=5,
+        n_iter_step=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    assert train_config.get_kwargs()["batch_size"] == 2
+
+def test_cosine_schedule(test_env):
+    assert test_env is not None
+    schedule_config = CosineScheduleConfig(
+        warning_step=100,
+        total_iters=1000
+    )
+    kwargs = schedule_config.get_kwargs()
+    assert kwargs["warning_step"] == 100
+    assert kwargs["lr_decay_iters"] == 900
+    
+
+def test_sgdr_schedule(test_env):
+    assert test_env is not None
+    schedule_config = SgdrScheduleConfig(
+        warning_step=100,
+        cycle_length=200,
+        T_mult=2
+    )
+    kwargs = schedule_config.get_kwargs()
+    assert kwargs["warning_step"] == 100
+    assert kwargs["cycle_length"] == 200
+    assert kwargs["T_mult"] == 2
+    
+def test_trainer_train(test_env):
+    optimizer = torch.optim.AdamW(test_env["model"].parameters())
+    train_config = TrainConfig(
+        train_type="seq",
+        dataset=test_env["dataset"],
+        optimizer=optimizer,
+        ckpt_dir=test_env["test_dir"],
+        n_epoch=1,
+        batch_size=2,
+        n_iter_ckpt=5,
+        n_iter_step=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    schedule_config = CosineScheduleConfig(
+        warning_step=100,
+        total_iters=1000
+    )
+    model_parameter = ModelParameter(
+        test_env["model"], 
+        test_env["tokenizer"], 
+        test_env["transformer_config"]
+    )
+    trainer = Trainer(model_parameter)
+    trainer.train(train_config, schedule_config)
+
+def test_checkpoint(test_env):
+    temp_dir = test_env["test_dir"]
+    config = test_env["transformer_config"]
+    model = test_env["model"]
+    tokenizer = test_env["tokenizer"]
+    
+    param = ModelParameter(model, tokenizer, config)
+    checkpoint = Checkpoint(
+        model=param.model, 
+        tokenizer=param.tokenizer, 
+        config=param.config, 
+        loss_list=[1.0, 2.0, 3.0], 
+        current_iter=3
+    )
+    ckpt_dir = os.path.join(temp_dir, "ckpt")
+    checkpoint.save(ckpt_dir)
+
+    loaded_ckpt = Checkpoint()
+    loaded_ckpt.load(ckpt_dir)
+    
+    assert loaded_ckpt.current_iter == 3
+    assert loaded_ckpt.loss_list == [1.0, 2.0, 3.0]
+
+    for p1, p2 in zip(model.parameters(), loaded_ckpt.model.parameters()):
+        assert torch.allclose(p1, p2)
+
+
+def test_checkpoint_train(test_env):
+    temp_dir = test_env["test_dir"]
+    config = test_env["transformer_config"]
+    model = test_env["model"]
+    tokenizer = test_env["tokenizer"]
+    dataset = test_env["dataset"]
+    
+    param = ModelParameter(model, tokenizer, config)
+    trainer = Trainer(param)
+    
+    optimizer = torch.optim.AdamW(test_env["model"].parameters())
+    train_config = TrainConfig(
+        train_type="seq",
+        dataset=dataset,
+        optimizer=optimizer,
+        ckpt_dir=test_env["test_dir"],
+        n_epoch=1,
+        batch_size=2,
+        n_iter_ckpt=5,
+        n_iter_step=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    schedule_config = CosineScheduleConfig(
+        warning_step=100,
+        total_iters=1000
+    )
+    
+    trainer.train(
+        train_config=train_config,
+        schedule_config=schedule_config,
+    )
+    
+