asymmetric embeddings

common/src/main/java/org/opensearch/ml/common/input/nlp/TextDocsMLInput.java

@@ -13,6 +13,7 @@
 import org.opensearch.ml.common.dataset.MLInputDataset;
 import org.opensearch.ml.common.dataset.TextDocsInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
+import org.opensearch.ml.common.input.parameter.MLAlgoParams;
 import org.opensearch.ml.common.output.model.ModelResultFilter;
 
 import java.io.IOException;
@@ -22,7 +23,7 @@
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 
 /**
- * ML input class which supports a list fo text docs.
+ * ML input class which supports a list of text docs.
  * This class can be used for TEXT_EMBEDDING model.
  */
 @org.opensearch.ml.common.annotation.MLInput(functionNames = {FunctionName.TEXT_EMBEDDING, FunctionName.SPARSE_ENCODING, FunctionName.SPARSE_TOKENIZE})
@@ -124,6 +125,9 @@
                 case RESULT_FILTER_FIELD:
                     resultFilter = ModelResultFilter.parse(parser);
                     break;
+                case ML_PARAMETERS_FIELD:
+                    this.parameters = parser.namedObject(MLAlgoParams.class, functionName.name(), null);
+                    break;
                 default:
                     parser.skipChildren();
                     break;

common/src/main/java/org/opensearch/ml/common/model/TextEmbeddingModelConfig.java

@@ -37,16 +37,25 @@
     public static final String POOLING_MODE_FIELD = "pooling_mode";
     public static final String NORMALIZE_RESULT_FIELD = "normalize_result";
     public static final String MODEL_MAX_LENGTH_FIELD = "model_max_length";
+    public static final String QUERY_PREFIX = "query_prefix";
+    public static final String PASSAGE_PREFIX = "passage_prefix";
 
     private final Integer embeddingDimension;
     private final FrameworkType frameworkType;
     private final PoolingMode poolingMode;
     private final boolean normalizeResult;
     private final Integer modelMaxLength;
+    private final String queryPrefix;
+    private final String passagePrefix;
+
+    public TextEmbeddingModelConfig(String modelType, Integer embeddingDimension, FrameworkType frameworkType, String allConfig,
+        PoolingMode poolingMode, boolean normalizeResult, Integer modelMaxLength) {
+        this(modelType, embeddingDimension, frameworkType, allConfig, poolingMode, normalizeResult, modelMaxLength, null, null);
+    }
 
     @Builder(toBuilder = true)
     public TextEmbeddingModelConfig(String modelType, Integer embeddingDimension, FrameworkType frameworkType, String allConfig,
-                                    PoolingMode poolingMode, boolean normalizeResult, Integer modelMaxLength) {
+                                    PoolingMode poolingMode, boolean normalizeResult, Integer modelMaxLength, String queryPrefix, String passagePrefix) {
         super(modelType, allConfig);
         if (embeddingDimension == null) {
             throw new IllegalArgumentException("embedding dimension is null");
@@ -59,6 +68,8 @@
         this.poolingMode = poolingMode;
         this.normalizeResult = normalizeResult;
         this.modelMaxLength = modelMaxLength;
+        this.queryPrefix = queryPrefix;
+        this.passagePrefix = passagePrefix;
     }
 
     public static TextEmbeddingModelConfig parse(XContentParser parser) throws IOException {
@@ -69,6 +80,8 @@
         PoolingMode poolingMode = null;
         boolean normalizeResult = false;
         Integer modelMaxLength = null;
+        String queryPrefix = null;
+        String passagePrefix = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -97,12 +110,18 @@
                 case MODEL_MAX_LENGTH_FIELD:
                     modelMaxLength = parser.intValue();
                     break;
+                case QUERY_PREFIX:
+                    queryPrefix = parser.text();
+                    break;
+                case PASSAGE_PREFIX:
+                    passagePrefix = parser.text();
+                    break;
                 default:
                     parser.skipChildren();
                     break;
             }
         }
-        return new TextEmbeddingModelConfig(modelType,  embeddingDimension, frameworkType, allConfig, poolingMode, normalizeResult, modelMaxLength);
+        return new TextEmbeddingModelConfig(modelType,  embeddingDimension, frameworkType, allConfig, poolingMode, normalizeResult, modelMaxLength, queryPrefix, passagePrefix);
     }
 
     @Override
@@ -121,6 +140,8 @@
         }
         normalizeResult = in.readBoolean();
         modelMaxLength = in.readOptionalInt();
+        queryPrefix = in.readOptionalString();
+        passagePrefix = in.readOptionalString();
     }
 
     @Override
@@ -136,6 +157,8 @@
         }
         out.writeBoolean(normalizeResult);
         out.writeOptionalInt(modelMaxLength);
