diff --git a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
index 6bda0379bc7d..195ce693a388 100644
--- a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
+++ b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
@@ -226,7 +226,7 @@ jobs:
             -Prunner=DataflowRunner \
             -PpythonVersion=3.10 \
             -PloadTest.requirementsTxtFile=apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt \
-            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_11 }} --job_name=benchmark-tests-mltransform-generate-vocab-batch-${{env.NOW_UTC}}'
+            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_11 }} --job_name=benchmark-tests-mltransform-generate-vocab-batch-${{env.NOW_UTC}} --artifact_location=gs://temp-storage-for-perf-tests/mltransform/vocab_artifacts_${{env.NOW_UTC}} --output_vocab=gs://temp-storage-for-perf-tests/mltransform/vocab_outputs/mltransform_generate_vocab_batch_${{env.NOW_UTC}}'
       - name: run MLTransform One-Hot Encoding Batch
         uses: ./.github/actions/gradle-command-self-hosted-action
         timeout-minutes: 180
@@ -236,6 +236,5 @@ jobs:
             -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_one_hot_encoding_benchmark \
             -Prunner=DataflowRunner \
             -PpythonVersion=3.10 \
-            -PbeamPythonExtra=ml_test \
-            -PloadTest.requirementsTxtFile=apache_beam/ml/transforms/mltransform_tests_requirements.txt \
-            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_12 }} --autoscaling_algorithm=NONE --metrics_table=mltransform_one_hot_encoding_batch --influx_measurement=mltransform_one_hot_encoding_batch --job_name=benchmark-tests-mltransform-one-hot-encoding-batch-${{env.NOW_UTC}} --output_file=gs://temp-storage-for-end-to-end-tests/mltransform/one_hot_output_${{env.NOW_UTC}} --artifact_location=gs://temp-storage-for-end-to-end-tests/mltransform/artifacts_${{env.NOW_UTC}}'
\ No newline at end of file
+            -PloadTest.requirementsTxtFile=apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt \
+            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_12 }} --autoscaling_algorithm=NONE --metrics_table=mltransform_one_hot_encoding_batch --influx_measurement=mltransform_one_hot_encoding_batch --job_name=benchmark-tests-mltransform-one-hot-encoding-batch-${{env.NOW_UTC}} --output_file=gs://temp-storage-for-end-to-end-tests/mltransform/one_hot_output_${{env.NOW_UTC}} --artifact_location=gs://temp-storage-for-end-to-end-tests/mltransform/artifacts_${{env.NOW_UTC}}'
diff --git a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_One_Hot_Encoding_Batch.txt b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_One_Hot_Encoding_Batch.txt
index 27648d0c0fb0..993dd0820822 100644
--- a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_One_Hot_Encoding_Batch.txt
+++ b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_One_Hot_Encoding_Batch.txt
@@ -22,7 +22,7 @@
 --staging_location=gs://temp-storage-for-perf-tests/loadtests
 --temp_location=gs://temp-storage-for-perf-tests/loadtests
 --sdk_location=container
---requirements_file=apache_beam/ml/transforms/mltransform_tests_requirements.txt
+--requirements_file=apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt
 --publish_to_big_query=true
 --metrics_dataset=beam_run_inference
 --metrics_table=mltransform_one_hot_encoding_batch
diff --git a/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
index 4e00b3b4316c..0eb79d9480ec 100644
--- a/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
+++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
@@ -18,9 +18,9 @@
 # MLTransform TFT operations need a consistent TensorFlow Transform stack;
 # otherwise workers can crash-loop with pandas/numpy ABI mismatches.
 google-cloud-monitoring>=2.27.0
-tensorflow_transform>=1.14.0,<1.15.0
-tensorflow-metadata>=1.14.0,<1.15.0
-tfx-bsl>=1.14.0,<1.15.0
-numpy<2
-pandas<2
+tensorflow_transform==1.21.0
+tensorflow-metadata==1.21.0
+tfx-bsl==1.21.0
 dill
+numpy
+pandas
diff --git a/sdks/python/apache_beam/ml/transforms/mltransform_tests_requirements.txt b/sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt
similarity index 82%
rename from sdks/python/apache_beam/ml/transforms/mltransform_tests_requirements.txt
rename to sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt
index 9f37e070a606..1e64bf29333c 100644
--- a/sdks/python/apache_beam/ml/transforms/mltransform_tests_requirements.txt
+++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt
@@ -18,12 +18,12 @@
 # Requirements for MLTransform tests on Dataflow workers.
 # Keep this aligned with CloudML benchmark stack to avoid worker import errors.
 dill==0.4.1
-tfx_bsl==1.16.1
-tensorflow-transform==1.16.0
-tensorflow>=2.16,<2.17
-numpy>=1.22.0,<2.0
-tensorflow-metadata>=1.16.1,<1.17.0
-pyarrow>=10,<11
-tensorflow-serving-api>=2.16.1,<2.20
-tf-keras>=2.16.0,<2.17
+tfx_bsl==1.21.0
+tensorflow-transform==1.21.0
+tensorflow-metadata==1.21.0
+tensorflow
+numpy
+pyarrow
+tensorflow-serving-api
+tf-keras
 google-cloud-monitoring>=2.27.0