YongganFu commited on
Commit
ad6359b
·
verified ·
1 Parent(s): 7542260

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +35 -2
README.md CHANGED
@@ -1,4 +1,37 @@
1
  ---
2
- {}
 
3
  ---
4
- Coming soon.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ library_name: transformers
3
+ tags: []
4
  ---
5
+
6
+ # Nemotron-Diffusion-Research-4B-v0
7
+
8
+ Developed by [DLER team](https://nv-dler.github.io/) @ NVR and will be updated actively. Contact Yonggan Fu, Lexington Whalen, and Pavlo Molchanov for any question.
9
+
10
+
11
+ # Environment
12
+
13
+ Docker path: `/lustre/fsw/portfolios/nvr/users/yongganf/docker/megatron_py25_dllm.sqsh` on OCI-ORD/OCI-NRT or `/lustre/fsw/nvr_lpr_llm/yongganf/docker/megatron_py25_dllm.sqsh` on EOS.
14
+
15
+
16
+ ## Chat with Our Model
17
+
18
+
19
+ ```
20
+ from transformers import AutoModel, AutoTokenizer
21
+ import torch
22
+
23
+ repo_name = "nvidia/Nemotron-Diffusion-Research-4B-v0"
24
+
25
+ tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
26
+ model = AutoModel.from_pretrained(repo_name, trust_remote_code=True)
27
+ model = model.cuda().to(torch.bfloat16)
28
+
29
+ user_input = input("User: ").strip()
30
+
31
+ prompt_ids = tokenizer(user_input,return_tensors='pt').input_ids.to(device='cuda')
32
+ out_ids, nfe = model.generate(prompt_ids, max_new_tokens=128, steps=128, block_length=32, threshold=0.9)
33
+
34
+ tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
35
+ print(f"Model: {tokenized_out}")
36
+ print(f"[Num Function Eval (NFE)={nfe}]")
37
+ ```