Updating dependencies as suggested

Author: tyagiva

dependencies/requirements/base_requirements/requirements.txt

@@ -2,7 +2,8 @@ absl-py
 aqtp
 array-record
 cloud-accelerator-diagnostics
-cloud-tpu-diagnostics
+# TODO(b/477743392): Remove the pinned versions once the yanked version detection problem is resolved
+cloud-tpu-diagnostics==0.1.5
 datasets
 flax
 gcsfs
@@ -12,6 +13,7 @@ google-cloud-mldiagnostics
 google-cloud-monitoring
 grain[parquet]
 huggingface_hub
+hf-xet==1.2.0
 jax
 jaxlib
 jaxtyping
@@ -38,7 +40,7 @@ tensorflow-datasets
 tensorflow-text
 tensorflow
 tiktoken
-tokamax
+tokamax==0.0.8
 transformers
 qwix
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip

dependencies/requirements/generated_requirements/tpu-requirements.txt

@@ -83,7 +83,7 @@ gviz-api>=1.10.0
 h11>=0.16.0
 h5py>=3.15.1
 hf-transfer>=0.1.9
-hf-xet>=1.2.1 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
+hf-xet>=1.2.0 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
 httpcore>=1.0.9
 httplib2>=0.31.1
 httpx>=0.28.1
@@ -217,7 +217,7 @@ tensorflow>=2.19.1
 tensorstore>=0.1.79
 termcolor>=3.2.0
 tiktoken>=0.12.0
-tokamax>=0.1.0
+tokamax>=0.0.8
 tokenizers>=0.22.2
 toml>=0.10.2
 tomlkit>=0.14.0