Added Pytype None checks to run_inference.py

PiperOrigin-RevId: 679607401
tensorflow · Sep 27, 2024 · a3f25df · a3f25df
1 parent b430343
commit a3f25df
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 46 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+version: '3.1'
+
 # We mount the TFX BSL project root at /build (which is the WORKDIR of the
 # image) in the container.
 services:

diff --git a/tfx_bsl/beam/run_inference.py b/tfx_bsl/beam/run_inference.py
@@ -539,6 +539,10 @@ def _run_inference(self, examples: List[InputType],
                      model) -> List[Mapping[Text, Any]]:
     self._check_examples(examples)
     body = {'instances': self._make_instances(examples, serialized_examples)}
+    if self._api_client is None:
+      raise ValueError(
+          'API client is not initialized. Call load_model() first.'
+      )
     request = self._api_client.projects().predict(
         name=self._full_model_name, body=body)
     response = self._execute_request(request)