update images (#113)

Author: morgandu

jetstream/tools/proxy_dev/base.Dockerfile

@@ -18,8 +18,20 @@ COPY  maxtext ./maxtext
 RUN cd maxtext/ && \
 pip install -r requirements.txt
 
-RUN pip install setuptools==58 fastapi==0.103.2 uvicorn nltk evaluate
+RUN pip install setuptools==58 fastapi==0.103.2 uvicorn
 
 RUN pip install ./JetStream
 
+COPY inference_mlperf4.1 ./inference_mlperf4.1
+RUN apt-get -y install python3-dev && apt-get -y install build-essential
+RUN pip install ./inference_mlperf4.1/loadgen
+RUN pip install \
+    transformers==4.31.0 \
+    nltk==3.8.1 \
+    evaluate==0.4.0 \
+    absl-py==1.4.0 \
+    rouge-score==0.1.2 \
+    sentencepiece==0.1.99 \
+    accelerate==0.21.0
+
 ENTRYPOINT ["bash"]

jetstream/tools/proxy_dev/dev.Dockerfile

@@ -13,17 +13,6 @@ COPY  JetStream ./JetStream
 COPY  maxtext ./maxtext
 
 RUN pip install ./JetStream
-
-COPY inference_mlperf4.1 ./inference_mlperf4.1
-RUN apt-get -y install python3-dev && apt-get -y install build-essential
-RUN pip install ./inference_mlperf4.1/loadgen
-RUN pip install \
-    transformers==4.31.0 \
-    nltk==3.8.1 \
-    evaluate==0.4.0 \
-    absl-py==1.4.0 \
-    rouge-score==0.1.2 \
-    sentencepiece==0.1.99 \
-    accelerate==0.21.0
+RUN pip install -r ./maxtext/requirements.txt
 
 ENTRYPOINT ["bash"]