having trouble loading other embedding models (GTE & BGE)

[sandys]

hi folks,
So we are already running minilm using DJL and it works perfectly.

However, we tried the same method to run BGE and GTE embedding models and it doesnt work.

The onnx is here https://huggingface.co/Supabase/bge-small-en/tree/main/onnx , but when we run it in djl, it gives an error "Input mismatch, looking for: [input_ids, attention_mask token_type_ids]"

anyone know how to run this ?

[frankfliu]

Can you share your java code?

[sandys]

thanks for replying - we are pretty much using this https://github.com/deepjavalibrary/djl-demo/blob/master/huggingface/nlp/src/main/java/com/examples/TextEmbedding.java except with optEngine as ONNX . I can pull out the code if you want (its split into a class structure).

[sandys]

i just pulled this out here

 public Acme<BgeSmallResponse> createEmbeddings(String input) {

    return new Acme<>(
        Observable.create(
            emitter -> {
              try {
                Predictor<String, float[]> predictor =
                          loadSmallBgeEn(endpoint.getModelPath()).newPredictor();
                  float[] predict = predictor.predict(input);
                  List<Float> floatList = new LinkedList<>();
                  for (float v : predict) {
                    floatList.add(v);
                  }

                  emitter.onNext(new BgeSmallResponse(floatList));
                  emitter.onComplete();
              } catch (final Exception e) {
                emitter.onError(e);
              }
            }),
        endpoint);
  }

  private ZooModel<String, float[]> loadSmallBgeEn(String modelPath) {

    ZooModel<String, float[]> r = bgeSmallEn;

    if (r == null) {
      synchronized (this) {
        r = bgeSmallEn;
        if (r == null) {
          Path path = Paths.get(modelPath);
          Criteria<String, float[]> criteria =
              Criteria.builder()
                  .setTypes(String.class, float[].class)
                  .optModelPath(path)
                  .optModelName("onnx/model.onnx")
                  .optEngine("OnnxRuntime")
                  .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                  .optProgress(new ProgressBar())
                  .build();

          try {
            r = criteria.loadModel();
            bgeSmallEn = r;
          } catch (IOException | ModelNotFoundException | MalformedModelException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
          }
        }
      }
    }
    return r;
  }

hope its not too messy. error occurs at predictor.predict(input)

[frankfliu]

@sandys
I'm able to reproduce your issue. Current TextEmbeddingTranslator only works for PyTorch model. We will add onnx support.

For the time being, you can create your own Translator to resolve this issue:

public class TextEmbedding {

    private TextEmbedding() {}

    public static void main(String[] args) throws IOException, ModelException, TranslateException {
        String text = "This is an example sentence";
        Path path = Paths.get("model");

        HuggingFaceTokenizer tokenizer =
                HuggingFaceTokenizer.builder()
                        .optTokenizerPath(path)
                        .optManager(NDManager.newBaseManager("PyTorch"))
                        .build();
        MyTextEmbeddingTranslator translator =
                new MyTextEmbeddingTranslator(tokenizer, Batchifier.STACK, "mean", true, true);


        Criteria<String, float[]> criteria =
                Criteria.builder()
                        .setTypes(String.class, float[].class)
                        .optModelPath(path)
                        .optEngine("OnnxRuntime")
                        .optTranslator(translator)
                        .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: " + Arrays.toString(res));
        }
    }


    static final class MyTextEmbeddingTranslator implements Translator<String, float[]> {

        private static final int[] AXIS = {0};

        private HuggingFaceTokenizer tokenizer;
        private Batchifier batchifier;
        private boolean normalize;
        private String pooling;
        private boolean includeTokenTypes;

        MyTextEmbeddingTranslator(
                HuggingFaceTokenizer tokenizer,
                Batchifier batchifier,
                String pooling,
                boolean normalize,
                boolean includeTokenTypes) {
            this.tokenizer = tokenizer;
            this.batchifier = batchifier;
            this.pooling = pooling;
            this.normalize = normalize;
            this.includeTokenTypes = includeTokenTypes;
        }

        /** {@inheritDoc} */
        @Override
        public Batchifier getBatchifier() {
            return batchifier;
        }

        /** {@inheritDoc} */
        @Override
        public NDList processInput(TranslatorContext ctx, String input) {
            Encoding encoding = tokenizer.encode(input);
            ctx.setAttachment("encoding", encoding);
            return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
        }

        /** {@inheritDoc} */
        @Override
        public float[] processOutput(TranslatorContext ctx, NDList list) {
            Encoding encoding = (Encoding) ctx.getAttachment("encoding");
            NDManager manager = ctx.getNDManager();
            NDArray embeddings = processEmbedding(manager, list, encoding, pooling);
            if (normalize) {
                embeddings = embeddings.normalize(2, 0);
            }

            return embeddings.toFloatArray();
        }

        static NDArray processEmbedding(
                NDManager manager, NDList list, Encoding encoding, String pooling) {
            NDArray embedding = list.get("last_hidden_state");
            if (embedding == null) {
                // For Onnx model, NDArray name is not present
                embedding = list.head();
            }
            long[] attentionMask = encoding.getAttentionMask();
            try (NDManager ptManager = NDManager.newBaseManager("PyTorch")) {
                NDArray inputAttentionMask = ptManager.create(attentionMask).toType(DataType.FLOAT32, true);
                switch (pooling) {
                    case "mean":
                        return meanPool(embedding, inputAttentionMask, false);
                    case "mean_sqrt_len":
                        return meanPool(embedding, inputAttentionMask, true);
                    case "max":
                        return maxPool(embedding, inputAttentionMask);
                    case "weightedmean":
                        return weightedMeanPool(embedding, inputAttentionMask);
                    case "cls":
                        return embedding.get(0);
                    default:
                        throw new AssertionError("Unexpected pooling mode: " + pooling);
                }
            }
        }

        private static NDArray meanPool(NDArray embeddings, NDArray attentionMask, boolean sqrt) {
            long[] shape = embeddings.getShape().getShape();
            attentionMask = attentionMask.expandDims(-1).broadcast(shape);
            NDArray inputAttentionMaskSum = attentionMask.sum(AXIS);
            NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
            NDArray prod = embeddings.mul(attentionMask);
            NDArray sum = prod.sum(AXIS);
            if (sqrt) {
                return sum.div(clamp.sqrt());
            }
            return sum.div(clamp);
        }

        private static NDArray maxPool(NDArray embeddings, NDArray inputAttentionMask) {
            long[] shape = embeddings.getShape().getShape();
            inputAttentionMask = inputAttentionMask.expandDims(-1).broadcast(shape);
            inputAttentionMask = inputAttentionMask.eq(0);
            embeddings = embeddings.duplicate();
            embeddings.set(inputAttentionMask, -1e9); // Set padding tokens to large negative value

            return embeddings.max(AXIS, true);
        }

        private static NDArray weightedMeanPool(NDArray embeddings, NDArray attentionMask) {
            long[] shape = embeddings.getShape().getShape();
            NDArray weight = embeddings.getManager().arange(1, shape[0] + 1);
            weight = weight.expandDims(-1).broadcast(shape);

            attentionMask = attentionMask.expandDims(-1).broadcast(shape).mul(weight);
            NDArray maskSum = attentionMask.sum(AXIS);
            NDArray embeddingSum = embeddings.mul(attentionMask).sum(AXIS);
            return embeddingSum.div(maskSum);
        }
    }

[sandys]

@frankfliu thank you so much! really appreciate this.

[frankfliu]

@sandys
There is PR that address your issue. You can try 0.24.0-SNAPSHOT release once the PR is built into nightly release.