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(">> ")