+        out.writeOptionalString(queryPrefix);
+        out.writeOptionalString(passagePrefix);
     }
 
     @Override
@@ -162,6 +185,12 @@
         if (normalizeResult) {
             builder.field(NORMALIZE_RESULT_FIELD, normalizeResult);
         }
+        if (queryPrefix != null) {
+            builder.field(QUERY_PREFIX, queryPrefix);
+        }
+        if (passagePrefix != null) {
+            builder.field(PASSAGE_PREFIX, passagePrefix);
+        }
         builder.endObject();
         return builder;
     }

ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java

@@ -135,6 +135,12 @@
                                     case TextEmbeddingModelConfig.MODEL_MAX_LENGTH_FIELD:
                                         configBuilder.modelMaxLength(((Double) configEntry.getValue()).intValue());
                                         break;
+                                    case TextEmbeddingModelConfig.QUERY_PREFIX:
+                                        configBuilder.queryPrefix(configEntry.getValue().toString());
+                                        break;
+                                    case TextEmbeddingModelConfig.PASSAGE_PREFIX:
+                                        configBuilder.passagePrefix(configEntry.getValue().toString());
+                                        break;
                                     default:
                                         break;
                                 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/DLModel.java

@@ -67,6 +67,8 @@ public abstract class DLModel implements Predictable {
     protected Device[] devices;
     protected AtomicInteger nextDevice = new AtomicInteger(0);
 
+    protected MLModelConfig modelConfig;
+
     @Override
     public MLOutput predict(MLInput mlInput, MLModel model) {
         throw new IllegalArgumentException("model not deployed");
@@ -183,6 +185,7 @@ protected void doLoadModel(
         IOException,
         TranslateException {
         devices = Engine.getEngine(engine).getDevices();
+        this.modelConfig = modelConfig;
         for (int i = 0; i < devices.length; i++) {
             log.debug("load model {} to device {}: {}", modelId, i, devices[i]);
             ZooModel<Input, Output> model;

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/TextEmbeddingModel.java

@@ -4,25 +4,35 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import org.opensearch.ml.common.dataset.MLInputDataset;
 import org.opensearch.ml.common.dataset.TextDocsInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
+import org.opensearch.ml.common.input.parameter.MLAlgoParams;
 import org.opensearch.ml.common.model.MLModelConfig;
 import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
 import org.opensearch.ml.common.output.model.ModelResultFilter;
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
+import org.opensearch.ml.engine.algorithms.text_embedding.AsymmetricTextEmbeddingParameters;
+import org.opensearch.ml.engine.algorithms.text_embedding.AsymmetricTextEmbeddingParameters.EmbeddingContentType;
 
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.translate.TranslateException;
 
 public abstract class TextEmbeddingModel extends DLModel {
+
     @Override
     public ModelTensorOutput predict(String modelId, MLInput mlInput) throws TranslateException {
-        MLInputDataset inputDataSet = mlInput.getInputDataset();
+        MLAlgoParams mlParams = mlInput.getParameters();
+
+        MLInputDataset inputDataSet = isAsymmetricModel(mlParams)
+            ? addPrefixesToData((AsymmetricTextEmbeddingParameters) mlParams, (TextDocsInputDataSet) mlInput.getInputDataset())
+            : mlInput.getInputDataset();
+
         List<ModelTensors> tensorOutputs = new ArrayList<>();
         Output output;
         TextDocsInputDataSet textDocsInput = (TextDocsInputDataSet) inputDataSet;
@@ -36,6 +46,52 @@
         return new ModelTensorOutput(tensorOutputs);
     }
 
+    private boolean isAsymmetricModel(MLAlgoParams mlParams) {
+        if (mlParams instanceof AsymmetricTextEmbeddingParameters) {
+            // Check for the necessary prefixes in modelConfig
+            if (modelConfig == null
+                || ((TextEmbeddingModelConfig) modelConfig).getPassagePrefix() == null
+                    && ((TextEmbeddingModelConfig) modelConfig).getQueryPrefix() == null) {
+                throw new IllegalArgumentException(
+                    "When passing AsymmetricTextEmbeddingParameters, the model requires to be "
+                        + "registered with at least one of `query_prefix` or `passage_prefix`."
+                );
+            }
+            // Passed all checks
+            return true;
+        } else if (mlParams != null) {
+            // mlParams is not null and not an instance of AsymmetricTextEmbeddingParameters.
+            // Should never happen.
+            throw new IllegalArgumentException("TEXT_EMBEDDING algorithm only supports AsymmetricTextEmbeddingParameters.");
+        }
+
+        // mlParams is null, but the model is asymmetric.
+        if (modelConfig != null
+            && (((TextEmbeddingModelConfig) modelConfig).getPassagePrefix() != null
+                || ((TextEmbeddingModelConfig) modelConfig).getQueryPrefix() != null)) {
+            throw new IllegalArgumentException(
+                "The embedding model chosen is asymmetric. To use it, you must declare whether the input is a query or a passage."
+            );
+        }
+
+        return false;
+    }
+
+    private TextDocsInputDataSet addPrefixesToData(AsymmetricTextEmbeddingParameters mlParams, TextDocsInputDataSet inputDataSet) {
+        // Asymmetric embedding models typically work with "mini-prompts" that prime the model to embed a text
+        // as a query or as a passage. Here we apply the prompt as defined in the model configuration. We default
+        // to passage embedding.
+        TextEmbeddingModelConfig modelConfig = (TextEmbeddingModelConfig) this.modelConfig;
+        String prefix = mlParams.getEmbeddingContentType() == EmbeddingContentType.QUERY
+            ? modelConfig.getQueryPrefix()
+            : modelConfig.getPassagePrefix();
+        if (prefix != null) {
+            List<String> prefixedDocs = inputDataSet.getDocs().stream().map(s -> prefix + s).collect(Collectors.toList());
+            return TextDocsInputDataSet.builder().docs(prefixedDocs).build();
+        }
+        return inputDataSet;
+    }
+
     public void warmUp(Predictor predictor, String modelId, MLModelConfig modelConfig) throws TranslateException {
         TextEmbeddingModelConfig textEmbeddingModelConfig = (TextEmbeddingModelConfig) modelConfig;
         String warmUpSentence = "warm up sentence";

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/text_embedding/AsymmetricTextEmbeddingParameters.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.engine.algorithms.text_embedding;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+
+import java.io.IOException;
+import java.util.Locale;
+
+import org.opensearch.core.ParseField;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.annotation.MLAlgoParameter;
+import org.opensearch.ml.common.input.parameter.MLAlgoParams;
+
+import lombok.Builder;
+import lombok.Data;
+
+/**
+ * This class defines the modes of operation of an asymmetric text embedding algorithm. The algorithm can be used to embed either a query or a passage.
+ */
+@Data
+@MLAlgoParameter(algorithms = { FunctionName.TEXT_EMBEDDING })
+public class AsymmetricTextEmbeddingParameters implements MLAlgoParams {
+
+    public enum EmbeddingContentType {
+        QUERY,
+        PASSAGE
+    }
+
+    public static final String PARSE_FIELD_NAME = FunctionName.TEXT_EMBEDDING.name();
+    public static final NamedXContentRegistry.Entry XCONTENT_REGISTRY = new NamedXContentRegistry.Entry(
+        MLAlgoParams.class,
+        new ParseField(PARSE_FIELD_NAME),
+        it -> parse(it)
+    );
+
+    @Builder(toBuilder = true)
+    public AsymmetricTextEmbeddingParameters(EmbeddingContentType embeddingContentType) {
+        this.embeddingContentType = embeddingContentType;
+    }
+
+    public AsymmetricTextEmbeddingParameters(StreamInput in) throws IOException {
+        this.embeddingContentType = EmbeddingContentType.valueOf(in.readOptionalString());
+    }
+
+    public static MLAlgoParams parse(XContentParser parser) throws IOException {
+        EmbeddingContentType embeddingContentType = null;
+
+        ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
+        while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+            String fieldName = parser.currentName();
+            parser.nextToken();
+
+            switch (fieldName) {
+                case EMBEDDING_CONTENT_TYPE_FIELD:
+                    String contentType = parser.text();
+                    embeddingContentType = EmbeddingContentType.valueOf(contentType.toUpperCase(Locale.ROOT));
+                    break;
+                default:
+                    parser.skipChildren();
+                    break;
+            }
+        }
+        return new AsymmetricTextEmbeddingParameters(embeddingContentType);
+    }
+
+    public static final String EMBEDDING_CONTENT_TYPE_FIELD = "content_type";
+
+    // The type of the content to be embedded
+    private EmbeddingContentType embeddingContentType;
+
+    @Override
+    public int getVersion() {
+        return 1;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return PARSE_FIELD_NAME;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalString(embeddingContentType.name());
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        xContentBuilder.startObject();
+        if (embeddingContentType != null) {
+            xContentBuilder.field(EMBEDDING_CONTENT_TYPE_FIELD, embeddingContentType.name());
+        }
+        xContentBuilder.endObject();
+        return xContentBuilder;
+    }
+
+    public EmbeddingContentType getEmbeddingContentType() {
+        return embeddingContentType;
+    }
+}