refactor: 更新inference 部分的实现

tests/module/test_module.py

@@ -3,8 +3,6 @@ import os
 import torch
 
 from astrai.config.param_config import ModelParameter
-from astrai.inference.generator import EmbeddingEncoderCore, GeneratorCore
-
 
 def test_model_parameter(test_env):
     save_dir = os.path.join(test_env["test_dir"], "save")
@@ -32,40 +30,4 @@ def test_transformer(test_env):
         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
+    assert output_logits.shape == target_shape