refactor: assemble inference engines through a shared composition root

- add build_engine() to astrai.inference.engine as the single load-place-wire path for InferenceEngine, accepting a checkpoint path or live model/tokenizer plus passthrough engine kwargs
- migrate the server lifespan, generate CLI, humaneval/ifeval evals, and all three demos to build_engine; app._create_engine collapses into a direct call
- export build_engine from astrai and astrai.inference
- parameterize the autoregressive demo with --prompt one-shot continuation plus model path and sampling knobs, exiting cleanly on !exit or EOF
- cover the composition root with unit tests for live-object assembly, kwargs passthrough, and argument validation
This commit is contained in:
2026-09-03 22:16:56 +08:00
parent 9d3ae76683
commit e13fe53475
12 changed files with 199 additions and 117 deletions
+5 -11
View File
@@ -6,9 +6,7 @@ import click
import torch
from tqdm import tqdm
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai.inference import build_engine
def processor(
@@ -28,17 +26,13 @@ def processor(
):
print(f"Loading model from {param_path} ...")
t0 = time.time()
model = AutoModel.from_pretrained(param_path)
tokenizer = AutoTokenizer.from_pretrained(param_path)
model.to(device="cuda", dtype=torch.bfloat16)
print(f" model loaded in {time.time() - t0:.1f}s")
engine = InferenceEngine(
model=model,
tokenizer=tokenizer,
engine = build_engine(
param_path=param_path,
max_batch_size=batch_size * num_samples,
max_seq_len=max_seq_len,
)
tokenizer = engine.tokenizer
print(f" model loaded in {time.time() - t0:.1f}s")
print(f"Reading {input_json_file} ...")
with open(input_json_file, "r", encoding="utf-8") as f: