Upload chat_sft.py with huggingface_hub

chat_sft.py

@@ -0,0 +1,135 @@
+"""
+chat_sft.py -- interactive multi-turn chat REPL for an SFT'd Quazimoto-LM.
+
+Renders the running conversation in the student's ChatML and generates the
+assistant reply. Two decode paths:
+
+  * SPECULATIVE (default) -- DeepSpec/DSpark-style self-speculative decoding: the
+    MTP heads draft the next few tokens, the main head verifies them in one pass,
+    accepting the longest correct prefix (see model.forward_drafts / generate.py
+    generate_speculative). Greedy, so output is deterministic.
+  * SAMPLED (--temperature > 0) -- the KV-cache sampler (generate.py generate):
+    the conversation is prefilled once, then each new token is a single-position
+    cached forward. Supports temperature / top-k / top-p.
+
+Both reuse the model's KV cache; speculative additionally drafts multiple tokens
+per verify. Usage:
+    python chat_sft.py                       # auto-find newest chkpt, speculative
+    python chat_sft.py --temperature 0.7     # sampled instead
+    python chat_sft.py --ckpt chkpt/quazimoto_sft.pt --system "You are concise."
+"""
+import argparse, os, sys
+# Force UTF-8 console so byte-merge tokens with non-cp1252 chars don't crash print().
+for _s in (sys.stdout, sys.stderr):
+    try:
+        _s.reconfigure(encoding="utf-8", errors="replace")
+    except (AttributeError, ValueError):
+        pass
+import torch
+
+from model import QuazimotoLM, QuazimotoConfig
+from generate import generate, generate_speculative, resolve_stop_ids
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+ROLE_TOK = {"system": "<|system|>", "user": "<|user|>", "assistant": "<|assistant|>"}
+
+
+def find_ckpt(path):
+    if path and os.path.isfile(path):
+        return path
+    folder = (os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")) or "."
+    if not os.path.isdir(folder):
+        return None
+    # prefer an SFT checkpoint, else newest *.pt
+    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
+    sft = [p for p in pts if "sft" in os.path.basename(p).lower()]
+    pool = sft or pts
+    return max(pool, key=os.path.getmtime) if pool else None
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+def render(history, tok):
+    """history: list of {role,content} -> token ids ending with the assistant header
+    so the model continues the assistant turn."""
+    parts = [f"<|im_start|>{ROLE_TOK.get(m['role'],'<|user|>')}\n{m['content']}<|im_end|>\n"
+             for m in history]
+    parts.append("<|im_start|><|assistant|>\n")
+    return tok.encode("".join(parts), add_special_tokens=False)
+
+
+def main():
+    p = argparse.ArgumentParser(description="Quazimoto-LM SFT chat (KV cache + speculative)")
+    p.add_argument("--ckpt", default="", help="SFT checkpoint (default: newest *sft*.pt in ./chkpt)")
+    p.add_argument("--tok_dir", default=PKG_DIR)
+    p.add_argument("--system", default="", help="optional system prompt")
+    p.add_argument("--max_new_tokens", type=int, default=256)
+    p.add_argument("--temperature", type=float, default=0.0, help="0 = speculative/greedy; >0 = sampled")
+    p.add_argument("--top_k", type=int, default=40)
+    p.add_argument("--top_p", type=float, default=0.95)
+    p.add_argument("--repetition_penalty", type=float, default=1.1)
+    p.add_argument("--no_speculative", action="store_true", help="force the cached sampler even at temp 0")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = p.parse_args()
+
+    path = find_ckpt(args.ckpt)
+    if path is None:
+        print("No checkpoint found; pass --ckpt or train one."); return
+    ckpt = torch.load(path, map_location=args.device, weights_only=False)
+    cfg = QuazimotoConfig(**ckpt["family_config"])
+    model = QuazimotoLM(cfg); model.load_state_dict(ckpt["model"], strict=False)
+    model.to(args.device).eval()
+    tok = load_tokenizer(args.tok_dir)
+    stop = resolve_stop_ids(tok, chat=True)          # <|im_end|>, <eos>, <|endoftext|>
+    imend = tok.get_vocab().get("<|im_end|>")
+
+    spec = (args.temperature <= 1e-4 and not args.no_speculative and model.mtp_heads is not None)
+    mode = "speculative (DSpark draft+verify)" if spec else f"sampled (temp {args.temperature}, KV cache)"
+    print(f"loaded {os.path.basename(path)} (step {ckpt.get('step')}) | decode: {mode}")
+    print("type your message; /reset clears history, /quit exits.\n")
+
+    history = []
+    if args.system:
+        history.append({"role": "system", "content": args.system})
+
+    while True:
+        try:
+            user = input("you> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print(); break
+        if not user:
+            continue
+        if user in ("/quit", "/exit"):
+            break
+        if user == "/reset":
+            history = ([{"role": "system", "content": args.system}] if args.system else [])
+            print("(history cleared)\n"); continue
+
+        history.append({"role": "user", "content": user})
+        ids = render(history, tok)
+        x = torch.tensor([ids], device=args.device)
+        with torch.no_grad():
+            if spec:
+                out = generate_speculative(model, cfg, x, args.max_new_tokens, stop_ids=stop)
+            else:
+                out = generate(model, cfg, x, args.max_new_tokens, temperature=max(args.temperature, 1e-6),
+                               top_k=args.top_k or None, top_p=args.top_p,
+                               repetition_penalty=args.repetition_penalty, stop_ids=stop,
+                               device=args.device, use_cache=True)
+        gen = out[0, len(ids):].tolist()
+        if imend in gen:                              # trim at the assistant turn end
+            gen = gen[:gen.index(imend)]
+        reply = tok.decode(gen, skip_special_tokens=True).strip()
+        try:
+            print(f"bot> {reply}\n")
+        except UnicodeEncodeError:
+            print("bot> " + reply.encode("ascii", "replace").decode() + "\n")
+        history.append({"role": "assistant", "content": reply})
+
+
+if __name__ == "__main__":
+    main()