fixes.

.gitignore

@@ -1,3 +1,4 @@
 ckpt/
 *.tar.gz
 *.swp
+pya0

README.md

@@ -26,8 +26,11 @@ Download your tokenizer, model checkpoints, and optionally the training logs (`e
 
 Optionally, test model using the MLM task:
 ```sh
-pip install pya0
-python test.py ./ckpt/to/tokenizer ./ckpt/to/model
+pip install pya0 # for math token preprocessing
+# testing local checkpoints:
+python test.py ./ckpt/math-tokenizer ./ckpt/2-2-0/encoder.ckpt
+# testing Model Hub checkpoints:
+python test.py approach0/coco-mae-220 approach0/coco-mae-220
 ```
 > **Note**  
 > Modify the test examples in `test.txt` to play with it.

test.py

@@ -43,6 +43,7 @@ def test(tokenizer_name_or_path, model_name_or_path, test_file='test.txt'):
             for pos in filter(lambda x: x!=0, maskpos):
                 tokens[pos-1] = '[MASK]'
             sentence = ' '.join(tokens)
+            sentence = sentence.replace('[mask]', '[MASK]')
             tokens = tokenizer(sentence,
                 padding=True, truncation=True, return_tensors="pt")
             #print(tokenizer.decode(tokens['input_ids'][0]))