| |
| """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() |
|
|
