Supporting sparse semantic retrieval in neural search #333
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,13 +8,15 @@ | |
| import java.util.ArrayList; | ||
| import java.util.Arrays; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| import lombok.NonNull; | ||
| import lombok.RequiredArgsConstructor; | ||
| import lombok.extern.log4j.Log4j2; | ||
|
|
||
| import org.opensearch.core.action.ActionListener; | ||
| import org.opensearch.core.common.util.CollectionUtils; | ||
| import org.opensearch.ml.client.MachineLearningNodeClient; | ||
| import org.opensearch.ml.common.FunctionName; | ||
| import org.opensearch.ml.common.dataset.MLInputDataset; | ||
|
|
@@ -100,10 +102,39 @@ public void inferenceSentences( | |
| @NonNull final List<String> inputText, | ||
| @NonNull final ActionListener<List<List<Float>>> listener | ||
| ) { | ||
| retryableInferenceSentencesWithVectorResult(targetResponseFilters, modelId, inputText, 0, listener); | ||
| } | ||
|
|
||
| public void inferenceSentencesWithMapResult( | ||
| @NonNull final String modelId, | ||
| @NonNull final List<String> inputText, | ||
| @NonNull final ActionListener<List<Map<String, ?>>> listener | ||
| ) { | ||
| retryableInferenceSentencesWithMapResult(modelId, inputText, 0, listener); | ||
|
There was a problem hiding this comment. why we want to do no retires? There was a problem hiding this comment. I think we are getting consistent with the existing There was a problem hiding this comment. The 0 here doesn't mean no retry, it's only a initial value, this value will be increased till max retry time(3), this can be optimized to decrease in the future to make it more intuitive though. |
||
| } | ||
|
|
||
| private void retryableInferenceSentencesWithMapResult( | ||
| final String modelId, | ||
| final List<String> inputText, | ||
| final int retryTime, | ||
| final ActionListener<List<Map<String, ?>>> listener | ||
| ) { | ||
| MLInput mlInput = createMLInput(null, inputText); | ||
| mlClient.predict(modelId, mlInput, ActionListener.wrap(mlOutput -> { | ||
| final List<Map<String, ?>> result = buildMapResultFromResponse(mlOutput); | ||
| log.debug("Inference Response for input sentence {} is : {} ", inputText, result); | ||
zhichao-aws marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| listener.onResponse(result); | ||
| }, e -> { | ||
| if (RetryUtil.shouldRetry(e, retryTime)) { | ||
| final int retryTimeAdd = retryTime + 1; | ||
| retryableInferenceSentencesWithMapResult(modelId, inputText, retryTimeAdd, listener); | ||
| } else { | ||
| listener.onFailure(e); | ||
| } | ||
| })); | ||
| } | ||
|
|
||
| private void retryableInferenceSentencesWithVectorResult( | ||
| final List<String> targetResponseFilters, | ||
| final String modelId, | ||
| final List<String> inputText, | ||
|
|
@@ -118,7 +149,7 @@ private void inferenceSentencesWithRetry( | |
| }, e -> { | ||
| if (RetryUtil.shouldRetry(e, retryTime)) { | ||
| final int retryTimeAdd = retryTime + 1; | ||
| retryableInferenceSentencesWithVectorResult(targetResponseFilters, modelId, inputText, retryTimeAdd, listener); | ||
| } else { | ||
| listener.onFailure(e); | ||
| } | ||
|
|
@@ -144,4 +175,22 @@ private List<List<Float>> buildVectorFromResponse(MLOutput mlOutput) { | |
| return vector; | ||
| } | ||
|
|
||
| private List<Map<String, ?>> buildMapResultFromResponse(MLOutput mlOutput) { | ||
| final ModelTensorOutput modelTensorOutput = (ModelTensorOutput) mlOutput; | ||
|
There was a problem hiding this comment. Is this cast safe? Or should we check? There was a problem hiding this comment. Same as the above one. Maybe we can use consistent code between the 2 use case. And If they need to be fixed, we can create another PR to fix them. There was a problem hiding this comment. We can fix the other one in the future, but I want to understand how we know for certain this wont cause class cast on invalid input. There was a problem hiding this comment. @zane-neo could you please help answer this? There was a problem hiding this comment. The two inference method we created in neural-search(inferenceWithVectorResult & inferenceWithMapResult) are for text-embedding and sparse-embedding only, and they both return ModelTensorOutput. For later one in the future, it could be expanded to other use cases, e.g. remote inference, but in this case, it's still returning ModelTensorOutput. So it's safe to cast this output to ModelTensorOutput. |
||
| final List<ModelTensors> tensorOutputList = modelTensorOutput.getMlModelOutputs(); | ||
| if (CollectionUtils.isEmpty(tensorOutputList) || CollectionUtils.isEmpty(tensorOutputList.get(0).getMlModelTensors())) { | ||
| throw new IllegalStateException( | ||
| "Empty model result produced. Expected at least [1] tensor output and [1] model tensor, but got [0]" | ||
| ); | ||
| } | ||
| List<Map<String, ?>> resultMaps = new ArrayList<>(); | ||
| for (ModelTensors tensors : tensorOutputList) { | ||
| List<ModelTensor> tensorList = tensors.getMlModelTensors(); | ||
| for (ModelTensor tensor : tensorList) { | ||
| resultMaps.add(tensor.getDataAsMap()); | ||
| } | ||
| } | ||
| return resultMaps; | ||
| } | ||
|
|
||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we typically have PR for one version, if that's for 2.11 it should be under 2.x
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure, will fix this.