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

scripts/demo/generate_ar.py

@@ -1,36 +1,68 @@
+from argparse import ArgumentParser
 from pathlib import Path
 
-import torch
-
-from astrai.inference import InferenceEngine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
+from astrai.inference import build_engine
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
-def generate_text():
-    model = AutoModel.from_pretrained(PARAMETER_ROOT)
-    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
-    model.to(device="cuda", dtype=torch.bfloat16)
-
-    query = input(">> ")
-
-    engine = InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
+def parse_args():
+    parser = ArgumentParser(
+        description="Autoregressive continuation demo: continue a prompt, "
+        "or drop into interactive mode without one"
     )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=None,
+        help="Text prefix to continue; omit to enter interactive mode",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=Path,
+        default=PARAMETER_ROOT,
+        help="Path to model weights",
+    )
+    parser.add_argument("--max_tokens", type=int, default=2048)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top_p", type=float, default=0.95)
+    parser.add_argument("--top_k", type=int, default=50)
+    return parser.parse_args()
+
+
+def run_once(engine, prompt, args):
+    print(prompt, end="", flush=True)
     for token in engine.generate(
-        prompt=query,
+        prompt=prompt,
         stream=True,
-        max_tokens=2048,
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k,
     ):
         print(token, end="", flush=True)
+    print()
+
+
+def main():
+    args = parse_args()
+    engine = build_engine(args.model_path)
+
+    if args.prompt is not None:
+        run_once(engine, args.prompt, args)
+        return
+
+    while True:
+        try:
+            query = input(">> ")
+        except EOFError:
+            break
+        if query == "!exit":
+            break
+        if query:
+            run_once(engine, query, args)
 
 
 if __name__ == "__main__":
-    generate_text()
+    main()

scripts/demo/generate_batch.py

@@ -1,9 +1,6 @@
 from pathlib import Path
 
-import torch
-
-from astrai.inference import InferenceEngine
-from astrai.model import AutoModel
+from astrai.inference import build_engine
 from astrai.tokenize import AutoTokenizer
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
@@ -11,10 +8,7 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def batch_generate():
-    # Load model using AutoModel
-    model = AutoModel.from_pretrained(PARAMETER_ROOT)
     tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
-    model.to(device="cuda", dtype=torch.bfloat16)
 
     inputs = [
         "你好",
@@ -33,10 +27,7 @@ def batch_generate():
         for q in inputs
     ]
 
-    engine = InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
-    )
+    engine = build_engine(PARAMETER_ROOT)
     responses = engine.generate(
         prompt=prompts,
         stream=False,

scripts/demo/stream_chat.py

@@ -1,11 +1,7 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
-import torch
-
-from astrai import InferenceEngine
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
+from astrai import build_engine
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 
@@ -68,10 +64,8 @@ def chat():
     args = parse_args()
     model_path = args.model_path
 
-    model = AutoModel.from_pretrained(model_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model.to(device="cuda", dtype=torch.bfloat16)
-    engine = InferenceEngine(model=model, tokenizer=tokenizer)
+    engine = build_engine(model_path)
+    tokenizer = engine.tokenizer
 
     while True:
         query = input(">> ")

scripts/eval/evaluate_humaneval.py

@@ -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)

scripts/eval/evaluate_ifeval.py

@@ -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,

scripts/tools/generate.py

@@ -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: