File size: 2,491 Bytes
dde517d 8414bb4 dde517d 8414bb4 dde517d 8414bb4 dde517d 8414bb4 dde517d | 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 58 59 60 61 62 63 64 65 | #!/usr/bin/env python3
"""Terminal demo for CBT-18.95M. Sums like `123+456` go to the adder; anything else is a chat turn.
python chat.py # interactive
python chat.py "98765432109876+1234" # one prompt
python chat.py --online "Who was Ada Lovelace?"
python chat.py --variant twin_1d "31415926535897932384626433832795+27182818284590452353602874713527"
"""
import argparse
import functools
import re
from cbt import add, chat, execute_tool, load
SUM = re.compile(r"^\s*(\d+)\s*\+\s*(\d+)\s*=?\s*$")
def respond(model, tok, text: str, online: bool, show_raw: bool) -> None:
m = SUM.match(text)
if m:
r = add(model, tok, m.group(1), m.group(2))
print(f" scratchpad : {r['scratchpad']}")
print(f" answer : {r['answer']} [{'PASS' if r['correct'] else 'FAIL'}]")
else:
r = chat(model, tok, text, tool_executor=functools.partial(execute_tool, online=online))
for key in ("think", "call", "result", "answer"):
if r[key] is not None:
print(f" {key:<10} : {r[key]}")
if show_raw:
print(f" raw : {r['raw']}")
print()
def main() -> None:
p = argparse.ArgumentParser(description="CBT-18.95M terminal demo")
p.add_argument("prompt", nargs="?", help="single prompt; omit for interactive mode")
p.add_argument("--model", default=None, help="local folder or Hugging Face repo id (default: this folder)")
p.add_argument("--variant", default="cbt_2d", choices=["cbt_2d", "twin_1d"], help="released 2D model or its 1D-RoPE twin")
p.add_argument("--device", default=None, help="cpu or cuda (default: auto)")
p.add_argument("--online", action="store_true", help="let the search tool query Wikipedia")
p.add_argument("--raw", action="store_true", help="also print the raw token stream")
args = p.parse_args()
model, tok = load(args.model, args.device, args.variant)
if args.prompt:
respond(model, tok, args.prompt, args.online, args.raw)
return
print(f"CBT-18.95M ({args.variant}). Type a sum (e.g. 123456789+987654321) or a question. Ctrl-D to quit.\n")
while True:
try:
text = input("> ").strip()
except (EOFError, KeyboardInterrupt):
print()
break
if text.lower() in ("exit", "quit"):
break
if text:
respond(model, tok, text, args.online, args.raw)
if __name__ == "__main__":
main()
|