build(docker): download in safetensors format for supported models (#186)

peakji · web-flow · commit 997f19a3f266 · 2023-04-26T19:41:42.000+08:00
diff --git a/deployments/bundle/bloomz-560m.Dockerfile b/deployments/bundle/bloomz-560m.Dockerfile
@@ -4,7 +4,7 @@ FROM hyperonym/basaran:0.17.1
 WORKDIR /app
 
 # Download the model to be bundled
-RUN python utils/download.py bigscience/bloomz-560m /model
+RUN TENSOR_FORMAT=safetensors python utils/download.py bigscience/bloomz-560m /model
 
 # Provide default environment variables
 ENV MODEL="/model"
diff --git a/deployments/bundle/llama-7b.Dockerfile b/deployments/bundle/llama-7b.Dockerfile
@@ -4,7 +4,7 @@ FROM hyperonym/basaran:0.17.1
 WORKDIR /app
 
 # Download the model to be bundled
-RUN python utils/download.py huggyllama/llama-7b /model
+RUN TENSOR_FORMAT=safetensors python utils/download.py huggyllama/llama-7b /model
 
 # Provide default environment variables
 ENV MODEL="/model"
