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

scripts/tools/generate.py

@@ -3,7 +3,8 @@ import json
 
 import torch
 
-from astrai.config.param_config import ModelParameter
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
 from astrai.inference import InferenceEngine
 
 
@@ -17,9 +18,9 @@ def processor(
     question_key: str,
     response_key: str,
 ):
-    param = ModelParameter.load(model_dir, disable_init=True)
-    param.to(device="cuda", dtype=torch.bfloat16)
-    engine = InferenceEngine(param)
+    # Load model using AutoModel
+    model = AutoModel.from_pretrained(model_dir, device="cuda", dtype=torch.bfloat16)
+    engine = InferenceEngine(model=model.model, tokenizer=model.tokenizer)
 
     with open(input_json_file, "r", encoding="utf-8") as f:
         input_data = [json.loads(line) for line in f]
@@ -29,7 +30,7 @@ def processor(
     responses = engine.generate(
         prompt=queries,
         stream=False,
-        max_tokens=param.config.max_len,
+        max_tokens=model.config.max_len,
         temperature=temperature,
         top_p=top_p,
         top_k=top_k,

scripts/tools/perplexity.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 import tqdm
 from torch import Tensor
 
-from astrai.config.param_config import ModelParameter
+from astrai.model import AutoModel
 
 
 def compute_perplexity(
@@ -20,7 +20,7 @@ def compute_perplexity(
     where PPL = exp(-(1/N) * sum(log P(w_i | w_<i))).
     """
 
-    output = model(input_ids, input_mask)
+    output = model(input_ids, input_mask=input_mask)
     logits = output["logits"]
 
     shifted_logits = logits[:, :-1, :]  # [batch_size, seq_len-1, vocab_size]
@@ -42,10 +42,9 @@ def compute_perplexity(
 def process_file(
     model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
 ):
-    param = ModelParameter.load(model_dir, disable_init=True)
-    param.to(device="cuda", dtype=torch.bfloat16)
-    model = param.model
-    tokenizer = param.tokenizer
+    # Load model using AutoModel
+    model = AutoModel.from_pretrained(model_dir, device="cuda", dtype=torch.bfloat16)
+    tokenizer = model.tokenizer
 
     with open(input_file, "r", encoding="utf-8") as f:
         input_data = [json.loads(line) for line in f]
@@ -54,7 +53,7 @@ def process_file(
     encoded_texts = [tokenizer.encode(text) for text in texts]
     output_data = []
 
-    for i in tqdm(
+    for i in tqdm.tqdm(
         range(0, len(encoded_texts), batch_size), desc="Computing perplexity"
     ):
         batch_encoded = encoded_texts[i : i + batch_size]
@@ -72,7 +71,7 @@ def process_file(
 
         input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
         input_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
-        perplexity = compute_perplexity(model, input_ids, input_mask)
+        perplexity = compute_perplexity(model.model, input_ids, input_mask)
 
         for text, ppl in zip(batch_texts, perplexity):
             output_data.append({text_key: text, "ppl": float(ppl.item())})

scripts/tools/train.py

@@ -5,10 +5,12 @@ from functools import partial
 import torch
 import torch.nn as nn
 import torch.optim as optim
+import safetensors.torch as st
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from astrai.config import ModelParameter, TrainConfig
+from astrai.config import ModelConfig, TrainConfig
 from astrai.dataset import DatasetFactory
+from astrai.model import Transformer
 from astrai.parallel import get_rank
 from astrai.trainer import SchedulerFactory, Trainer
 
@@ -196,12 +198,23 @@ def train(
     assert train_type in ["seq", "sft", "dpo"]
     assert os.path.exists(param_path)
 
-    parameter = ModelParameter.load(param_path)
+    # Load config
+    config = ModelConfig()
+    config_path = os.path.join(param_path, "config.json")
+    if os.path.exists(config_path):
+        config.load(config_path)
 
     if window_size is None:
-        window_size = parameter.config.max_len
+        window_size = config.max_len
 
-    model = parameter.model
+    # Create bare Transformer (for training, no tokenizer needed)
+    model = Transformer(config)
+
+    # Load weights if available
+    weights_path = os.path.join(param_path, "model.safetensors")
+    if os.path.exists(weights_path):
+        state_dict = st.load_file(weights_path)
+        model.load_state_dict(state_dict, strict=False)
 
     strategy_kwargs = {"dpo_beta": dpo_beta, "label_smoothing": label_smoothing}