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
+6 -6
View File
@@ -8,16 +8,16 @@ import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from astrai.config import ModelConfig, TrainConfig
from astrai.config import AutoRegressiveLMConfig, TrainConfig
from astrai.dataset import DatasetFactory
from astrai.model import Transformer
from astrai.model import AutoRegressiveLM
from astrai.parallel import get_rank
from astrai.trainer import SchedulerFactory, Trainer
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Train the Transformer model.")
parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
parser.add_argument(
"--train_type",
@@ -246,13 +246,13 @@ def train(
# Load config
config_path = os.path.join(param_path, "config.json")
config = ModelConfig.from_file(config_path)
config = AutoRegressiveLMConfig.from_file(config_path)
if window_size is None:
window_size = config.max_len
# Create bare Transformer (for training, no tokenizer needed)
model = Transformer(config)
# Create bare AutoRegressiveLM (for training, no tokenizer needed)
model = AutoRegressiveLM(config)
# Load weights if available
weights_path = os.path.join(param_path, "model.safetensors")