refactor: Transformer更名为AutoRegressiveLM并新增EmbeddingEncoder

- AutoRegressiveLM 注册名改为 autoregressive_lm
- 新增 EmbeddingEncoder 支持 mean/cls/last pooling
- ModelConfig 增加 pooling_type / normalize_embeddings 字段
- 导入、注释、测试全部同步更新
This commit is contained in:
2026-05-17 15:29:20 +08:00
parent 8f1b32f2b6
commit 97c7ac0f4f
13 changed files with 374 additions and 72 deletions
+5 -5
View File
@@ -8,8 +8,8 @@ import torch
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from torch.utils.data import Dataset
from astrai.config.model_config import ModelConfig
from astrai.model.transformer import Transformer
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.tokenize import AutoTokenizer
@@ -104,8 +104,8 @@ def test_tokenizer():
@pytest.fixture(scope="session")
def test_model():
"""Session-scoped small Transformer model, created once."""
config = ModelConfig(
"""Session-scoped small AutoRegressiveLM model, created once."""
config = AutoRegressiveLMConfig(
vocab_size=1000,
dim=8,
n_heads=2,
@@ -116,7 +116,7 @@ def test_model():
norm_eps=1e-5,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Transformer(config).to(device=device)
model = AutoRegressiveLM(config).to(device=device)
return {
"model": model,