Transformers custommodeling · unknown
publisher-documented, not executed source
from transformers import AutoModel, AutoTokenizer
import torch
repo_name = "nvidia/Nemotron-Labs-Diffusion-8B"
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModel.from_pretrained(repo_name, 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')
## Chat in AR Mode
out_ids, nfe = model.ar_generate(inputs.input_ids, max_new_tokens=512)
## Chat in dLM Mode
out_ids, nfe = model.generate(prompt_ids, max_new_tokens=512, block_length=32, threshold=0.9, eos_token_id=tokenizer.eos_token_id)
## Chat in Linear Self-Speculation Mode
out_ids, nfe = model.linear_spec_generate(prompt_ids, max_new_tokens=512, block_length=32, eos_token_id=tokenizer.eos_token_id)
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}]")Prerequisites: transformers>=5.0.0; custom code as shown in card.
- Tri-mode autoregressive/diffusion/block diffusion generation; compare speed only with same mode, steps, block size.
- Custom remote-code cache/denoising state is not generic KV.