File size: 1,157 Bytes
13c5606 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | """Single-turn chat in Linear Self-Speculation mode. Mirrors the HF README snippet.
linear_spec_generate draws a diffusion draft under bidirectional attention,
then verifies it autoregressively, accepting the longest matching prefix plus
one bonus token per iteration. No LoRA — see chat_linear_spec_lora.py for the
LoRA-enhanced draft variant.
"""
import torch
from transformers import AutoModel, AutoTokenizer
REPO = "nvidia/Nemotron-Labs-Diffusion-8B"
tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
model = AutoModel.from_pretrained(REPO, trust_remote_code=True).cuda().to(torch.bfloat16)
user_input = input("User: ").strip()
history = [{"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("cuda")
out_ids, nfe = model.linear_spec_generate(
prompt_ids,
max_new_tokens=512,
block_length=32,
eos_token_id=tokenizer.eos_token_id,
)
reply = tokenizer.decode(out_ids[0, prompt_ids.shape[1]:], skip_special_tokens=True)
print(f"Model: {reply}")
print(f"[NFE={nfe}]")
|