Have Tokenizers extend tf.Module

This fixes #224 where tokenizers were unable to be saved from within custom Keras layers. PiperOrigin-RevId: 315784812
tensorflow · Jun 16, 2020 · 97c922a · 97c922a
1 parent b72c60d
commit 97c922a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tensorflow_text/BUILD b/tensorflow_text/BUILD
@@ -498,6 +498,9 @@ py_test(
 py_library(
     name = "tokenization",
     srcs = ["python/ops/tokenization.py"],
+    deps = [
+        # python/module tensorflow dep,
+    ],
 )
 
 py_tf_text_library(

diff --git a/tensorflow_text/python/ops/tokenization.py b/tensorflow_text/python/ops/tokenization.py
@@ -21,9 +21,10 @@
 
 import abc
 
+from tensorflow.python.module import module
 
-# TODO(broken): Have this extend Module when it becomes public
-class Tokenizer(object):
+
+class Tokenizer(module.Module):
   """Base class for tokenizer implementations."""
 
   __metaclass__ = abc.ABCMeta