#!/usr/bin/env python3
"""Smoke-test tokenizer/model loading and generation."""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from heapr.constants import DEFAULT_SMOKE_MODEL
from heapr.model_utils import load_causal_lm, load_tokenizer
from heapr.utils import require_torch


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", default=DEFAULT_SMOKE_MODEL)
    parser.add_argument("--revision")
    parser.add_argument("--prompt", default="Write a tiny Python function that adds two numbers.")
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--attn-implementation")
    parser.add_argument("--cache-implementation", default="static")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch = require_torch()
    tokenizer = load_tokenizer(args.model_id, revision=args.revision)
    model = load_causal_lm(
        args.model_id,
        revision=args.revision,
        dtype=args.dtype,
        attn_implementation=args.attn_implementation,
        use_cache=True,
        cache_implementation=args.cache_implementation,
    )
    inputs = tokenizer(args.prompt, return_tensors="pt")
    first_device = next(model.parameters()).device
    inputs = {key: value.to(first_device) for key, value in inputs.items()}
    start = time.time()
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=args.max_new_tokens)
    elapsed = time.time() - start
    print(tokenizer.decode(output_ids[0], skip_special_tokens=True))
    print(f"\n[smoke] elapsed={elapsed:.2f}s")


if __name__ == "__main__":
    main()