Cannot load Saved model including Tensorflow-text ops with `force_tf_compat_v1 = False`

[jeongukjae]

I'm using tfx Transform component with tensorflow-text op in dataflow, and facing this issue. (Env: tensorflow-text==2.8.2, tfx==1.7.1, tensorflow-transform==1.7.0)

    File "apache_beam/runners/common.py", line 1198, in apache_beam.runners.common.DoFnRunner.process 
    File "apache_beam/runners/common.py", line 718, in apache_beam.runners.common.PerWindowInvoker.invoke_process 
    File "apache_beam/runners/common.py", line 836, in apache_beam.runners.common.PerWindowInvoker._invoke_process_per_window 
    File "apache_beam/runners/common.py", line 1334, in apache_beam.runners.common._OutputProcessor.process_outputs 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow_transform/beam/impl.py", line 402, in process self._graph_state = self._shared_graph_state_handle.acquire( 
    File "/usr/local/lib/python3.8/dist-packages/apache_beam/utils/shared.py", line 305, in acquire return _shared_map.acquire(self._key, constructor_fn, tag) 
    File "/usr/local/lib/python3.8/dist-packages/apache_beam/utils/shared.py", line 246, in acquire result = control_block.acquire(constructor_fn, tag) 
    File "/usr/local/lib/python3.8/dist-packages/apache_beam/utils/shared.py", line 139, in acquire result = constructor_fn() 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow_transform/beam/impl.py", line 403, in <lambda> lambda: self._make_graph_state(saved_model_dir)) 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow_transform/beam/impl.py", line 369, in _make_graph_state result = self._GraphStateV2(saved_model_dir, 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow_transform/beam/impl.py", line 219, in __init__ saved_model_loader = saved_transform_io_v2.SavedModelLoader( 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow_transform/saved/saved_transform_io_v2.py", line 126, in __init__ imported = tf.compat.v2.saved_model.load(saved_model_dir) 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/saved_model/load.py", line 936, in load result = load_internal(export_dir, tags, options)["root"] 
    File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/saved_model/load.py", line 977, in load_internal raise FileNotFoundError( FileNotFoundError: Op type not registered 'NormalizeUTF8' in binary running on user-demographic-model-ta-06130957-qubq-harness-cp40. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed. You may be trying to load on a different device from the computational device. Consider setting the `experimental_io_device` option in `tf.saved_model.LoadOptions` to the io_device such as '/job:localhost'. 

I think the below code is the main reason for this issue.

transform/tensorflow_transform/saved/saved_transform_io_v2.py

Lines 112 to 154 in 74789ff

	class SavedModelLoader:
	"""Handles a SavedModel exported using TF 1.x APIs in TF 2.x."""
	def __init__(self, saved_model_dir: str):
	"""Init method for SavedModelLoader.
	Args:
	saved_model_dir: A SavedModel directory providing a transform graph. The
	MetaGraphDef and signature are selected from the SavedModel using keys
	defined in `../constants.py` ('transform' and 'transform_signature',
	respectively).
	"""
	# TODO(b/160294509): Stop using tf.compat.v2 when TF1.15 support is
	# dropped.
	imported = tf.compat.v2.saved_model.load(saved_model_dir)
	load_v2_in_compat = constants.TRANSFORM_SIGNATURE in imported.signatures
	if load_v2_in_compat:
	restored_function = imported.signatures[constants.TRANSFORM_SIGNATURE]
	wrapped, structured_inputs, structured_outputs = (
	_restore_from_v1_saved_model(restored_function, saved_model_dir))
	else:
	# transform_fn is now a ConcreteFunction, but was a tf.function. We need
	# to handle both to maintain backward compatiblity. If it's a tf.function,
	# since `input_signature` was specified when exporting the tf function to
	# `SavedModel`, there should be exactly one concrete function present on
	# loading the `SavedModel`.
	if hasattr(imported.transform_fn, 'concrete_functions'):
	concrete_functions = imported.transform_fn.concrete_functions
	assert len(concrete_functions) == 1, concrete_functions
	wrapped = concrete_functions[0]
	else:
	wrapped = imported.transform_fn
	func_graph = wrapped.graph
	structured_inputs = (
	tf2_utils.get_structured_inputs_from_func_graph(func_graph))
	structured_outputs = tf.nest.pack_sequence_as(
	func_graph.structured_outputs,
	func_graph.outputs,
	expand_composites=True)
	outputs_to_inputs_map = _get_output_to_inputs_map(structured_outputs)
	self._initialize(load_v2_in_compat, imported, wrapped, structured_inputs,
	structured_outputs, outputs_to_inputs_map)
	saved_transform_io._maybe_register_addon_ops() # pylint: disable=protected-access

As you can see from the code, tf-transform tries to register required ops in L154, but it also tries to load the model in L126 before the required ops are registered.