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
+6 -19
View File
@@ -18,13 +18,10 @@ from math import prod
from typing import Dict, Iterator, List, Optional, Sequence, Tuple
import numpy as np
import torch
import tqdm
from datasets import load_dataset
from astrai.inference import InferenceEngine
from astrai.model import AutoModel
from astrai.tokenize import AutoTokenizer
from astrai.inference import build_engine
# ---------------------------------------------------------------------------
# Config
@@ -91,20 +88,6 @@ def save_json(path: str, data):
json.dump(data, f, indent=2, ensure_ascii=False)
def create_engine(
param_path: str, batch_size: int, max_seq_len: int
) -> InferenceEngine:
model = AutoModel.from_pretrained(param_path)
tokenizer = AutoTokenizer.from_pretrained(param_path)
model.to(device="cuda", dtype=torch.bfloat16)
return InferenceEngine(
model=model,
tokenizer=tokenizer,
max_batch_size=batch_size,
max_seq_len=max_seq_len,
)
def trim_stop(text: str) -> str:
for stop in STOP_SEQUENCES:
idx = text.find(stop)
@@ -322,7 +305,11 @@ def run_pipeline(cfg: EvalConfig) -> Dict:
if cfg.problem_indices:
problems = [problems[i] for i in cfg.problem_indices if i < len(problems)]
engine = create_engine(cfg.param_path, cfg.batch_size, cfg.max_seq_len)
engine = build_engine(
cfg.param_path,
max_batch_size=cfg.batch_size,
max_seq_len=cfg.max_seq_len,
)
try:
generated = generate_all(engine, problems, cfg)