feat: 实现模型动态注册机制

2026-04-05 19:38:12 +08:00
parent ff43a2fab8
commit fc278d17ab
25 changed files with 686 additions and 651 deletions
@@ -1,34 +0,0 @@
-import os
-
-import torch
-
-from astrai.config.param_config import ModelParameter
-
-
-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"]
-    )
-    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),
-    )
-    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