BAAI/bge-reranker-v2-m3

[tommasofavaron1]

I'm trying on running reranking model on ai-djl but i was not able to under stard the data format required to perform this operation:
for now i stucked here but i dont know how to go on
I convert the model using djl/huggingFace converter and save in local forlder
here my code:

Criteria<String[], Classifications[]> criteria =
    Criteria.builder()
            .setTypes(String[].class, Classifications[].class)
            .optModelPath(Paths.get("MYPATH"))
            .optEngine("PyTorch")
            .optProgress(new ProgressBar())
            .optTranslatorFactory(new TextClassificationTranslatorFactory())
            .build();
ZooModel<String[], Classifications[]> rankerModel = criteria.loadModel();
Predictor<String[], Classifications[]> predictorRerank = rankerModel.newPredictor();
String[] array = {"pair1", "pair2"};
Classifications[] a = predictorRerank.predict(array);

has someone come in my same problem?

[frankfliu]

You should use CrossEncoderTranslatorFactory:

        Criteria<StringPair, float[]> criteria =
                Criteria.builder()
                        .setTypes(StringPair.class, float[].class)
                        .optModelPath(path)
                        .optEngine("PyTorch")
                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
                        .optArgument("maxLength", 512)
                        .optArgument("sigmoid", false)
                        .build();

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

I created a PR that allows you to disable sigmoid: #3087

[AzeemWaqarRao]

Hi,
I wanted to use "djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2" as cross-encoder to perform reranking so I checked out to this commit (#3087) and used your code.
Everything worked fine but the result i got from predictor.predict() is an Array of length 384 when instead it should be a single score representing the similarity score between the pair of strings.

[frankfliu]

sentence-transformers/all-MiniLM-L6-v2 is text embedding model, the output of the model is 384 embedding. It doesn't have a classifier layer. You need use a CorssEncoder (e.g. BAAI/bge-reranker-v2-m3) model for your case.

[frankfliu]

@AzeemWaqarRao

You can use our command line tool to convert the model: https://docs.djl.ai/master/extensions/tokenizers/index.html#use-command-line

python3 src/main/python/model_converter.py -f PyTorch -m BAAI/bge-reranker-v2-m3

[AzeemWaqarRao]

Thanks for your help, I was able to fix the issue and got the task done.
Below is my code,

        Criteria<StringPair, float[]> criteria =
                Criteria.builder()
                        .setTypes(StringPair.class, float[].class)

                        .optModelPath(Paths.get("D:\\Java\\djl\\extensions\\tokenizers\\src\\main\\python\\djl_converter\\model\\ms-marco-MiniLM-L-6-v2\\ms-marco-MiniLM-L-6-v2.pt"))
//                        .optModelPath(Paths.get("D:\\Java\\djl\\extensions\\tokenizers\\src\\main\\python\\djl_converter\\model\\ms-marco-MiniLM-L-2-v2\\ms-marco-MiniLM-L-2-v2.pt"))
//                        .optModelPath(Paths.get("D:\\Java\\djl\\extensions\\tokenizers\\src\\main\\python\\djl_converter\\model\\bge-reranker-v2-m3\\bge-reranker-v2-m3.pt"))
                        .optEngine("PyTorch")
                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                        .optProgress(new ProgressBar())
                        .optArgument("reranking", true)
                        .optArgument("sigmoid", false)
                        .build();

        try (ZooModel<StringPair, float[]> model = criteria.loadModel();
             Predictor<StringPair, float[]> predictor = model.newPredictor()) {
            float[] res = predictor.predict(input);
            return res[0];

        }

Just one small thing, the scores I am getting from this code vary from the scores I got from python code when tested on same input and same model.
Even though the model ranks the chunks correctly, I was just wondering why it happens.

[frankfliu]

@AzeemWaqarRao

It turned out the convert script has bug, it didn't include token_type_id during jit trace.
I created a PR to fix this issue: #3318

[AzeemWaqarRao]

Thanks for fixing the issue, it works perfect now. Although my code(which i provided above) does inference on one string pair at at time which makes my overall application slow.
Is there a way to achieve batching or just simply passing the complete list of StringPairs at once?

[frankfliu]

You can use Predictor.batchPredict() function