Update app.py

app.py

@@ -11,17 +11,14 @@ model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-PAD_TOKEN = "[PAD]"
-MASK_TOKEN = "[MASK]"
+with open("vocab.json") as f:
+    vocab = json.load(f)
 
-vocab = {PAD_TOKEN: 0, MASK_TOKEN: 1}
+# Reverse mapping (IDs → tokens)
+id_to_word = {int(v): k for k, v in vocab.items()}
 
-# Reverse mapping
-id_to_word = {i: w for w, i 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]"]