#!/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")