Fixes python3 related unicode errors in wiki_revision_utils.

T2T Team · copybara-github · commit 0b3ee79c460a · 2021-10-21T13:41:17.000-07:00
PiperOrigin-RevId: 404870305
diff --git a/tensor2tensor/data_generators/wiki_revision_utils.py b/tensor2tensor/data_generators/wiki_revision_utils.py
@@ -33,6 +33,10 @@
 import tensorflow.compat.v1 as tf
 
 
+def to_unicode(s):
+  return s.decode("utf-8")
+
+
 def include_revision(revision_num, skip_factor=1.1):
   """Decide whether to include a revision.
 
@@ -76,6 +80,10 @@ def file_page_generator(my_file, max_page_size=2**28):
   leftovers = ""
   while True:
     chunk = my_file.read(chunk_size)
+    try:
+      chunk = to_unicode(chunk)
+    except UnicodeDecodeError:
+      chunk = ""
     if not chunk:
       break
     chunk = leftovers + chunk
@@ -112,7 +120,7 @@ def get_title(page):
   assert start_pos != -1
   assert end_pos != -1
   start_pos += len("<title>")
-  return text_encoder.to_unicode_utf8(page[start_pos:end_pos])
+  return page[start_pos:end_pos]
 
 
 def get_id(page):
@@ -251,7 +259,6 @@ def get_text(revision, strip=True):
     ret = revision[end_tag_pos:end_pos]
   if strip:
     ret = strip_text(ret)
-  ret = text_encoder.to_unicode_utf8(ret)
   return ret
 
 
