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

scripts/demo/generate_ar.py

@@ -2,24 +2,32 @@ from pathlib import Path
 
 import torch
 
-from astrai.config.param_config import ModelParameter
 from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def generate_text():
-    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
-    param.to(device="cuda", dtype=torch.bfloat16)
+    # Load model from pretrained
+    model = AutoModel.from_pretrained(PARAMETER_ROOT)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    # Load tokenizer from pretrained
+    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT / "tokenizer")
 
     query = input(">> ")
 
-    engine = InferenceEngine(param)
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+    )
     response = engine.generate(
         prompt=query,
         stream=False,
-        max_tokens=param.config.max_len,
+        max_tokens=2048,
         temperature=0.8,
         top_p=0.95,
         top_k=50,

scripts/demo/generate_batch.py

@@ -2,7 +2,7 @@ from pathlib import Path
 
 import torch
 
-from astrai.config.param_config import ModelParameter
+from astrai.model import AutoModel
 from astrai.inference import InferenceEngine
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
@@ -10,8 +10,10 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def batch_generate():
-    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
-    param.to(device="cuda", dtype=torch.bfloat16)
+    # Load model using AutoModel
+    model = AutoModel.from_pretrained(
+        PARAMETER_ROOT, device="cuda", dtype=torch.bfloat16
+    )
 
     inputs = [
         "你好",
@@ -21,11 +23,14 @@ def batch_generate():
         "请问什么是显卡",
     ]
 
-    engine = InferenceEngine(param)
+    engine = InferenceEngine(
+        model=model.model,
+        tokenizer=model.tokenizer,
+    )
     responses = engine.generate(
         prompt=inputs,
         stream=False,
-        max_tokens=param.config.max_len,
+        max_tokens=model.config.max_len,
         temperature=0.8,
         top_p=0.95,
         top_k=50,

scripts/demo/stream_chat.py

@@ -1,32 +1,39 @@
 from pathlib import Path
 
 import torch
-
-from astrai.config.param_config import ModelParameter
 from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def chat():
-    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
-    param.to(device="cuda", dtype=torch.bfloat16)
+    model = AutoModel.from_pretrained(PARAMETER_ROOT)
+    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
+    model.to(device="cuda", dtype=torch.bfloat16)
 
-    history = []
-    engine = InferenceEngine(param)
+    messages = []
+    engine = InferenceEngine(model=model, tokenizer=tokenizer)
 
     while True:
         query = input(">> ")
         if query == "!exit":
             break
 
+        # Add user message
+        messages.append({"role": "user", "content": query})
+
+        # Generate response
         full_response = ""
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
 
         for token in engine.generate(
-            prompt=query,
+            prompt=prompt,
             stream=True,
-            max_tokens=param.config.max_len,
+            max_tokens=model.config.max_len,
             temperature=0.8,
             top_p=0.95,
             top_k=50,
@@ -35,7 +42,8 @@ def chat():
             full_response += token
 
         print()
-        history.append((query, full_response.strip()))
+        # Add assistant response to messages
+        messages.append({"role": "assistant", "content": full_response.strip()})
 
 
 if __name__ == "__main__":

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}