feat: adding support for hf_pretrained_model option

containers · Jun 21, 2024 · 25595da · 25595da
1 parent 49a3320
commit 25595da
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/model_servers/llamacpp_python/src/requirements.txt b/model_servers/llamacpp_python/src/requirements.txt
@@ -1,2 +1,3 @@
 llama-cpp-python[server]==0.2.57
+transformers==4.41.2
 pip==24.0
diff --git a/model_servers/llamacpp_python/src/run.sh b/model_servers/llamacpp_python/src/run.sh
@@ -5,7 +5,13 @@ if [ ${CONFIG_PATH} ] || [[ ${MODEL_PATH} && ${CONFIG_PATH} ]]; then
 fi
 
 if [ ${MODEL_PATH} ]; then
-    python -m llama_cpp.server --model ${MODEL_PATH} --host ${HOST:=0.0.0.0} --port ${PORT:=8001} --n_gpu_layers ${GPU_LAYERS:=0} --clip_model_path ${CLIP_MODEL_PATH:=None} --chat_format ${MODEL_CHAT_FORMAT:="llama-2"}
+    CMD="python -m llama_cpp.server --model ${MODEL_PATH} --host ${HOST:=0.0.0.0} --port ${PORT:=8001} --clip_model_path ${CLIP_MODEL_PATH:=None} --n_gpu_layers ${GPU_LAYERS:=0} --chat_format ${MODEL_CHAT_FORMAT:="llama-2"}"
+
+    if [ ! -z "${HF_PRETRAINED_MODEL}" ] && [ "${HF_PRETRAINED_MODEL}" != "None" ]; then
+        CMD+=" --hf_pretrained_model_name_or_path ${HF_PRETRAINED_MODEL}"
+    fi
+
+    eval $CMD
     exit 0
 fi