ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

style: 使用ruff 工具优化代码风格

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 23:32:28 +08:00
parent 345fd2f091
commit 426af2d75f
52 changed files with 1836 additions and 1493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/module/test_module.py
+32 -23
View File
@@ -12,6 +12,7 @@ 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__
@@ -19,7 +20,7 @@ def test_env(request: pytest.FixtureRequest):
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,
@@ -28,20 +29,20 @@ def test_env(request: pytest.FixtureRequest):
"dim_ffn": 256,
"max_len": 64,
"n_layers": 2,
"norm_eps": 1e-5
"norm_eps": 1e-5,
}
with open(config_path, 'w') as f:
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,
@@ -51,47 +52,55 @@ def test_env(request: pytest.FixtureRequest):
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 = ModelParameter(
test_env["model"], test_env["tokenizer"], test_env["transformer_config"]
)
ModelParameter.save(model_param, 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))
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)
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"]
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"]
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))
@@ -102,8 +111,8 @@ def test_generator_core(test_env):
top_p=0.95,
attn_mask=None,
kv_caches=None,
start_pos=0
start_pos=0,
)
assert next_token_id.shape == (4, 1)
assert cache_increase == 10
tests/module/test_tie_weight.py
+30 -34
View File
@@ -13,7 +13,7 @@ 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,
@@ -22,18 +22,14 @@ def transformer_test_env():
"dim_ffn": 256,
"max_len": 64,
"n_layers": 2,
"norm_eps": 1e-5
"norm_eps": 1e-5,
}
with open(config_path, 'w') as f:
with open(config_path, "w") as f:
json.dump(config, f)
yield {
"test_dir": test_dir,
"config_path": config_path,
"config": config
}
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):
@@ -46,74 +42,75 @@ def transformer_test_env():
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:
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:
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:
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:
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))
@@ -121,4 +118,3 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
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()