feat: add chat_loop.py

chat_loop.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""
+MicroLLM2 Interactive Chat Loop
+- Loads MLVXN/MicroLLM2 (or local ./microllm2-checkpoints/final_merged)
+- ChatML: <|im_start|>user / assistant
+- Works on H100 (bf16) and local CPU
+- Run: python chat_loop.py [--local] [--temp 0.7]
+
+No token hardcoded — uses HF_TOKEN env if private, else public pull.
+"""
+import os, sys, torch
+from pathlib import Path
+
+# Use local checkpoint if available (faster on H100), else HF
+LOCAL = Path("/home/zeus/microllm2/microllm2-checkpoints/final_merged")
+HF_ID = "MLVXN/MicroLLM2"
+MODEL_ID = str(LOCAL) if LOCAL.exists() else HF_ID
+
+# Allow override
+if "--local" in sys.argv and LOCAL.exists():
+    MODEL_ID = str(LOCAL)
+elif "--hf" in sys.argv:
+    MODEL_ID = HF_ID
+
+print(f"[*] Loading MicroLLM2 from {MODEL_ID} ...")
+try:
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+except ImportError:
+    print("pip install transformers accelerate torch"); sys.exit(1)
+
+tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=False)
+if tok.pad_token is None:
+    tok.pad_token = tok.eos_token
+# Ensure ChatML tokens exist
+if "<|im_start|>" not in tok.get_vocab():
+    tok.add_special_tokens({"additional_special_tokens": ["<|im_start|>", "<|im_end|>"]})
+
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+device_map = "auto" if torch.cuda.is_available() else None
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, torch_dtype=dtype, device_map=device_map,
+        trust_remote_code=False, attn_implementation="sdpa"
+    )
+except Exception as e:
+    print(f"[!] sdpa load failed {e}, retry without attn arg")
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=dtype, device_map=device_map)
+
+model.eval()
+device = next(model.parameters()).device
+print(f"[+] Loaded on {device} ({dtype}) — {model.num_parameters()/1e9:.2f}B params")
+print(f"[+] MicroLLM2 by Maximalist Labs — type 'exit' to quit, 'clear' to reset history\n")
+
+# Chat history as list of dicts for ChatML
+history = []
+
+def format_prompt(history, user_msg):
+    # Build ChatML prompt
+    msgs = history + [{"role": "user", "content": user_msg}]
+    parts = []
+    for m in msgs:
+        parts.append(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>")
+    parts.append("<|im_start|>assistant\n")
+    return "\n".join(parts)
+
+# Generation defaults — tuned for GPT2-XL 1.5B chat
+temp = 0.7
+top_p = 0.9
+max_new = 120
+if "--temp" in sys.argv:
+    try: temp = float(sys.argv[sys.argv.index("--temp")+1])
+    except: pass
+
+while True:
+    try:
+        user = input("\nYou: ").strip()
+    except (EOFError, KeyboardInterrupt):
+        print("\nbye"); break
+    if not user:
+        continue
+    if user.lower() in ("exit","quit","q"):
+        break
+    if user.lower() in ("clear","reset","new"):
+        history = []; print("[*] history cleared"); continue
+
+    prompt = format_prompt(history, user)
+    inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=900).to(device)
+
+    # Warn if truncated (1024 limit)
+    if inputs.input_ids.shape[1] >= 900:
+        print("[!] near 1024 ctx — consider 'clear'")
+
+    with torch.no_grad():
+        out = model.generate(
+            **inputs, max_new_tokens=max_new, do_sample=(temp>0),
+            temperature=temp if temp>0 else 1.0, top_p=top_p,
+            repetition_penalty=1.1, pad_token_id=tok.eos_token_id,
+            eos_token_id=tok.convert_tokens_to_ids("<|im_end|>") if "<|im_end|>" in tok.get_vocab() else tok.eos_token_id,
+        )
+    # Decode only new tokens
+    gen = out[0][inputs.input_ids.shape[1]:]
+    text = tok.decode(gen, skip_special_tokens=False)
+    # Strip ChatML tail
+    if "<|im_end|>" in text:
+        text = text.split("<|im_end|>")[0]
+    text = text.replace("<|endoftext|>", "").strip()
+    print(f"\nMicroLLM2: {text}")
+
+    # Keep history (trim to last 6 turns to stay <1024)
+    history.append({"role": "user", "content": user})
+    history.append({"role": "assistant", "content": text})
+    if len(history) > 12:
+        history = history[-12:]
+
+print("done")