diff --git a/tensorflow_text/python/ops/sentencepiece_tokenizer.py b/tensorflow_text/python/ops/sentencepiece_tokenizer.py index f286abab5..4023d45b6 100644 --- a/tensorflow_text/python/ops/sentencepiece_tokenizer.py +++ b/tensorflow_text/python/ops/sentencepiece_tokenizer.py @@ -19,8 +19,6 @@ from __future__ import division from __future__ import print_function -import uuid - from tensorflow.python.eager import monitoring from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -45,24 +43,18 @@ class _SentencepieceModelResource(resource.TrackableResource): """Utility to track the model resource tensor (for SavedModel support).""" - def __init__(self, model, name, use_unique_shared_resource_name): + def __init__(self, model, name): super(_SentencepieceModelResource, self).__init__() self._model = model self._name = name - self._use_unique_shared_resource_name = use_unique_shared_resource_name _ = self.resource_handle # Accessing this property creates the resource. def _create_resource(self): model, name = self._model, self._name with ops.name_scope(name, "SentenceTokenizerInitializer", [model]): - # TODO(b/318839908): Switch to ref-counted resources and remove this. - shared_name = "" - if self._use_unique_shared_resource_name: - shared_name = str(uuid.uuid4()) - return gen_sentencepiece_tokenizer.sentencepiece_op( - model=model, - shared_name=shared_name, - ) + # TODO(b/318839908): Switch to using a ref-counted resource instead of + # this kernel-owned resource. + return gen_sentencepiece_tokenizer.sentencepiece_op(model=model) class SentencepieceTokenizer(TokenizerWithOffsets, Detokenizer): @@ -89,8 +81,7 @@ def __init__(self, add_bos=False, add_eos=False, return_nbest=False, - name=None, - use_unique_shared_resource_name=False): + name=None): """Creates & initializes a Sentencepiece processor. Args: @@ -115,10 +106,6 @@ def __init__(self, of a single one. The returned tensor has shape `[batch * nbest, (tokens)]`. name: The name argument that is passed to the op function. - use_unique_shared_resource_name: Whether to set a unique `shared_name` for - the Sentencepiece resource. This is useful for allowing the resource to - outlive the kernel that created it in eager mode. The resource's - lifetime will instead be tied to the eager context. Returns: pieces: A SentencepieceTokenizer. @@ -132,8 +119,7 @@ def __init__(self, self.add_bos = add_bos self.add_eos = add_eos self.return_nbest = return_nbest - self._model_resource = _SentencepieceModelResource( - model, name, use_unique_shared_resource_name) + self._model_resource = _SentencepieceModelResource(model, name) def tokenize(self, input, name=None): # pylint: disable=redefined-builtin """Tokenizes a tensor of UTF-8 strings. diff --git a/tensorflow_text/python/ops/sentencepiece_tokenizer_test.py b/tensorflow_text/python/ops/sentencepiece_tokenizer_test.py index 0651a4b57..53c506288 100644 --- a/tensorflow_text/python/ops/sentencepiece_tokenizer_test.py +++ b/tensorflow_text/python/ops/sentencepiece_tokenizer_test.py @@ -485,15 +485,11 @@ def testSavedModel(self): self.assertAllEqual(restored_model.tokenize(inputs), expected_result) file_io.delete_recursively(temp_dir) - @parameterized.parameters([True, False]) - def testBasicPipeline(self, use_unique_shared_resource_name): + def testBasicPipeline(self): if not context.executing_eagerly(): self.skipTest('testBasicPipeline only supported in eager mode.') - sp = SentencepieceTokenizer( - self.model, - use_unique_shared_resource_name=use_unique_shared_resource_name, - ) + sp = SentencepieceTokenizer(self.model) strings = ['hello', 'world'] dataset = dataset_ops.Dataset.from_tensor_slices(strings)