ViperEkura
e13fe53475
- 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
102 lines
2.6 KiB
Python
102 lines
2.6 KiB
Python
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
|
|
from astrai import build_engine
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser(description="Interactive streaming chat")
|
|
parser.add_argument(
|
|
"--model_path",
|
|
type=Path,
|
|
default=PROJECT_ROOT / "params",
|
|
help="Path to model weights (params/ or checkpoint/epoch_N_step_M/)",
|
|
)
|
|
parser.add_argument(
|
|
"--temperature",
|
|
type=float,
|
|
default=0.8,
|
|
help="Sampling temperature (default: 0.8)",
|
|
)
|
|
parser.add_argument(
|
|
"--top_p",
|
|
type=float,
|
|
default=0.95,
|
|
help="Top-p sampling threshold",
|
|
)
|
|
parser.add_argument(
|
|
"--top_k",
|
|
type=int,
|
|
default=50,
|
|
help="Top-k sampling threshold",
|
|
)
|
|
parser.add_argument(
|
|
"--max_tokens",
|
|
type=int,
|
|
default=2048,
|
|
help="Maximum tokens to generate",
|
|
)
|
|
parser.add_argument(
|
|
"--frequency_penalty",
|
|
type=float,
|
|
default=0.5,
|
|
help="Penalty per occurrence for repeated tokens (0.0 disables, "
|
|
"range -2.0~2.0, typical 0.3-1.0)",
|
|
)
|
|
parser.add_argument(
|
|
"--rep_window",
|
|
type=int,
|
|
default=64,
|
|
help="Number of recent prompt tokens to include in penalty history",
|
|
)
|
|
parser.add_argument(
|
|
"--system_prompt",
|
|
type=str,
|
|
default="",
|
|
help="Optional system prompt (default: empty, model not SFT-trained on system role)",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def chat():
|
|
args = parse_args()
|
|
model_path = args.model_path
|
|
|
|
engine = build_engine(model_path)
|
|
tokenizer = engine.tokenizer
|
|
|
|
while True:
|
|
query = input(">> ")
|
|
if query == "!exit":
|
|
break
|
|
|
|
msgs = []
|
|
if args.system_prompt:
|
|
msgs.append({"role": "system", "content": args.system_prompt})
|
|
msgs.append({"role": "user", "content": query})
|
|
prompt = tokenizer.apply_chat_template(
|
|
msgs, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
full_response = ""
|
|
for token in engine.generate(
|
|
prompt=prompt,
|
|
stream=True,
|
|
max_tokens=args.max_tokens,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
frequency_penalty=args.frequency_penalty,
|
|
rep_window=args.rep_window,
|
|
):
|
|
print(token, end="", flush=True)
|
|
full_response += token
|
|
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
chat()
|