How to use Pretrained Hugging face all-MiniLM-L6-v2 mode using java Deep Java Library

Question

How to use Pretrained Hugging face all-MiniLM-L6-v2 mode using java. Was able to load the model but facing issues when predicting.Tried writing a custom translator with String input and float output but didnt work .Any examples with Translator would help.

Answer 1

You can use DJL's built-in TextEmbeddingTranslatorFactory:

String text = "This is an example sentence";

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())
        .build();

try (ZooModel<String, float[]> model = criteria.loadModel();
        Predictor<String, float[]> predictor = model.newPredictor()) {
     float[] res = predictor.predict(text);
     System.out.println("Embedding: " + Arrays.toString(res));
}

See djl-demo for more huggingface examples.