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:
@@ -16,12 +16,10 @@ import os
|
||||
import re
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import tqdm
|
||||
from datasets import load_dataset
|
||||
|
||||
from astrai.inference import InferenceEngine
|
||||
from astrai.model import AutoModel
|
||||
from astrai.inference import InferenceEngine, build_engine
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
IFEVAL_HF_DATASET = "google/IFEval"
|
||||
@@ -536,17 +534,12 @@ def main():
|
||||
print(f"Loaded {len(problems)} problems")
|
||||
print(f"Supported constraint types: {len(CONSTRAINT_VERIFIERS)}")
|
||||
|
||||
model = AutoModel.from_pretrained(args.param_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.param_path)
|
||||
model.to(device="cuda", dtype=torch.bfloat16)
|
||||
model.eval()
|
||||
|
||||
engine = InferenceEngine(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
engine = build_engine(
|
||||
args.param_path,
|
||||
max_batch_size=args.batch_size,
|
||||
max_seq_len=args.max_seq_len,
|
||||
)
|
||||
tokenizer = engine.tokenizer
|
||||
|
||||
results = evaluate(
|
||||
engine=engine,
|
||||
|
||||
Reference in New Issue
Block a user