File size: 4,815 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """Interactive chat with a Nemotron-Labs diffusion-LM, mirroring the snippets
in the HF model card (https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B).
Usage:
python chat.py --mode {ar|dlm|linear_spec|linear_spec_lora}
[--model nvidia/Nemotron-Labs-Diffusion-8B]
[--max-new-tokens 512]
[--block-length 32]
[--threshold 0.9]
The four modes mirror the model's four decoding paths:
ar -> model.ar_generate(...)
dlm -> model.generate(...) block-diffusion sampling
linear_spec -> model.linear_spec_generate(...) no LoRA
linear_spec_lora-> model.linear_spec_generate(...) with the bundled
`linear_spec_lora/` PEFT adapter attached as the draft
After load, type a message and press Enter. Type :q or Ctrl-C to quit. The
multi-turn chat history is fed back to the model via the chat template each
turn — same behaviour as the HF README snippet.
"""
import argparse
import sys
import torch
from transformers import AutoModel, AutoTokenizer
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--mode", required=True, choices=["ar", "dlm", "linear_spec", "linear_spec_lora"])
p.add_argument("--model", default="nvidia/Nemotron-Labs-Diffusion-8B")
p.add_argument("--max-new-tokens", type=int, default=512)
p.add_argument("--block-length", type=int, default=32, help="dLM / linear_spec block size")
p.add_argument("--threshold", type=float, default=0.9, help="dLM confidence threshold")
p.add_argument("--lora-subfolder", default="linear_spec_lora",
help="subfolder under --model that holds the PEFT adapter (linear_spec_lora mode)")
p.add_argument("--device", default="cuda")
p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
return p.parse_args()
def load_model_and_tokenizer(args: argparse.Namespace):
dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
model = AutoModel.from_pretrained(args.model, trust_remote_code=True)
model = model.to(args.device).to(dtype)
# linear_spec_lora: attach the bundled draft-side LoRA adapter and unwrap
# the PeftModel so we can call linear_spec_generate directly (the method
# toggles the adapter internally between the diffusion draft and the
# causal verify pass).
if args.mode == "linear_spec_lora":
from peft import PeftModel
model = PeftModel.from_pretrained(model, args.model, subfolder=args.lora_subfolder).eval()
model = model.model
return model, tokenizer
def generate(args, model, tokenizer, prompt_ids: torch.Tensor):
"""Dispatch to the right per-mode call. Returns (out_ids, nfe)."""
eos = tokenizer.eos_token_id
if args.mode == "ar":
return model.ar_generate(prompt_ids, max_new_tokens=args.max_new_tokens)
if args.mode == "dlm":
return model.generate(
prompt_ids,
max_new_tokens=args.max_new_tokens,
block_length=args.block_length,
threshold=args.threshold,
eos_token_id=eos,
)
# linear_spec / linear_spec_lora share the call path; the LoRA attach has
# already happened in load_model_and_tokenizer if needed.
return model.linear_spec_generate(
prompt_ids,
max_new_tokens=args.max_new_tokens,
block_length=args.block_length,
eos_token_id=eos,
)
def main():
args = parse_args()
print(f"Loading {args.model} ({args.mode})...", file=sys.stderr)
model, tokenizer = load_model_and_tokenizer(args)
print(f"Ready. Type :q to quit.\n", file=sys.stderr)
history = []
while True:
try:
user_input = input("User: ").strip()
except (EOFError, KeyboardInterrupt):
print()
break
if not user_input or user_input in {":q", ":quit", "exit", "quit"}:
break
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(model.device)
out_ids, nfe = generate(args, model, tokenizer, prompt_ids)
new_tokens = out_ids[0, prompt_ids.shape[1]:]
reply = tokenizer.decode(new_tokens, skip_special_tokens=True)
history.append({"role": "assistant", "content": reply})
print(f"Model [{args.mode}, NFE={nfe}]: {reply}\n")
if __name__ == "__main__":
main()
|