FILL-MASK task code question

[LMartinezEXEX]

Hi everyone!

I am learning DJL and thought it will help me a lot if I started a little project. My main (noob) goal is to get the top K words of preference of a BERT-based model to fill a sentence (i.e. fill mask task). Digging around I found this test of different Translators, one of which achieves what I was trying to do, but I don't quite understand the logits array and hardcoded numbers used within it.

String text = "Hello I'm a [MASK] model.";
Block block =  new LambdaBlock(
                        a -> {
                            NDManager manager = a.getManager();
                            float[][] logits = new float[10][4828];
                            logits[6][4827] = 5;
                            logits[6][2535] = 4;
                            logits[6][2047] = 3;
                            logits[6][3565] = 2;
                            logits[6][2986] = 1;
                            NDArray arr = manager.create(logits);
                            arr = arr.expandDims(0);
                            return new NDList(arr);
                        },
                        "model");
[...]

I suspect it has a close relationship with the text variable defined before but I don't found clear documentation of this because when I change the text the results seem to be always some mathematical symbols.

My goal will be to be able to provide any length sentence and get some more coherent results. If you know more documentation on this would be greatly appreciated!

[frankfliu]

@LMartinezEXEX
This is a mock test, we created a Block that mock the real model's output.

In real case, you need load a huggingface model from the folder, or using existing DJL's huggingface model zoo:

String text = "Hello I'm a [MASK] model.";
Criteria<String, Classifications> criteria =
        Criteria.builder()
                .setTypes(String.class, Classifications.class)
                .optModelUrls("djl://ai.djl.huggingface.pytorch/bert-base-uncased")
                .optEngine("PyTorch")
                .build();

try (ZooModel<String, Classifications> model = criteria.loadModel();
        Predictor<String, Classifications> predictor = model.newPredictor()) {
    Classifications res = predictor.predict(text);
}