Expose match_batch_size as envvar for TGI entrypoint

huggingface · Apr 17, 2024 · 2eb609e · 2eb609e
1 parent 30153d5
commit 2eb609e
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/text-generation-inference/entrypoint.sh b/text-generation-inference/entrypoint.sh
@@ -22,6 +22,12 @@ else
   export TGI_MAX_CONCURRENT_REQUESTS 4
 fi
 
+if [[ -n "${TGI_MAX_BATCH_SIZE}" ]]; then
+  export TGI_MAX_BATCH_SIZE="${TGI_MAX_BATCH_SIZE}"
+else
+  export TGI_MAX_BATCH_SIZE 1
+fi
+
 if [[ -n "${TGI_MAX_INPUT_LENGTH}" ]]; then
   export TGI_MAX_INPUT_LENGTH="${TGI_MAX_INPUT_LENGTH}"
 else
@@ -48,6 +54,7 @@ fi
 
 text-generation-launcher --port 8080 \
   --max-concurrent-requests ${TGI_MAX_CONCURRENT_REQUESTS}
+  --max-batch-size ${TGI_MAX_BATCH_SIZE}
   --max-input-length ${TGI_MAX_INPUT_LENGTH} \
   --max-total-tokens ${TGI_MAX_TOTAL_TOKENS} \
   --max-batch-prefill-tokens ${TGI_MAX_BATCH_PREFILL_TOKENS} \