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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user