Added Multimodal semantic search feature (opensearch-project#359)

* Adding inference processor and factory, register that in plugin class Signed-off-by: Martin Gaievski <[email protected]>
martin-gaievski · Oct 3, 2023 · 841f280 · 841f280
1 parent 8d749b8
commit 841f280
Show file tree

Hide file tree

Showing 24 changed files with 1,315 additions and 48 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x](https://github.com/opensearch-project/neural-search/compare/2.10...2.x)
 ### Features
 Support sparse semantic retrieval by introducing `sparse_encoding` ingest processor and query builder ([#333](https://github.com/opensearch-project/neural-search/pull/333))
+Added Multimodal semantic search feature ([#359](https://github.com/opensearch-project/neural-search/pull/359))
 ### Enhancements
 Add `max_token_score` parameter to improve the execution efficiency for `neural_sparse` query clause ([#348](https://github.com/opensearch-project/neural-search/pull/348))
 ### Bug Fixes

diff --git a/build.gradle b/build.gradle
@@ -151,6 +151,7 @@ dependencies {
     runtimeOnly group: 'org.reflections', name: 'reflections', version: '0.9.12'
     runtimeOnly group: 'org.javassist', name: 'javassist', version: '3.29.2-GA'
     runtimeOnly group: 'org.opensearch', name: 'common-utils', version: "${opensearch_build}"
+    runtimeOnly group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
     runtimeOnly group: 'com.google.code.gson', name: 'gson', version: '2.10.1'
     runtimeOnly group: 'org.json', name: 'json', version: '20230227'
 }

diff --git a/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java b/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java
@@ -5,6 +5,9 @@
 
 package org.opensearch.neuralsearch.ml;
 
+import static org.opensearch.neuralsearch.processor.TextImageEmbeddingProcessor.INPUT_IMAGE;
+import static org.opensearch.neuralsearch.processor.TextImageEmbeddingProcessor.INPUT_TEXT;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -113,13 +116,30 @@ public void inferenceSentencesWithMapResult(
         retryableInferenceSentencesWithMapResult(modelId, inputText, 0, listener);
     }
 
+    /**
+     * Abstraction to call predict function of api of MLClient with provided targetResponse filters. It uses the
+     * custom model provided as modelId and run the {@link FunctionName#TEXT_EMBEDDING}. The return will be sent
+     * using the actionListener which will have a list of floats in the order of inputText.
+     *
+     * @param modelId {@link String}
+     * @param inputObjects {@link Map} of {@link String}, {@link String} on which inference needs to happen
+     * @param listener {@link ActionListener} which will be called when prediction is completed or errored out.
+     */
+    public void inferenceSentences(
+        @NonNull final String modelId,
+        @NonNull final Map<String, String> inputObjects,
+        @NonNull final ActionListener<List<Float>> listener
+    ) {
+        retryableInferenceSentencesWithSingleVectorResult(TARGET_RESPONSE_FILTERS, modelId, inputObjects, 0, listener);
+    }
+
     private void retryableInferenceSentencesWithMapResult(
         final String modelId,
         final List<String> inputText,
         final int retryTime,
         final ActionListener<List<Map<String, ?>>> listener
     ) {
-        MLInput mlInput = createMLInput(null, inputText);
+        MLInput mlInput = createMLTextInput(null, inputText);
         mlClient.predict(modelId, mlInput, ActionListener.wrap(mlOutput -> {
             final List<Map<String, ?>> result = buildMapResultFromResponse(mlOutput);
             listener.onResponse(result);
@@ -140,7 +160,7 @@ private void retryableInferenceSentencesWithVectorResult(
         final int retryTime,
         final ActionListener<List<List<Float>>> listener
     ) {
-        MLInput mlInput = createMLInput(targetResponseFilters, inputText);
+        MLInput mlInput = createMLTextInput(targetResponseFilters, inputText);
         mlClient.predict(modelId, mlInput, ActionListener.wrap(mlOutput -> {
             final List<List<Float>> vector = buildVectorFromResponse(mlOutput);
             listener.onResponse(vector);
@@ -154,7 +174,7 @@ private void retryableInferenceSentencesWithVectorResult(
         }));
     }
 
-    private MLInput createMLInput(final List<String> targetResponseFilters, List<String> inputText) {
+    private MLInput createMLTextInput(final List<String> targetResponseFilters, List<String> inputText) {
         final ModelResultFilter modelResultFilter = new ModelResultFilter(false, true, targetResponseFilters, null);
         final MLInputDataset inputDataset = new TextDocsInputDataSet(inputText, modelResultFilter);
         return new MLInput(FunctionName.TEXT_EMBEDDING, null, inputDataset);
@@ -191,4 +211,41 @@ private List<List<Float>> buildVectorFromResponse(MLOutput mlOutput) {
         return resultMaps;
     }
 
+    private List<Float> buildSingleVectorFromResponse(final MLOutput mlOutput) {
+        final List<List<Float>> vector = buildVectorFromResponse(mlOutput);
+        return vector.isEmpty() ? new ArrayList<>() : vector.get(0);
+    }
+
+    private void retryableInferenceSentencesWithSingleVectorResult(
+        final List<String> targetResponseFilters,
+        final String modelId,
+        final Map<String, String> inputObjects,
+        final int retryTime,
+        final ActionListener<List<Float>> listener
+    ) {
+        MLInput mlInput = createMLMultimodalInput(targetResponseFilters, inputObjects);
+        mlClient.predict(modelId, mlInput, ActionListener.wrap(mlOutput -> {
+            final List<Float> vector = buildSingleVectorFromResponse(mlOutput);
+            log.debug("Inference Response for input sentence is : {} ", vector);
+            listener.onResponse(vector);
+        }, e -> {
+            if (RetryUtil.shouldRetry(e, retryTime)) {
+                final int retryTimeAdd = retryTime + 1;
+                retryableInferenceSentencesWithSingleVectorResult(targetResponseFilters, modelId, inputObjects, retryTimeAdd, listener);
+            } else {
+                listener.onFailure(e);
+            }
+        }));
+    }
+
+    private MLInput createMLMultimodalInput(final List<String> targetResponseFilters, final Map<String, String> input) {
+        List<String> inputText = new ArrayList<>();
+        inputText.add(input.get(INPUT_TEXT));
+        if (input.containsKey(INPUT_IMAGE)) {
+            inputText.add(input.get(INPUT_IMAGE));
+        }
+        final ModelResultFilter modelResultFilter = new ModelResultFilter(false, true, targetResponseFilters, null);
+        final MLInputDataset inputDataset = new TextDocsInputDataSet(inputText, modelResultFilter);
+        return new MLInput(FunctionName.TEXT_EMBEDDING, null, inputDataset);
+    }
 }
diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
@@ -33,11 +33,13 @@
 import org.opensearch.neuralsearch.processor.NormalizationProcessorWorkflow;
 import org.opensearch.neuralsearch.processor.SparseEncodingProcessor;
 import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor;
+import org.opensearch.neuralsearch.processor.TextImageEmbeddingProcessor;
 import org.opensearch.neuralsearch.processor.combination.ScoreCombinationFactory;
 import org.opensearch.neuralsearch.processor.combination.ScoreCombiner;
 import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory;
 import org.opensearch.neuralsearch.processor.factory.SparseEncodingProcessorFactory;
 import org.opensearch.neuralsearch.processor.factory.TextEmbeddingProcessorFactory;
+import org.opensearch.neuralsearch.processor.factory.TextImageEmbeddingProcessorFactory;
 import org.opensearch.neuralsearch.processor.normalization.ScoreNormalizationFactory;
 import org.opensearch.neuralsearch.processor.normalization.ScoreNormalizer;
 import org.opensearch.neuralsearch.query.HybridQueryBuilder;
@@ -106,7 +108,9 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
             TextEmbeddingProcessor.TYPE,
             new TextEmbeddingProcessorFactory(clientAccessor, parameters.env),
             SparseEncodingProcessor.TYPE,
-            new SparseEncodingProcessorFactory(clientAccessor, parameters.env)
+            new SparseEncodingProcessorFactory(clientAccessor, parameters.env),
+            TextImageEmbeddingProcessor.TYPE,
+            new TextImageEmbeddingProcessorFactory(clientAccessor, parameters.env, parameters.ingestService.getClusterService())
         );
     }
 

diff --git a/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
@@ -249,7 +249,7 @@ protected void setVectorFieldsToDocument(IngestDocument ingestDocument, Map<Stri
     @SuppressWarnings({ "unchecked" })
     @VisibleForTesting
     Map<String, Object> buildNLPResult(Map<String, Object> processorMap, List<?> results, Map<String, Object> sourceAndMetadataMap) {
-        InferenceProcessor.IndexWrapper indexWrapper = new InferenceProcessor.IndexWrapper(0);
+        IndexWrapper indexWrapper = new IndexWrapper(0);
         Map<String, Object> result = new LinkedHashMap<>();
         for (Map.Entry<String, Object> knnMapEntry : processorMap.entrySet()) {
             String knnKey = knnMapEntry.getKey();
@@ -270,7 +270,7 @@ private void putNLPResultToSourceMapForMapType(
         String processorKey,
         Object sourceValue,
         List<?> results,
-        InferenceProcessor.IndexWrapper indexWrapper,
+        IndexWrapper indexWrapper,
         Map<String, Object> sourceAndMetadataMap
     ) {
         if (processorKey == null || sourceAndMetadataMap == null || sourceValue == null) return;
@@ -291,11 +291,7 @@ private void putNLPResultToSourceMapForMapType(
         }
     }
 
-    private List<Map<String, Object>> buildNLPResultForListType(
-        List<String> sourceValue,
-        List<?> results,
-        InferenceProcessor.IndexWrapper indexWrapper
-    ) {
+    private List<Map<String, Object>> buildNLPResultForListType(List<String> sourceValue, List<?> results, IndexWrapper indexWrapper) {
         List<Map<String, Object>> keyToResult = new ArrayList<>();
         IntStream.range(0, sourceValue.size())
             .forEachOrdered(x -> keyToResult.add(ImmutableMap.of(listTypeNestedMapKey, results.get(indexWrapper.index++))));