refactor: 优化工具脚本接口并修复批处理问题

This commit is contained in:
2026-04-05 21:56:22 +08:00
parent e58dbd7c57
commit 9b22b1651e
8 changed files with 140 additions and 64 deletions
+8 -7
View File
@@ -2,8 +2,9 @@ from pathlib import Path
import torch
from astrai.model import AutoModel
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")
@@ -11,9 +12,9 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def batch_generate():
# Load model using AutoModel
model = AutoModel.from_pretrained(
PARAMETER_ROOT, device="cuda", dtype=torch.bfloat16
)
model = AutoModel.from_pretrained(PARAMETER_ROOT)
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
model.to(device="cuda", dtype=torch.bfloat16)
inputs = [
"你好",
@@ -24,13 +25,13 @@ def batch_generate():
]
engine = InferenceEngine(
model=model.model,
tokenizer=model.tokenizer,
model=model,
tokenizer=tokenizer,
)
responses = engine.generate(
prompt=inputs,
stream=False,
max_tokens=model.config.max_len,
max_tokens=2048,
temperature=0.8,
top_p=0.95,
top_k=50,