DJL very slow compared to Python, is this normal?

[grossjohann]

I am trying to compute embeddings based on the HuggingFace model cambridgeltl/SapBERT-from-PubMedBERT-fulltext. I wrote some short Python code that processes a file with about 100,000 sentences and it computes the embeddings in 570s. I also wrote some short Java code using DJL, processing the same file, and it takes about 2,000s, i.e. there is about a factor of four.

I am running this on my M2 MacBook Pro. Is this kind of performance to be expected, or have I done something wrong? If you suspect I've done something wrong, then I'll provide the source code.

In production, we'll be running this on some sort of x64 server. So perhaps the M2 performance characteristics that we observed are irrelevant. Is there documentation that could guide me to the right hardware to use for this, and the corresponding DJL settings?

[frankfliu]

By default, DJL is not using MPS device on M2 machine, on linux machine you should expect the same performance.

[grossjohann]

Thank you, I will try it out.

[frankfliu]

@grossjohann

You can try the following code to use mps device, it should work, but we didn't really run benchmark for mps:

        Criteria<String, float[]> criteria =
                Criteria.builder()
                        .setTypes(String.class, float[].class)
                        .optModelUrls(
                                "djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2")
                        .optEngine("PyTorch")
                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                        .optProgress(new ProgressBar())
                        .optDevice(Device.of("mps", -1))
                        .build();