Update tpu single-host-inference demo deps (GoogleCloudPlatform#577)

* Update tpu single-host-inference demo deps * Update tpu single-host-inference tf demo deps * Simplify deps and path * Simplify deps for tf demo * maintain workable deps for jax bert demo
blackzlq · Apr 10, 2024 · 9bb746a · 9bb746a
1 parent 395b7fa
commit 9bb746a
Show file tree

Hide file tree

Showing 7 changed files with 5 additions and 398 deletions.
diff --git a/tutorials-and-examples/tpu-examples/single-host-inference/jax/bert/export_bert_model.py b/tutorials-and-examples/tpu-examples/single-host-inference/jax/bert/export_bert_model.py
@@ -14,7 +14,6 @@
 # ==============================================================================
 import logging
 import os
-import subprocess
 
 from jax.experimental import jax2tf
 import tensorflow as tf

diff --git a/tutorials-and-examples/tpu-examples/single-host-inference/jax/bert/install-bert.yaml b/tutorials-and-examples/tpu-examples/single-host-inference/jax/bert/install-bert.yaml
@@ -32,7 +32,7 @@ spec:
     - |
       cd /tmp/
       git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git
-      cp -r ai-on-gke/gke-tpu-examples/single-host-inference/jax .
+      cp -r ai-on-gke/tutorials-and-examples/tpu-examples/single-host-inference/jax .
       echo got jax bert resources
   - name: export-bert
     image: python:3.10
@@ -47,7 +47,8 @@ spec:
     - -c
     - |
       cd /tmp/
-      pip install -r ./jax/jax-requirements.txt
+      pip install -U "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+      pip install jaxlib==0.4.6 jax==0.4.6 flax==0.6.2 flatbuffers==23.3.3 tensorflow==2.11.1 tensorflow-serving-api==2.11.0 transformers==4.36.0
       cd jax/bert
       python3 export_bert_model.py
       echo exported bert model

diff --git a/tutorials-and-examples/tpu-examples/single-host-inference/jax/jax-requirements.txt b/tutorials-and-examples/tpu-examples/single-host-inference/jax/jax-requirements.txt
diff --git a/...-and-examples/tpu-examples/single-host-inference/tf/resnet50/export-model-requirements.in b/...-and-examples/tpu-examples/single-host-inference/tf/resnet50/export-model-requirements.in