chat_lilo.py

#!/usr/bin/env python3
"""LILO chat CLI — CPU-optimized, KV-cache generation.

Usage:
    python chat_lilo.py "prompt"                       # single prompt
    python chat_lilo.py                                # interactive
    python chat_lilo.py --temp 0.7 --top-p 0.9 --max 200 "prompt"
"""
import argparse, os, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from lilo_common import load_lilo, generate_lilo


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("prompt", nargs="*")
    ap.add_argument("--variant", default="fp16", choices=["fp16", "int8"])
    ap.add_argument("--max", type=int, default=200)
    ap.add_argument("--temp", type=float, default=0.7)
    ap.add_argument("--top-p", type=float, default=0.9)
    ap.add_argument("--rep", type=float, default=1.1)
    ap.add_argument("--threads", type=int, default=1)
    ap.add_argument("--path", default=HERE)
    a = ap.parse_args()

    tok, model = load_lilo(a.path, variant=a.variant, threads=a.threads)

    def run(prompt):
        ids = tok.encode(prompt, add_special_tokens=False)
        out = generate_lilo(model, tok, ids, max_new_tokens=a.max, temperature=a.temp,
                            top_p=a.top_p, repetition_penalty=a.rep, stream=True)
        return tok.decode(out, skip_special_tokens=True)

    if a.prompt:
        run(" ".join(a.prompt))
        print()
        return

    print("LILO siap. Ketik pesan (Ctrl-C/Ctrl-D untuk keluar).")
    while True:
        try:
            p = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not p:
            continue
        run(p)
        print()


if __name__ == "__main__":
    main()