Update README.md

README.md

@@ -67,6 +67,24 @@ with torch.no_grad():
 
     out = tokenizer.decode(out.cpu().numpy().tolist())
 
+    if "###" in out:
+        special_index = out.index("###")
+        out = out[: special_index]
+
+    if "[UNK]" in out:
+        special_index = out.index("[UNK]")
+        out = out[:special_index]
+        
+    if "</s>" in out:
+        special_index = out.index("</s>")
+        out = out[: special_index]
+
+    if len(out) > 0 and out[0] == " ":
+        out = out[1:]
+
+        convert_tokens = convert_tokens[1:]
+        probs = probs[1:]
+
     print(out)
 ```