Router Configuration on Amazon Sagemaker

[kandakji]

How are the TEI router configurations enforced in Sagemaker Endpoints?

Is below the correct way? 👇

from sagemaker.huggingface import HuggingFaceModel

hub = {
  'HF_MODEL_ID':'BAAI/bge-reranker-base',
  'MAX_BATCH_TOKENS': 1000
}
# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
	image_uri=get_huggingface_llm_image_uri("huggingface-tei",version="1.2.3"),
	env=hub,
	role=role, 
)
# deploy model to SageMaker Inference
huggingface_model.deploy(initial_instance_count=1,instance_type="ml.g5.2xlarge")

[OlivierDehaene]

What do you mean by "router configurations"?
@philschmid maybe you have more info?

[kandakji]

More precisely, I want to change the maximum amount of potential tokens within a batch. How can I do this on a SageMaker Endpoint using TEI? Is it through that env variables?

[OlivierDehaene]

Yes, you can just specify a new value or use the default one (16384) by removing the MAX_BATCH_TOKENS arg.