"""Quick smoke-test for the exported HuggingFace model."""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def main():
    model_dir = "."
    print("Loading model and tokenizer …")

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_dir, trust_remote_code=True, torch_dtype=torch.float32
    )
    if torch.cuda.is_available():
        model = model.to("cuda")

    prompts = [
        "एक समय की बात है",
        "एक जंगल में",
        "एक छोटी लड़की",
    ]

    for prompt in prompts:
        print(f"\n{'='*60}")
        print(f"Prompt: {prompt}")
        print(f"{'='*60}")
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            top_k=40,
            top_p=0.95,
            temperature=0.8,
        )
        print(tokenizer.decode(outputs[0], skip_special_tokens=True))


if __name__ == "__main__":
    main()