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
@@ -4,7 +4,7 @@ import torch
import torch.nn as nn
from torch import Tensor
from astrai.config.model_config import ModelConfig
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.inference.core.cache import KvcacheView
from astrai.model.automodel import AutoModel
from astrai.model.components.decoder_block import DecoderBlock
@@ -46,11 +46,11 @@ def process_attention_mask(
).masked_fill_(attend.unsqueeze(1), 0.0)
@AutoModel.register("transformer")
class Transformer(AutoModel):
"""Transformer language model with paged KV cache."""
@AutoModel.register("autoregressive_lm")
class AutoRegressiveLM(AutoModel):
"""Autoregressive language model with paged KV cache."""
def __init__(self, config: ModelConfig):
def __init__(self, config: AutoRegressiveLMConfig):
super().__init__(config)
self.config = config
rope_dim = (