DJL support for embedding models using sentence-transformers

[pchamart]

Hi,

Per this doc seems like only the below tasks are supported.

Any plans on including feature-extraction task as well in the future?

I'd be great if we can use text embedding models (both bi and cross encoders) from huggingface for e.g

• sentence-transformers/all-MiniLM-L6-v2 (bi-encoder) and
• cross-encoders/mmarco-mMiniLMv2-L12-H384-v1 (cross-encoder).

Thanks!

[frankfliu]

@pchamart What you need is text-embedding task

[frankfliu]

See this demo: https://github.com/deepjavalibrary/djl-demo/blob/master/huggingface/nlp/src/main/java/com/examples/TextEmbedding.java

[pchamart]

Thanks @frankfliu

Using DJL DLC on Amazon SageMaker for inference - in the serving.properties if we specify translatorFactory=TextEmbedding would that suffice.

engine=MPI
translatorFactory=TextEmbedding
option.model_id=sentence-transformers/all-MiniLM-L6-v2
option.trust_remote_code=true
option.tensor_parallel_degree=1
...

Also, can you please confirm if you plan to support cross-encoders as well in the future.
for e.g. cross-encoders/mmarco-mMiniLMv2-L12-H384-v1 (cross-encoder).

[frankfliu]

@pchamart

We do have this model in our model zoo, but this is a text-classification model not for TextEmbedding.

If you want to serve sentence-transformers/all-MiniLM-L6-v2 model, we really recommend you to Java engine instead of using python engine, it at least 2X the throughput.

And you can use CPU container for this model as well: https://github.com/aws/deep-learning-containers/blob/master/available_images.md#djl-cpu-full-inference-containers

image_uri=f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.23.0-cpu-full"
env = {
    "SERVING_LOAD_MODELS":
    "djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2"
}
endpoint_name = sagemaker.utils.name_from_base("textembedding")

model = Model(
    image_uri=image_uri,
    env=env,
    role=role,
)

# deploy the endpoint endpoint
model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge",
    endpoint_name=endpoint_name,
)

[frankfliu]

feel free to reopen this issue if you still have questions.