#!/usr/bin/env python3
"""
Chat with this model on your own machine.

    pip install -r requirements.txt
    python run_model.py                            # multi-turn chat
    python run_model.py "When did WWII end?"       # one answer, then exit
    python run_model.py --model <dir or repo id>   # a model somewhere else

By default the model is loaded from the directory this script sits in, so a
plain download of the repo is enough:

    hf download <repo-id> --local-dir model && cd model && python run_model.py

Runs on CPU (no GPU, no flash-attn needed); uses CUDA automatically when
available. Greedy decoding (temperature 0) is the default and gives this
model's most reliable answers.

In-chat commands:
    /reset       start a new conversation (forget the history)
    /temp N      set the sampling temperature (0 = greedy)
    /length N    set the maximum tokens per reply
    /help        show the commands
    /quit        exit (Ctrl-C or Ctrl-D also work)
"""

import argparse
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))

COMMANDS = """Commands:
  /reset      start a new conversation
  /temp N     sampling temperature (0 = greedy, the default)
  /length N   maximum tokens per reply
  /help       show this list
  /quit       exit"""


def resolve_model(model):
    """A local directory must hold a config.json; anything that looks like
    `user/name` is treated as a Hub repo id (downloaded on first use)."""
    if os.path.isdir(model):
        if os.path.exists(os.path.join(model, "config.json")):
            return model
        sys.exit(f"No model at {model} (no config.json there).\n"
                 "Run this script from a downloaded model directory, or pass "
                 "--model <directory or repo id>.")
    if "/" in model and not model.startswith((".", "/", "~")):
        return model
    sys.exit(f"No model at {model} (not a directory, not a repo id).")


def main():
    p = argparse.ArgumentParser(
        description="Chat with this model locally.",
        epilog=COMMANDS, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("prompt", nargs="*", help="A single question (omit for interactive chat).")
    p.add_argument("--model", default=HERE,
                   help="Model directory or Hub repo id (default: this script's directory).")
    p.add_argument("--temperature", type=float, default=0.0,
                   help="0 = greedy (recommended). Higher adds errors and drift.")
    p.add_argument("--max-new-tokens", type=int, default=2048,
                   help="Maximum tokens per reply (default: 2048; generation stops earlier "
                        "when the answer ends).")
    p.add_argument("--fp32", action="store_true",
                   help="Run the model in float32 on CUDA instead of the default bf16 "
                        "(CPU is always float32).")
    args = p.parse_args()
    model_path = resolve_model(args.model)

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    device = "cuda" if torch.cuda.is_available() else "cpu"
    # bf16 on CUDA unless --fp32; CPU is always float32.
    dtype = torch.bfloat16 if (device == "cuda" and not args.fp32) else torch.float32
    print(f"Loading {model_path} ({device}, {str(dtype).split('.')[-1]})...")
    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, trust_remote_code=True, dtype=dtype).to(device).eval()

    temperature = args.temperature
    max_new_tokens = args.max_new_tokens

    def reply(messages):
        ids = torch.tensor(tok.build_chatml_ids(messages, add_generation_prompt=True))
        out = model.chat_generate(ids, max_new_tokens=max_new_tokens, temperature=temperature)
        return tok.decode_response(out)

    if args.prompt:
        print(reply([{"role": "user", "content": " ".join(args.prompt)}]))
        return

    print("Ready. Type a message; /help lists the commands.\n")
    messages = []  # the whole conversation is re-sent each turn
    while True:
        try:
            text = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return
        if not text:
            continue
        if text.startswith("/"):
            cmd, _, arg = text.partition(" ")
            cmd = cmd.lower()
            if cmd == "/quit":
                return
            elif cmd == "/reset":
                messages = []
                print("New conversation.")
            elif cmd == "/temp":
                try:
                    temperature = float(arg)
                    print(f"Temperature {temperature:g}" + (" (greedy)" if temperature < 0.01 else ""))
                except ValueError:
                    print("Usage: /temp N   (e.g. /temp 0 or /temp 0.7)")
            elif cmd == "/length":
                try:
                    max_new_tokens = int(arg)
                    print(f"Max tokens per reply {max_new_tokens}")
                except ValueError:
                    print("Usage: /length N   (e.g. /length 512)")
            elif cmd == "/help":
                print(COMMANDS)
            else:
                print(f"Unknown command {cmd}.\n{COMMANDS}")
            continue
        messages.append({"role": "user", "content": text})
        try:
            answer = reply(messages)
        except KeyboardInterrupt:
            messages.pop()
            print("\n(interrupted)\n")
            continue
        messages.append({"role": "assistant", "content": answer})
        print(f"\nModel: {answer}\n")


if __name__ == "__main__":
    main()