| """Interactive chat with a North ML Flare model. |
| |
| REMOTE (recommended — runs the model on an SSH device, nothing to download): |
| python3 chat.py --ssh mac-mini |
| python3 chat.py --ssh mac-mini --model fiber-1 |
| |
| LOCAL (downloads HF weights, runs on this machine): |
| python3 chat.py --model North-ML1/fiber-pro |
| |
| Remote needs: ssh access to the host + NORTH_API_KEY or ~/.config/northml/api.key |
| Local needs: pip install torch transformers safetensors huggingface_hub |
| Commands: /model /search /think /new /quit |
| """ |
| import argparse, json, os, sys, threading, importlib.util |
|
|
| IM_S = "<" + "|im_start" + "|>" |
| IM_E = "<" + "|im_end" + "|>" |
| TH_S = "<" + "think" + ">" |
| TH_E = "<" + "/think" + ">" |
| EOT = "<" + "|endoftext" + "|>" |
|
|
| DIM = "\033[2m"; CYAN = "\033[36m"; BOLD = "\033[1m"; RESET = "\033[0m"; GRAY = "\033[90m" |
| if not sys.stdout.isatty(): |
| DIM = CYAN = BOLD = RESET = GRAY = "" |
|
|
|
|
| |
|
|
| REMOTE_PORT = 8095 |
| LOCAL_PORT = 18095 |
|
|
|
|
| def get_key(): |
| k = os.environ.get("NORTH_API_KEY", "").strip() |
| if k: |
| return k |
| p = os.path.expanduser("~/.config/northml/api.key") |
| if os.path.exists(p): |
| return open(p).read().strip() |
| sys.exit("No API key. Set NORTH_API_KEY or write it to " + p) |
|
|
|
|
| def health(): |
| import urllib.request |
| try: |
| return urllib.request.urlopen( |
| "http://127.0.0.1:%d/health" % LOCAL_PORT, timeout=3).status == 200 |
| except Exception: |
| return False |
|
|
|
|
| def ensure_tunnel(host): |
| import subprocess, time |
| if health(): |
| return |
| print("%sopening ssh tunnel to %s...%s" % (DIM, host, RESET)) |
| subprocess.Popen( |
| ["ssh", "-N", "-L", "%d:127.0.0.1:%d" % (LOCAL_PORT, REMOTE_PORT), |
| "-o", "ExitOnForwardFailure=yes", "-o", "ServerAliveInterval=15", host], |
| stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
| for _ in range(40): |
| if health(): |
| print("%stunnel up on :%d%s" % (DIM, LOCAL_PORT, RESET)) |
| return |
| time.sleep(0.5) |
| sys.exit("Tunnel failed — is `ssh %s` working and the API running?" % host) |
|
|
|
|
| def ask(key, model, messages, search, think): |
| import urllib.request |
| req = urllib.request.Request( |
| "http://127.0.0.1:%d/v1/chat/completions" % LOCAL_PORT, |
| data=json.dumps({"model": model, "messages": messages, "stream": True, |
| "search": search, "think": think}).encode(), |
| headers={"Authorization": "Bearer " + key, |
| "Content-Type": "application/json"}) |
| return urllib.request.urlopen(req, timeout=180) |
|
|
|
|
| def handle_sse(resp): |
| answer, thinking = [], False |
| for raw in resp: |
| line = raw.decode("utf-8", "replace").strip() |
| if not line.startswith("data: "): |
| continue |
| payload = line[6:] |
| if payload == "[DONE]": |
| break |
| try: |
| ev = json.loads(payload) |
| except json.JSONDecodeError: |
| continue |
| agent = ev.get("agent") |
| if agent: |
| t = agent.get("type") |
| if t == "status": |
| print("\n%s ▸ %s%s" % (GRAY, agent.get("stage", t), RESET), flush=True) |
| elif t == "route": |
| print("%s ▸ route: search=%s think=%s%s" % ( |
| GRAY, agent.get("search"), agent.get("think"), RESET), flush=True) |
| elif t == "search": |
| print("%s ▸ query: %s [%s]%s" % ( |
| GRAY, agent.get("query", ""), agent.get("engine", ""), RESET), flush=True) |
| elif t == "source": |
| mark = "✓" if agent.get("kept") else "·" |
| print("%s %s %s %s%s" % (GRAY, mark, agent.get("title", "?"), |
| agent.get("url", ""), RESET)) |
| elif t == "thinking": |
| print("%s%s%s" % (DIM, agent.get("text", ""), RESET), end="", flush=True) |
| thinking = True |
| elif t == "verify": |
| print("%s ▸ verify: %s%s" % ( |
| GRAY, str(agent.get("verdict", agent.get("text", "")))[:160], RESET), |
| flush=True) |
| elif t == "error": |
| print("\033[31m error: %s%s" % (agent.get("message", agent), RESET)) |
| continue |
| for ch in ev.get("choices") or []: |
| d = ch.get("delta", {}).get("content") |
| if d: |
| if thinking: |
| print() |
| thinking = False |
| print(d, end="", flush=True) |
| answer.append(d) |
| print() |
| return "".join(answer) |
|
|
|
|
| def remote_loop(host, model): |
| import urllib.error |
| key = get_key() |
| ensure_tunnel(host) |
| search, think = "auto", "auto" |
| history = [] |
| print("%sNorth ML chat — model on %s%s" % (BOLD, host, RESET)) |
| print("model: %s%s%s · /model /search /think /new /quit\n" % (CYAN, model, RESET)) |
| while True: |
| try: |
| q = input("%syou > %s" % (BOLD, RESET)).strip() |
| except (EOFError, KeyboardInterrupt): |
| print() |
| break |
| if not q: |
| continue |
| if q.startswith("/"): |
| parts = q.split() |
| cmd = parts[0] |
| if cmd in ("/quit", "/q"): |
| break |
| if cmd == "/new": |
| history = [] |
| print("%s(history cleared)%s" % (DIM, RESET)) |
| continue |
| if cmd == "/model" and len(parts) > 1: |
| model = parts[1] |
| continue |
| if cmd == "/search" and len(parts) > 1: |
| search = parts[1] |
| continue |
| if cmd == "/think" and len(parts) > 1: |
| think = parts[1] |
| continue |
| print("%scommands: /model ink-1|fiber-1|fiber-1-pro " |
| "/search auto|on|off /think auto|on|off /new /quit%s" % (DIM, RESET)) |
| continue |
| history.append({"role": "user", "content": q}) |
| try: |
| ans = handle_sse(ask(key, model, history, search, think)) |
| history.append({"role": "assistant", "content": ans}) |
| except urllib.error.HTTPError as e: |
| print("%sHTTP %s: %s%s" % (DIM, e.code, e.read()[:200].decode("utf-8", "replace"), RESET)) |
| history.pop() |
| except Exception as e: |
| print("%serror: %s%s" % (DIM, e, RESET)) |
| history.pop() |
|
|
|
|
| |
|
|
| def pick_device(): |
| import torch |
| if torch.cuda.is_available(): |
| return "cuda", torch.bfloat16 |
| if torch.backends.mps.is_available(): |
| return "mps", torch.float16 |
| return "cpu", torch.float32 |
|
|
|
|
| def load(repo): |
| import torch |
| from huggingface_hub import snapshot_download |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| local = snapshot_download(repo) if not os.path.isdir(repo) else repo |
| device, dtype = pick_device() |
| tok = AutoTokenizer.from_pretrained(local) |
| model = AutoModelForCausalLM.from_pretrained(local, dtype=dtype).to(device).eval() |
|
|
| arch_p = os.path.join(local, "starlight_arch.py") |
| reg_p = os.path.join(local, "agent_register.pt") |
| meta_p = os.path.join(local, "starlight.json") |
| if all(os.path.exists(p) for p in (arch_p, reg_p, meta_p)): |
| spec = importlib.util.spec_from_file_location("starlight_arch", arch_p) |
| arch = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(arch) |
| meta = json.load(open(meta_p)) |
| regs = arch.attach_registers(model, arch.hidden_size_of(model), |
| slots=int(meta.get("slots", 16))) |
| regs.load_state_dict(torch.load(reg_p, map_location="cpu")["registers"]) |
| regs.to(device=device, dtype=dtype) |
| print("registers attached (%s slots)" % meta.get("slots", 16)) |
| print("loaded on %s (%s)" % (device, dtype)) |
| return model, tok, device |
|
|
|
|
| def render(history, think): |
| s = "" |
| for role, text in history: |
| s += IM_S + role + "\n" + text + IM_E + "\n" |
| s += IM_S + "assistant\n" |
| if think: |
| s += TH_S + "\n" |
| return s |
|
|
|
|
| def gen_stream(model, tok, device, input_ids, cap, eos_ids): |
| import torch |
| from transformers import TextIteratorStreamer |
| streamer = TextIteratorStreamer(tok, skip_special_tokens=True, skip_prompt=True) |
| holder = {} |
|
|
| def _run(): |
| holder["out"] = model.generate( |
| input_ids=input_ids, max_new_tokens=cap, do_sample=False, |
| repetition_penalty=1.05, eos_token_id=eos_ids, |
| pad_token_id=tok.eos_token_id, streamer=streamer) |
|
|
| t = threading.Thread(target=_run) |
| t.start() |
| text = "" |
| for chunk in streamer: |
| text += chunk |
| sys.stdout.write(chunk) |
| sys.stdout.flush() |
| t.join() |
| return text, holder["out"] |
|
|
|
|
| def chat_turn(model, tok, device, history, think): |
| import torch |
| im_end = tok.convert_tokens_to_ids(IM_E) |
| eot = tok.convert_tokens_to_ids(EOT) |
| th_e = tok.convert_tokens_to_ids(TH_E) |
| eos = list({i for i in (tok.eos_token_id, im_end, eot) if i is not None}) |
|
|
| prompt = render(history, think) |
| ids = tok(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device) |
|
|
| if think: |
| sys.stdout.write("\n[think] ") |
| th_text, out1 = gen_stream(model, tok, device, ids, 220, eos + [th_e]) |
| if th_e not in out1[0].tolist()[ids.shape[1]:]: |
| extra = tok(TH_E + "\n", add_special_tokens=False, |
| return_tensors="pt").input_ids.to(device) |
| ids2 = torch.cat([out1, extra], dim=1) |
| else: |
| ids2 = out1 |
| sys.stdout.write("\n[answer] ") |
| ans_text, _ = gen_stream(model, tok, device, ids2, 500, eos) |
| sys.stdout.write("\n") |
| return ans_text |
| ans_text, _ = gen_stream(model, tok, device, ids, 500, eos) |
| sys.stdout.write("\n") |
| return ans_text |
|
|
|
|
| def local_loop(repo, think): |
| print("downloading %s ..." % repo) |
| model, tok, device = load(repo) |
| history = [] |
| print("\nReady. Ask anything. /think on|off /new /quit\n") |
| while True: |
| try: |
| q = input("you> ").strip() |
| except (EOFError, KeyboardInterrupt): |
| print() |
| break |
| if not q: |
| continue |
| if q in ("/quit", "/q", "exit"): |
| break |
| if q == "/new": |
| history = [] |
| print("(cleared)\n") |
| continue |
| if q.startswith("/think"): |
| think = "off" not in q |
| print("think =", "on" if think else "off", "\n") |
| continue |
| history.append(("user", q)) |
| ans = chat_turn(model, tok, device, history, think) |
| history.append(("assistant", ans.strip())) |
| sys.stdout.write("\n") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ssh", metavar="HOST", default="", |
| help="ssh host running the North ML API (e.g. mac-mini)") |
| ap.add_argument("--model", default="", |
| help="remote: ink-1|fiber-1|fiber-1-pro · local: HF repo id") |
| ap.add_argument("--no-think", action="store_true") |
| a = ap.parse_args() |
|
|
| if a.ssh: |
| remote_loop(a.ssh, a.model or "fiber-1-pro") |
| else: |
| local_loop(a.model or "North-ML1/fiber-pro", not a.no_think) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
