Sentence Embedding difference between DJL Java code and equivalent Python code for Bert model

[legalizenet]

I have a following code in python which accepts sentence and computes embedding for the sentence..

import torch
from transformers import BertTokenizer, BertModel
sen = ['Hello World']
tokenizer = BertTokenizer.from_pretrained('<local path to bert-large-uncased huggingface sentence transformer model>')
model = BertModel.from_pretrained('<local path to bert-large-uncased huggingface sentence transformer model>')
encoded_input = tokenizer.batch_encode_plus(sen, padding=True, truncation=True, max_length=128, return_tensors='pt', add_special_tokens=True)
input_ids = encoded_input['input_ids']
attention_mask = encoded_input['attention_mask']
with torch.no_grad(): 
    outputs = model(input_ids, attention_mask=attention_mask)
    word_embeddings = outputs.last_hidden_state.mean(dim=1)
    print(word_embeddings.shape)    

This prints torch.size([1,1024])

I'm trying to replicate the same with Java code to get the embedding output with following code..

String text = "Hello World";
Criteria<String, float[]> criteria =
    Criteria.builder()
            .setTypes(String.class, float[].class)
            .optModelUrls(
                    "<path to bert-large-uncased huggingface sentence transformer model>")
            .optEngine("PyTorch")
            .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
            .optProgress(new ProgressBar())
            .build();

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

I see the length printed as Embedding length: 30522 <- which is length of vocab instead of what I see in Python code output.

What am I missing, how do I get average vector of 1024 for the same in Java ?

Im pretty sure there is something simple Im missing.

[legalizenet]

In short I would expect embedding float[] to be 768 for bert base model and 1024 for bert large model. Why is it returning 30522 embedding array.

[frankfliu]

Can you share your model_id?
You can try the models in our model zoo, they are generating 1024 embeddings:

        String text = "Hello world";
        Criteria<String, float[]> criteria =
                Criteria.builder()
                        .setTypes(String.class, float[].class)
                        .optModelUrls("djl://ai.djl.huggingface.pytorch/sentence-transformers/bert-large-nli-cls-token")
                        .optEngine("PyTorch")
                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                        .build();

        try (ZooModel<String, float[]> model = criteria.loadModel();
                Predictor<String, float[]> predictor = model.newPredictor()) {
            float[] res = predictor.predict(text);
            Assert.assertEquals(res.length, 1024);
        }

[legalizenet]

@frankfliu ,

I have tried large and base models both.. e.g. try this model -> "djl://ai.djl.huggingface.pytorch/bert-base-uncased" , I get res.length as 30522 , instead of 768

[frankfliu]

This djl://ai.djl.huggingface.pytorch/bert-base-uncased model is NOT an sentence transformers model. This is a fill musk model. The output is logits, a probability for each vocab.

[legalizenet]

@frankfliu - I used the same model on following Python code and I get 768 as length.

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('bert-base-uncased')
embedding = model.encode("Hello world")
print (len(embedding))

prints 768

and if I use BertModel, BertTokenizer , after saving above model to my local machine

from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained(<local path to bert-base-uncased-sf>)
model = BertModel.from_pretrained(<local path to bert-base-uncased-sf>)
encoded_input = tokenizer.batch_encode_plus("Hello world", padding=True, truncation=True, max_length=128, return_tensors='pt', add_special_tokens=True)
input_ids = encoded_input['input_ids']
attention_mask = encoded_input['attention_mask']
with torch.no_grad():
     outputs = model(input_ids, attention_mask=attention_mask)
      word_embeddings = outputs.last_hidden_state.mean(dim=1)
      print(word_embeddings.shape)

i see torch.Size([11, 768])

So the question is if its same model , how come in python library it prints the right dimension and does not print a right dimension in java DJL library ?

[frankfliu]

When we import bert-base-uncased to our model zoo, we use pipeline, it actually load the model as BertForMaskedLM. So the output will be different.

If you want to use base bert model. You can trace the model with BertModel, and it will have the same output.

[legalizenet]

How do I trace the model with BertModel to get the same output in DJL Java library ? , any help is appreciated ...

[frankfliu]

See this document: https://docs.djl.ai/master/docs/pytorch/how_to_convert_your_model_to_torchscript.html

