refactor(tests): 重构测试文件目录结构

2026-01-08 21:34:52 +08:00
parent d407962ffa
commit 4da70785b5
9 changed files with 0 additions and 0 deletions
@@ -0,0 +1,109 @@
+import os
+import json
+import torch
+import shutil
+import pytest
+import tempfile
+import safetensors.torch as st
+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
+def test_env(request: pytest.FixtureRequest):
+    func_name = request.function.__name__
+    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
+    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,
+        "dim": 128,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "dim_ffn": 256,
+        "max_len": 64,
+        "n_layers": 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 = ModelConfig().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)
+
+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"].max_len))
+    output_logits = model(input_ids)["logits"]
+    target_shape = (4, test_env["transformer_config"].max_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"].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)
+    input_ids = torch.randint(0, test_env["transformer_config"].vocab_size, (4, 10))
+    next_token_id, cache_increase = generator.generate_iterator(
+        input_ids=input_ids,
+        temperature=0.8,
+        top_k=50,
+        top_p=0.95,
+        attn_mask=None,
+        kv_caches=None,
+        start_pos=0
+    )
+    
+    assert next_token_id.shape == (4, 1)
+    assert cache_increase == 10
@@ -0,0 +1,124 @@
+import os
+import json
+import torch
+import pytest
+import tempfile
+import safetensors.torch as st
+from khaosz.model.transformer import Transformer
+from khaosz.config.model_config import ModelConfig
+
+
+@pytest.fixture
+def transformer_test_env():
+    """创建Transformer测试专用环境"""
+    test_dir = tempfile.mkdtemp(prefix="transformer_test_")
+    config_path = os.path.join(test_dir, "config.json")
+    
+    config = {
+        "vocab_size": 1000,
+        "dim": 128,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "dim_ffn": 256,
+        "max_len": 64,
+        "n_layers": 2,
+        "norm_eps": 1e-5
+    }
+    
+    with open(config_path, 'w') as f:
+        json.dump(config, f)
+    
+    yield {
+        "test_dir": test_dir,
+        "config_path": config_path,
+        "config": config
+    }
+    
+    if os.path.exists(test_dir):
+        try:
+            for file in os.listdir(test_dir):
+                os.remove(os.path.join(test_dir, file))
+            os.rmdir(test_dir)
+        except:
+            pass
+
+
+def test_tie_weight_init(transformer_test_env):
+    config_path = transformer_test_env["config_path"]
+    config_data = transformer_test_env["config"].copy()
+    
+    # case 1: tie weight
+    config_data["tie_weight"] = True
+    
+    with open(config_path, 'w') as f:
+        json.dump(config_data, f)
+    
+    config = ModelConfig().load(config_path)
+    model = Transformer(config)
+    
+    assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
+    
+    original_weight = model.embed_tokens.weight.clone()
+    model.embed_tokens.weight.data[0, 0] = 100.0
+    
+    assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert not torch.equal(model.lm_head.weight, original_weight)
+    
+    # case 2: not tie weight
+    config_data["tie_weight"] = False
+    
+    with open(config_path, 'w') as f:
+        json.dump(config_data, f)
+    
+    config = ModelConfig().load(config_path)
+    model = Transformer(config)
+    
+    assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
+
+    original_weight = model.embed_tokens.weight.clone()
+    model.embed_tokens.weight.data[0, 0] = 100.0
+    
+    assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert not torch.equal(model.lm_head.weight, original_weight)
+
+def test_model_save_load_with_tie_weight(transformer_test_env):
+    test_dir = transformer_test_env["test_dir"]
+    model_path = os.path.join(test_dir, "model.safetensors")
+    
+    config_data = transformer_test_env["config"].copy()
+    
+    # case 1: tie weight
+    config_data["tie_weight"] = True
+    config_path = os.path.join(test_dir, "config.json")
+    
+    with open(config_path, 'w') as f:
+        json.dump(config_data, f)
+    
+    config = ModelConfig().load(config_path)
+    original_model = Transformer(config)
+    
+    st.save_file(original_model.state_dict(), model_path)
+
+    loaded_config = ModelConfig().load(config_path)
+    model = Transformer(loaded_config)
+    model.load_state_dict(st.load_file(model_path))
+    
+    assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
+    assert "lm_head.weight" not in model.state_dict()
+
+    # case 2: not tie weight (form tie-weight state dict load)
+    config_data["tie_weight"] = False
+    with open(config_path, 'w') as f:
+        json.dump(config_data, f)
+    
+    loaded_config = ModelConfig().load(config_path)
+    model = Transformer(loaded_config)
+    model.load_state_dict(st.load_file(model_path))
+
+    assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
+    assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
+    assert "lm_head.weight" in model.state_dict()
+