from transformers import AutoModel, AutoTokenizer, AutoConfig
import torch

repo_name = "src/models/Nemotron-Diffusion-Exp-Ministral-8B-Instruct-final-template-merged"

tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
config = AutoConfig.from_pretrained(repo_name, trust_remote_code=True)
config.dlm_paradigm = "bidirectional"
model = AutoModel.from_pretrained(repo_name, config=config, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)

history = []

user_input = input("User: ").strip()
history.append({"role": "user", "content": user_input})

prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
prompt_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device='cuda')
out_ids, nfe = model.generate(prompt_ids, max_new_tokens=128, steps=128, block_length=32, shift_logits=False, causal_context=True, threshold=0.9)

tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
print(f"Model: {tokenized_out}")
print(f"[Num Function Eval (NFE)={nfe}]")
