| """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) |
|
|
| |
| |
| |
| |
| 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, |
| ) |
| |
| |
| 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() |
|
|
