Fix whitespace replacement bug (#24)

Author: FanhaiLu1

jetstream/engine/token_utils.py

@@ -36,7 +36,8 @@ def mix_decode(vocab: Vocabulary, tok_id: int):
       utilize IdToPiece to convert it into a string, likely with a space placeholder (' ') for the corresponding tokens.
     """
     p_token = vocab.tokenizer.IdToPiece(tok_id)
-    p_token = p_token.replace('▁', ' ').replace('_', ' ')
+    # SentencePiece escapes the whitespace with a meta symbol "▁" (U+2581)
+    p_token = p_token.replace('▁', ' ')
     d_token = vocab.tokenizer.decode([tok_id])
     return p_token if p_token.lstrip() == d_token else d_token 
 

jetstream/tests/engine/test_token_utils.py

@@ -28,7 +28,7 @@ def __init__(self, tokenizer_path: str):
 
    def decode(self, t: int) -> str:
     token = self.vocab.tokenizer.IdToPiece(t)
-    token = token.replace('▁', ' ').replace('_', ' ')
+    token = token.replace('▁', ' ')
     return token      
 
 
@@ -77,6 +77,14 @@ def test_sp_vs_seqio(self):
           seqio_t = self.jt_tokenizer.vocab.tokenizer.decode([n])
           self.assertEqual(sp_t, seqio_t)
 
+    def test_underscore_in_output(self):
+       self.setup()
+       n = 21326
+       mix_output = token_utils.mix_decode(vocab = self.jt_tokenizer.vocab, tok_id = n)
+       decode_output = self.sp_tokenizer.decode([n])  
+       self.assertEqual(mix_output, " `__")  
+       self.assertEqual(mix_output.lstrip(), decode_output)
+
 
 if __name__ == '__main__':
    unittest.main()