tensorflow_text operations not registered in BulkInferrer component

System information

• Have I specified the code to reproduce the issue (Yes, No): No
• Environment in which the code is executed (e.g., Local(Linux/macOS/Windows), Interactive Notebook, Google Cloud, etc): Kubeflow
• TensorFlow version: 2.4.1
• TFX Version: 1.0.0rc0
• Python version: 3.7.1

I’m trying to use the BulkInferrer component in my TFX pipeline. My model has a BERT component and this relies on tensorflow-text. Components that require this module to run i.e. the Transform and Trainer components complete successfully. The exported model is passed directly to the BulkInferrer component at which point the following error is received, Op type not registered 'RegexSplitWithOffsets' in binary . This is an operation that is imported with tensorflow-text so I feel like something is going wrong here. tensorflow-text is installed in the image using the following dockerfile:

FROM tensorflow/tfx:1.0.0rc0
WORKDIR /pipeline
COPY ./ ./
ENV PYTHONPATH="/pipeline:${PYTHONPATH}"
RUN /usr/bin/python3 -m pip install --upgrade pip
RUN pip install -r requirements.txt
RUN pip install tensorflow-text --no-dependencies
RUN python -c 'import tensorflow_text # Registers the ops.'

I see in the executor of BulkInferrer there’s this line here, this import tensorflow_text as _ should be all that is required to register the operations with tensorflow. I’ve looked at the logs of my component and I don’t see any exception raised signaling that this failed.

I’m fairly sure this isn’t a TFX version issue as I’ve experienced the same error using TFX 0.27.0.

Could anyone from the TFX team please advise where I can look to try and figure out this issue? At this point, I’m stumped, since this has worked throughout the rest of my pipeline.

[jiyongjung0]

tensorflow_text support in tensorflow/serving is a different story and it seems like supported. Please see https://github.com/tensorflow/serving/issues/1490

The problem in BulkInferrer is that ‘tensorflow_text’ package is not loaded well in the bulk inferrer workers although it is imported in the BulkInferrer executor.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No