Save & Package a custom PyTorch model

[s-bernard]

Hi,

TLDR: I want to create my own private Zoo.

I’ve successfully saved a PyTorch model (actually just a standard Bert) with:

torch.jit.save(traced_model, "traced_bert.pt")

Then I load it with (I had to define BertUncasedTranslator first):

val criteria = Criteria.builder()
        .setTypes(classOf[String], classOf[Array[Float]])
        .optModelPath(Path.of("traced_bert.pt"))
        .optEngine("PyTorch")
        .optProgress(ProgressBar())
        .optTranslator(BertUncasedTranslator())
        .build()
val model = criteria.loadModel()

Then I can create a predictor and use it and everything works as expected:

val predictor = model.newPredictor()
predictor.predict("...")

Now, I want to save it in a more standard DJL way so I don’t have to redefine a Translator or the engine each time I load it. Is there a way? Ideally, I want to create my own private Zoo. I’ve looked at the documentation but it’s quite scarce.

Moreover, if I try to call model.save it fails with java.lang.IllegalStateException: Model has not be trained or loaded yet.
If I try to get model.describeInput() or model.describeOutput(), I just get null.

Thanks!

Note: my code examples are in Scala

[frankfliu]

The model.save() function is only meant to save trained weights. What you need actually is to package a serving ready model. Please take a look this document: https://docs.djl.ai/docs/create_serving_ready_model.html

Actually we have quite a few Huggingface models already imported in DJL model zoo, you can use our built-in translators by adding a serving.properties file in your model zip file. You can take a look our example: https://mlrepo.djl.ai/model/nlp/text_embedding/ai/djl/huggingface/pytorch/sentence-transformers/all-MiniLM-L12-v2/0.0.1/all-MiniLM-L12-v2.zip

[s-bernard]

Thanks a lot, I’ll look at the doc and the example you gave me.

If I understand well, for a base bert uncased, I should be able to reuse sentence transformers and not have to code my custom Translator, right?

[frankfliu]

We imported the quite a few popular sentence transformers model in DJL model zoo, you can use the following url to load them:

• djl://ai.djl.huggingface.pytorch/shibing624/text2vec-base-chinese
• djl://ai.djl.huggingface.pytorch/sentence-transformers/distilbert-multilingual-nli-stsb-quora-ranking/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/LaBSE/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/all-mpnet-base-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/all-mpnet-base-v1/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/paraphrase-MiniLM-L6-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L12-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/paraphrase-MiniLM-L3-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/paraphrase-xlm-r-multilingual-v1/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/bert-base-nli-mean-tokens/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/distilbert-base-nli-mean-tokens/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/multi-qa-MiniLM-L6-cos-v1/
• djl://ai.djl.huggingface.pytorch/sentence-transformers/paraphrase-multilingual-mpnet-base-v2/

Use the following command the list models in our model zoo:

cd djl
./gradlew listModels

We also has python code that help you download huggingface model and generate the .zip file: https://github.com/deepjavalibrary/djl/blob/master/extensions/tokenizers/src/main/python/model_zoo_importer.py

[s-bernard]

My goal here is to get bert embeddings (but I don’t want to be limited to bert, it’s just an example here).

So I used ai.djl.huggingface.tokenizers.HuggingFaceTokenizer and ai.djl.huggingface.translator.TextEmbeddingTranslator named as autoTranslator and then:

val autoCriteria = Criteria.builder()
        .optApplication(Application.NLP.TEXT_EMBEDDING)
        .setTypes(classOf[String], classOf[Array[Float]])
        .optModelPath(Path.of("traced_bert.pt"))
        .optEngine("PyTorch")
        .optProgress(ProgressBar())
        .optTranslator(autoTranslator)
        .build()
val autoModel = autoCriteria.loadModel()

Note: traced_bert.pt is bert_base_uncased that I traced and saved on disk.

It seems to work well but the predictor is not working:

autoModel.newPredictor().predict("hello")
ai.djl.translate.TranslateException: java.lang.NullPointerException: Cannot invoke "ai.djl.ndarray.NDArray.getShape()" because "embeddings" is null
  ai.djl.inference.Predictor.batchPredict(Predictor.java:189)
  ai.djl.inference.Predictor.predict(Predictor.java:126)
  ammonite.$sess.cmd42$Helper.<init>(cmd42.sc:1)
  ammonite.$sess.cmd42$.<clinit>(cmd42.sc:64952)
java.lang.NullPointerException: Cannot invoke "ai.djl.ndarray.NDArray.getShape()" because "embeddings" is null
  ai.djl.huggingface.translator.TextEmbeddingTranslator.processOutput(TextEmbeddingTranslator.java:70)
  ai.djl.huggingface.translator.TextEmbeddingTranslator.processOutput(TextEmbeddingTranslator.java:30)
  ai.djl.inference.Predictor.processOutputs(Predictor.java:225)
  ai.djl.inference.Predictor.batchPredict(Predictor.java:183)
  ai.djl.inference.Predictor.predict(Predictor.java:126)
  ammonite.$sess.cmd42$Helper.<init>(cmd42.sc:1)
  ammonite.$sess.cmd42$.<clinit>(cmd42.sc:64952)

Any idea why list.get("last_hidden_state") is not defined? The way I was doing it in my custom Translator was to take the last element of list.getResourceNDArrays.

[frankfliu]

If you use Huggingface AutoModel or pipeline, they should have "last_hidden_state" as model output.
If you are using a customized model, you need to create your own Translator

[s-bernard]

Ok thanks a lot!

[s-bernard]

I understood why my output did not have "last_hidden_state". It’s because I was loading my model in Python with BertModel.from_pretrained("bert-base-uncased", torchscript=True) which define a tuple as output instead of a dict with the names.

By removing the torchscript=True and using strict=False in torch.jit.trace, it works as expected.