from pathlib import Path

import torch
from transformers import Gemma3ForCausalLM, PreTrainedTokenizerFast


MODEL_DIR = Path(__file__).resolve().parent / "hf"
PROMPT = "Once upon"


def main() -> None:
    tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_DIR)
    model = Gemma3ForCausalLM.from_pretrained(
        MODEL_DIR,
        dtype=torch.float32,
    ).eval()

    input_ids = torch.tensor(
        [
            [tokenizer.bos_token_id]
            + tokenizer.encode(PROMPT, add_special_tokens=False)
        ]
    )

    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=100,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    print(tokenizer.decode(output[0], skip_special_tokens=True))


if __name__ == "__main__":
    main()