#!/usr/bin/env python3
from __future__ import annotations

import argparse

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--prompt", action="append", default=None)
    parser.add_argument("--max-new-tokens", type=int, default=32)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--fix-mistral-regex", action="store_true")
    parser.add_argument("--print-token-details", action="store_true")
    parser.add_argument("--top-k-debug", type=int, default=0)
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_dir,
        trust_remote_code=True,
        fix_mistral_regex=args.fix_mistral_regex,
    )
    model = AutoModelForCausalLM.from_pretrained(
        args.model_dir,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    model.to(args.device)
    model.eval()
    prompts = args.prompt or ["The capital of France is", "中国的首都是"]
    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(args.device)
        with torch.inference_mode():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=args.max_new_tokens,
                do_sample=False,
                use_cache=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        print("PROMPT:", prompt)
        print(tokenizer.decode(output_ids[0], skip_special_tokens=False))
        if args.print_token_details:
            print("GENERATE_OUTPUT_TOKENS:", output_ids[0].tolist())
            ids = inputs["input_ids"].clone()
            generated = []
            logprobs = []
            pieces = []
            with torch.inference_mode():
                for _ in range(args.max_new_tokens):
                    logits = model(input_ids=ids).logits[:, -1, :]
                    logp = F.log_softmax(logits.float(), dim=-1)
                    next_id = torch.argmax(logits, dim=-1)
                    generated.append(int(next_id.item()))
                    logprobs.append(float(logp[0, next_id].item()))
                    pieces.append(tokenizer.decode([int(next_id.item())], skip_special_tokens=False))
                    ids = torch.cat([ids, next_id[:, None]], dim=1)
            print("INPUT_TOKENS:", inputs["input_ids"][0].tolist())
            print("GENERATED_TOKENS:", generated)
            print("GENERATED_PIECES:", pieces)
            print("GENERATED_LOGPROBS:", logprobs)
        if args.top_k_debug > 0:
            ids = inputs["input_ids"].clone()
            with torch.inference_mode():
                for step in range(args.max_new_tokens):
                    logits = model(input_ids=ids).logits[:, -1, :]
                    logp = F.log_softmax(logits.float(), dim=-1)
                    vals, toks = torch.topk(logp[0], k=args.top_k_debug)
                    print(f"TOPK_STEP {step} PREFIX_IDS {ids[0].tolist()}")
                    for token_id, value in zip(toks.tolist(), vals.tolist()):
                        print(token_id, repr(tokenizer.decode([token_id], skip_special_tokens=False)), value)
                    ids = torch.cat([ids, toks[:1][None, :]], dim=1)


if __name__ == "__main__":
    main()
