File size: 1,743 Bytes
36a88b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | #!/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()
|