"""Interactive chat with the finetuned model. Model loads once, then you can ask questions back-to-back."""

import argparse
import torch
from pathlib import Path
from litgpt import Tokenizer
from litgpt.config import Config
from litgpt.model import GPT


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint_dir", type=str, default="Base/out/finetune/custom-100m-english-instruct/final")
    parser.add_argument("--tokenizer_dir", type=str, default="Base/checkpoints/EleutherAI/pythia-160m")
    parser.add_argument("--max_new_tokens", type=int, default=50)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--top_k", type=int, default=1)
    args = parser.parse_args()

    ckpt = Path(args.checkpoint_dir)
    print(f"Loading model from {ckpt}...")
    cfg = Config.from_checkpoint(ckpt)
    model = GPT(cfg)
    sd = torch.load(str(ckpt / "lit_model.pth"), map_location="cpu", weights_only=False)
    if "model" in sd:
        sd = sd["model"]
    model.load_state_dict(sd, strict=False)
    model = model.to("cuda").eval()
    tok = Tokenizer(Path(args.tokenizer_dir))
    print(f"Model loaded! ({sum(p.numel() for p in model.parameters()):,} params)")
    print(f"Settings: temp={args.temperature}, top_k={args.top_k}, max_tokens={args.max_new_tokens}")
    print("Type your question and press Enter. Type 'quit' or 'exit' to stop.\n")

    while True:
        try:
            question = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break
        if not question:
            continue
        if question.lower() in ("quit", "exit", "q"):
            print("Bye!")
            break

        alpaca = (
            "Below is an instruction that describes a task. "
            "Write a response that appropriately completes the request.\n\n"
            f"### Instruction:\n{question}\n\n### Response:\n"
        )
        ids = tok.encode(alpaca, device="cuda").unsqueeze(0)
        with torch.no_grad():
            for _ in range(args.max_new_tokens):
                logits = model(ids[:, -cfg.block_size:])
                logits = logits[:, -1, :] / args.temperature
                v, _ = torch.topk(logits, min(args.top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
                probs = torch.softmax(logits, dim=-1)
                nxt = torch.multinomial(probs, num_samples=1)
                ids = torch.cat([ids, nxt], dim=1)
                if nxt.item() == tok.eos_id:
                    break
        resp = tok.decode(ids[0]).split("### Response:")[-1].strip()
        print(f"\nLUNA: {resp}\n")


if __name__ == "__main__":
    main()