change the detokenization thread to return the actual eos token. (#108)

* change the detokenization thread to return the actual eos token.

Author: jwyang-google

jetstream/engine/mock_utils.py

@@ -63,7 +63,7 @@ def _encode(self, s: str) -> Sequence[int]:
 
   def _decode(self, ids: np.ndarray):
     """Converts a numpy array into a string."""
-    return "".join([chr(r) for r in list(ids)])
+    return "".join([chr(r) for r in list(ids) if r not in self.stop_tokens])
 
   def _encode_tf(self, s: str) -> np.ndarray:
     """Converts a string into a numpy array."""

jetstream/engine/token_utils.py

@@ -214,6 +214,7 @@ def process_result_tokens(
           )
         if tok_id in stop_tokens or not valid:
           complete[idx] = True
+          tok_id_so_far.append(tok_id)
           break
         else:
           if not is_client_side_tokenization: