"""Interactive chat with the fine-tuned Gemma4-E2B reasoning model (base +
LoRA adapter) — streaming output.

Repo: xbruce22/gemma-4-e2b-reasoning-lora
Base: unsloth/gemma-4-E2B-it

Auto-detects the accelerator (CUDA / Intel XPU / CPU). Loads the base model,
applies this LoRA adapter, merges it for fast inference, and runs a
multi-turn chat using the Gemma4 chat template with thinking ON — the model
emits a <|channel>thought ... <channel|> reasoning block (concise bullets,
as it was trained) before the final answer.

Install:
  pip install torch transformers peft

Run:
  python chat.py
  python chat.py --repo xbruce22/gemma-4-e2b-reasoning-lora
  python chat.py --device cpu            # force CPU

In-chat commands:
  /q       quit          /reset   clear history
  /raw     toggle raw output (show <|channel>/<channel|>/<turn|> markers)
  /think   toggle thinking on/off (default ON)
"""
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoProcessor, TextStreamer
from peft import PeftModel

DEFAULT_REPO = "xbruce22/gemma-4-e2b-reasoning-lora"
BASE_MODEL = "unsloth/gemma-4-E2B-it"

# Build special-token strings from chr() so this source file never contains
# literal angle-bracket markers (avoids editor/toolchain mangling).
CHAN_OPEN = chr(60) + "|channel>thought" + chr(10)
CHAN_CLOSE = chr(60) + "channel|" + chr(62)
TURN_END = chr(60) + "turn|" + chr(62)
THINK = chr(60) + "|think|" + chr(62)


def pick_device():
    if torch.cuda.is_available():
        return "cuda", torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        return "xpu", torch.bfloat16
    return "cpu", torch.float32


def clean_display(text):
    if CHAN_OPEN in text and CHAN_CLOSE in text:
        _, _, rest = text.partition(CHAN_OPEN)
        thought, _, answer = rest.partition(CHAN_CLOSE)
        return ("\n── thinking ──\n" + thought.strip() +
                "\n── answer ──\n" + answer.strip())
    for m in (TURN_END, THINK):
        text = text.replace(m, "")
    return text.strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=DEFAULT_REPO,
                    help="HF repo id of the LoRA adapter")
    ap.add_argument("--device", default=None,
                    help="force device: cuda | xpu | cpu")
    args = ap.parse_args()

    device, dtype = pick_device() if args.device is None else (args.device, torch.float32)
    print(f"device={device}  dtype={dtype}")

    print("Loading processor...")
    processor = AutoProcessor.from_pretrained(BASE_MODEL)
    tokenizer = processor.tokenizer
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Loading base model {BASE_MODEL} ...")
    base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, dtype=dtype)
    base = base.to(device)
    base.config.use_cache = True
    base.eval()

    print(f"Applying + merging LoRA adapter {args.repo} ...")
    model = PeftModel.from_pretrained(base, args.repo)
    model = model.merge_and_unload()
    model.eval()
    print("Ready.\n")

    show_raw = False
    thinking = True
    messages = [{"role": "system", "content": "You are a helpful, concise assistant."}]

    print(f"Chat ready. /q quit · /reset · /raw · /think   "
          f"(thinking={'ON' if thinking else 'OFF'})\n")

    while True:
        try:
            user = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nbye."); break
        if not user:
            continue
        if user == "/q":
            print("bye."); break
        if user == "/reset":
            messages = [{"role": "system", "content": "You are a helpful, concise assistant."}]
            print("(reset)\n"); continue
        if user == "/raw":
            show_raw = not show_raw
            print(f"(display={'raw' if show_raw else 'clean'})\n"); continue
        if user == "/think":
            thinking = not thinking
            print(f"(thinking={'ON' if thinking else 'OFF'})\n"); continue

        messages.append({"role": "user", "content": user})
        try:
            text = processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True,
                enable_thinking=thinking)
        except TypeError:
            text = processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True)

        inputs = processor(text=[text], return_tensors="pt").to(device)
        for k in list(inputs.keys()):
            if "token_type" in k or "pixel" in k or "audio" in k:
                inputs.pop(k)

        print("model> ", end="", flush=True)
        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        with torch.inference_mode():
            out_ids = model.generate(
                **inputs, max_new_tokens=2048, do_sample=True,
                temperature=1.0, top_p=0.95, top_k=64,
                pad_token_id=tokenizer.pad_token_id, streamer=streamer)
        gen_ids = out_ids[0][inputs["input_ids"].shape[1]:]
        gen_text = tokenizer.decode(gen_ids, skip_special_tokens=False)
        messages.append({"role": "assistant",
                         "content": tokenizer.decode(gen_ids, skip_special_tokens=True)})
        print()
        if show_raw:
            print("--- raw ---"); print(gen_text); print("--- end raw ---")
        print()


if __name__ == "__main__":
    main()