Update README.md

README.md

@@ -64,17 +64,18 @@ model.to(device)
 ### Vocabulary Initialization
 
 ```python
-PAD_TOKEN = "[PAD]"
-MASK_TOKEN = "[MASK]"
+import json
+from huggingface_hub import hf_hub_download
 
-vocab = {PAD_TOKEN: 0, MASK_TOKEN: 1}
+vocab_file = hf_hub_download("yasserrmd/diffusion-text-demo", "vocab.json")
+with open(vocab_file) as f:
+    vocab = json.load(f)
 
-# Reverse mapping
-id_to_word = {i: w for w, i in vocab.items()}
+# Reverse mapping (IDs → tokens)
+id_to_word = {int(v): k for k, v in vocab.items()}
 
-# Special token IDs
-pad_id = vocab[PAD_TOKEN]
-mask_id = vocab[MASK_TOKEN]
+# Special IDs
+pad_id, mask_id = vocab["[PAD]"], vocab["[MASK]"]
 
 ```