#!/usr/bin/env python3
"""Minimal text-generation example for Qwick-3.5-9B."""

from __future__ import annotations

import argparse

import torch
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="horiuchinobuyuki/Qwick-3.5-9B")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--max-new-tokens", type=int, default=8192)
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = Qwen3_5ForConditionalGeneration.from_pretrained(
        args.model, dtype=torch.bfloat16, device_map="auto"
    ).eval()
    inputs = tokenizer.apply_chat_template(
        [{"role": "user", "content": args.prompt}],
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=True,
        return_tensors="pt",
        return_dict=True,
    ).to(next(model.parameters()).device)
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=True,
            temperature=1.0,
            top_p=0.95,
            top_k=20,
            min_p=0.0,
            repetition_penalty=1.0,
        )
    completion = output[0, inputs.input_ids.shape[-1]:]
    print(tokenizer.decode(completion, skip_special_tokens=True))


if __name__ == "__main__":
    main()