0

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.

Dinesh
  • 11
  • Can you share your code, please? I only know some Java but might be able to help with a starting project/code. – cronoik May 18 '23 at 09:00

1 Answers1

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.

Frank Liu
  • 281
  • 1
  • 4
  • Is there any way we could download the model locally and achieve the same – Dinesh Jul 18 '23 at 20:31
  • Yes, you can trace the model into TorchScript model and load the model for local folder. Use can use the following code convert huggingface model: https://github.com/deepjavalibrary/djl/blob/master/extensions/tokenizers/src/main/python/model_zoo_importer.py Once you converted the model (you will get a .zip file), then you can load the model: `.optModelPath(Paths.get("mymodel.zip"))` – Frank Liu Jul 19 '23 at 22:54