#!/usr/bin/env python3
"""Small CLI launcher for LULUV2 native-bf16 local inference."""
from __future__ import annotations
import argparse
import torch
from luluv2_live_inference import LULUV2LiveEngine, GenerationConfig


def main():
    p = argparse.ArgumentParser("LULUV2 local inference")
    p.add_argument("--ckpt", default="LULUV2-bf16.pt", help="Path to the native-bf16 checkpoint file")
    p.add_argument("--tokenizer-dir", default="tokenizer", help="Local tokenizer directory")
    p.add_argument("--prompt", required=True, help="User prompt")
    p.add_argument("--system", default="You are LuluV2, a helpful local AI assistant.")
    p.add_argument("--max-new-tokens", type=int, default=512)
    p.add_argument("--temperature", type=float, default=0.65)
    p.add_argument("--top-p", type=float, default=0.90)
    p.add_argument("--top-k", type=int, default=40)
    p.add_argument("--device", default=None)
    p.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])
    args = p.parse_args()

    engine = LULUV2LiveEngine(
        ckpt_path=args.ckpt,
        model_py="luluv2_inference_runtime.py",
        tokenizer_dir=args.tokenizer_dir,
        device=args.device,
        dtype=args.dtype,
        local_files_only=True,
        no_config_download=True,
    )
    cfg = GenerationConfig(
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
    )
    history = []
    for text in engine.generate_stream(args.prompt, history, args.system, cfg):
        print(text, end="", flush=True)
    print()


if __name__ == "__main__":
    torch.set_grad_enabled(False)
    main()