[frankfliu]

    sen = ['Hello World']
    tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
    model = BertModel.from_pretrained('bert-large-uncased')
    encoded_input = tokenizer.batch_encode_plus(sen, padding=True, truncation=True, max_length=128, return_tensors='pt',
                                                add_special_tokens=True)
    input_ids = encoded_input['input_ids']
    attention_mask = encoded_input['attention_mask']

    traced_script_module = torch.jit.trace(model, (input_ids, attention_mask), strict=False)
    # Save the TorchScript model
    traced_script_module.save("traced_bert_model.pt")

[legalizenet]

@frankfliu - Thanks for the info. I really appreciate it , I will try this and get back soon.

Meanwhile one question....

if the purpose of tracing and generating traced .pt file for DJL library to produce same output as Python code , then in above example we are passing sen="Hello World" to get input_ids and attention_mask specific to encoded input for a given input ("Hello World"). and using these input_ids and attention_mask to generate traced bert model which can be used for any text ?
Should not tracing / traced bert model be very generic ?, almost without input_id & attention_mask (because they are specific to a given text input) ?

[legalizenet]

OK.. so its a dummy input as per the documentation here - https://huggingface.co/docs/transformers/torchscript#dummy-inputs-and-standard-lengths . Which also means dummy input must be as long as your potential input would be for text inferencing. which is impossible to know what kind of text or how many words would be in that text.

BTW, that worked. Thanks for pointing me in a right direction. I appreciate it..
I created a traced bert model used it in DJL for getting embedding and now the length of embedding is 1024 for bert large .

But .....

• Average vector produced by Python code above (Sentence Embedding difference between DJL Java code and equivalent Python code for Bert model #2895 (reply in thread)) is like -> [-0.42412397 -0.13317452 -0.4207553 0.28615868 -0.112585396....]
• Average vector produced by DJL Java code above (Sentence Embedding difference between DJL Java code and equivalent Python code for Bert model #2895 (comment)) with traced bert mode is like -> [-0.0334651, -0.010507952, -0.03319929, 0.022579163....]

They do not match at all, not even close.. so im not sure if I run cosine similarity on DJL would this even be accurate ?

[frankfliu]

The default implementation in DJL TextEmbeddingTranslatorFactory did an normalize, but you didn't normalize in your python code.

You can disable normalize in DJL to get the same output:

        String text = "Hello world";
        Path path = Paths.get("traced_bert_model.pt");
        Criteria<String, float[]> criteria =
                Criteria.builder()
                        .setTypes(String.class, float[].class)
                        .optModelPath(path)
                        .optEngine("PyTorch")
                        .optArgument("normalize", "false")
                        .optArgument("tokenizer", "bert-large-uncased")
                        .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(Arrays.toString(res));
        }

[legalizenet]

That works great now!. Thank you very much for helping out here. .Follow up questions

• .optArgument("tokenizer", "bert-large-uncased") - > does that make any difference , if I add it or I dont add it ?, by adding it is it indicating engine to do ssomething special ?

• My question about Dummy input is still valid (Sentence Embedding difference between DJL Java code and equivalent Python code for Bert model #2895 (reply in thread)) . is there any way to create traced model without dummy input ?.

[frankfliu]

1. For the "tokenizer"

The TextEnbdeedingTranslater requires a tokenizer, you need to use a matching tokenizer for your model.

When we import huggingface model into our model zoo, we also save the tokenizer into model folder. When you trace your model, you only saved torchscript model (.pt file). You have to save the tokenizer into the same folder, or specify a tokenizer name, so we ca download the tokenizer from huggingface.

2. the dummy input, this is required for tracing a model. The content of dummy input doesn't really matter as long as the shape is valid for the model. You need execute the model, then the trace will record the model graph into a serialized form and save the graph into a file. The input is not recorded. You can read PyTorch official document for jit trace works: https://pytorch.org/tutorials/advanced/cpp_export.html

[legalizenet]

Yes .. I had whole folder for bert model with tokenizer.json in it already , I just swapped base model with traced based model and it worked without me specifying .optArgument("tokenizer", "bert-large-uncased").

Thanks for the link on cpp_export. I will read thru it. I ran couple of examples on dummy input traced model with Python and DLJ , output embeddings are very close. So would not worry too much about accuracy of embedding elements.

Once again! appreciate all the help

[frankfliu]

@legalizenet

DJL and python use the same underly PyTorch native library, the model output should be identical. If the result are different, it can only be we are using different posting processing in the translator. You can always implement your own translator to match the python code.