Skip to content

Commit 9aacc78

Browse files
rbhavnarbhavna
committed
inference processors with local models
1 parent c620964 commit 9aacc78

File tree

5 files changed

+1137
-1087
lines changed

5 files changed

+1137
-1087
lines changed

‎plugin/src/main/java/org/opensearch/ml/processor/InferenceProcessorAttributes.java

+4
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ public class InferenceProcessorAttributes {
1919
protected List<Map<String, String>> outputMaps;
2020

2121
protected String modelId;
22+
protected String functionName;
2223
protected int maxPredictionTask;
2324

2425
protected Map<String, String> modelConfigMaps;
@@ -59,6 +60,7 @@ public class InferenceProcessorAttributes {
5960
*
6061
*/
6162
public static final String OUTPUT_MAP = "output_map";
63+
public static final String FUNCTION_NAME = "function_name";
6264
public static final String MODEL_CONFIG = "model_config";
6365
public static final String MAX_PREDICTION_TASKS = "max_prediction_tasks";
6466

@@ -68,12 +70,14 @@ public class InferenceProcessorAttributes {
6870

6971
public InferenceProcessorAttributes(
7072
String modelId,
73+
String functionName,
7174
List<Map<String, String>> inputMaps,
7275
List<Map<String, String>> outputMaps,
7376
Map<String, String> modelConfigMaps,
7477
int maxPredictionTask
7578
) {
7679
this.modelId = modelId;
80+
this.functionName = functionName;
7781
this.modelConfigMaps = modelConfigMaps;
7882
this.inputMaps = inputMaps;
7983
this.outputMaps = outputMaps;

‎plugin/src/main/java/org/opensearch/ml/processor/MLInferenceIngestProcessor.java

+6-1
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ public class MLInferenceIngestProcessor extends AbstractProcessor implements Mod
6464

6565
protected MLInferenceIngestProcessor(
6666
String modelId,
67+
String functionName,
6768
List<Map<String, String>> inputMaps,
6869
List<Map<String, String>> outputMaps,
6970
Map<String, String> modelConfigMaps,
@@ -78,6 +79,7 @@ protected MLInferenceIngestProcessor(
7879
super(tag, description);
7980
this.inferenceProcessorAttributes = new InferenceProcessorAttributes(
8081
modelId,
82+
functionName,
8183
inputMaps,
8284
outputMaps,
8385
modelConfigMaps,
@@ -184,12 +186,13 @@ private void processPredictions(
184186
}
185187
}
186188

187-
ActionRequest request = getRemoteModelInferenceRequest(modelParameters, inferenceProcessorAttributes.getModelId());
189+
ActionRequest request = getRemoteModelInferenceRequest(modelParameters, inferenceProcessorAttributes.getModelId(), inferenceProcessorAttributes.getFunctionName());
188190

189191
client.execute(MLPredictionTaskAction.INSTANCE, request, new ActionListener<>() {
190192

191193
@Override
192194
public void onResponse(MLTaskResponse mlTaskResponse) {
195+
logger.info("Received response", mlTaskResponse.getOutput().toString());
193196
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) mlTaskResponse.getOutput();
194197
if (processOutputMap == null || processOutputMap.isEmpty()) {
195198
appendFieldValue(modelTensorOutput, null, DEFAULT_OUTPUT_FIELD_NAME, ingestDocument);
@@ -404,6 +407,7 @@ public MLInferenceIngestProcessor create(
404407
Map<String, Object> config
405408
) throws Exception {
406409
String modelId = ConfigurationUtils.readStringProperty(TYPE, processorTag, config, MODEL_ID);
410+
String functionName = ConfigurationUtils.readStringProperty(TYPE, processorTag, config, FUNCTION_NAME);
407411
Map<String, Object> modelConfigInput = ConfigurationUtils.readOptionalMap(TYPE, processorTag, config, MODEL_CONFIG);
408412
List<Map<String, String>> inputMaps = ConfigurationUtils.readOptionalList(TYPE, processorTag, config, INPUT_MAP);
409413
List<Map<String, String>> outputMaps = ConfigurationUtils.readOptionalList(TYPE, processorTag, config, OUTPUT_MAP);
@@ -433,6 +437,7 @@ public MLInferenceIngestProcessor create(
433437

434438
return new MLInferenceIngestProcessor(
435439
modelId,
440+
functionName,
436441
inputMaps,
437442
outputMaps,
438443
modelConfigMaps,

‎plugin/src/main/java/org/opensearch/ml/processor/ModelExecutor.java

+35-3
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,25 @@
66
package org.opensearch.ml.processor;
77

88
import java.io.IOException;
9+
import java.lang.reflect.Type;
910
import java.util.ArrayList;
1011
import java.util.Arrays;
1112
import java.util.List;
13+
import java.util.Locale;
1214
import java.util.Map;
1315
import java.util.stream.Collectors;
1416

17+
import com.google.gson.Gson;
18+
import org.apache.logging.log4j.LogManager;
19+
import org.apache.logging.log4j.Logger;
1520
import org.opensearch.action.ActionRequest;
1621
import org.opensearch.ml.common.FunctionName;
22+
import org.opensearch.ml.common.dataset.TextDocsInputDataSet;
1723
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
1824
import org.opensearch.ml.common.input.MLInput;
25+
import org.opensearch.ml.common.input.parameter.MLAlgoParams;
26+
import org.opensearch.ml.common.input.parameter.textembedding.AsymmetricTextEmbeddingParameters;
27+
import org.opensearch.ml.common.output.model.ModelResultFilter;
1928
import org.opensearch.ml.common.output.model.ModelTensor;
2029
import org.opensearch.ml.common.output.model.ModelTensorOutput;
2130
import org.opensearch.ml.common.output.model.ModelTensors;
@@ -25,12 +34,17 @@
2534
import com.jayway.jsonpath.Configuration;
2635
import com.jayway.jsonpath.JsonPath;
2736
import com.jayway.jsonpath.Option;
37+
import org.opensearch.ml.repackage.com.google.common.reflect.TypeToken;
38+
39+
import static org.opensearch.ml.common.utils.StringUtils.gson;
2840

2941
/**
3042
* General ModelExecutor interface.
3143
*/
3244
public interface ModelExecutor {
3345

46+
Logger logger = LogManager.getLogger(ModelExecutor.class);
47+
3448
Configuration suppressExceptionConfiguration = Configuration
3549
.builder()
3650
.options(Option.SUPPRESS_EXCEPTIONS, Option.DEFAULT_PATH_LEAF_TO_NULL)
@@ -45,13 +59,31 @@ public interface ModelExecutor {
4559
* @return an ActionRequest instance for remote model inference
4660
* @throws IllegalArgumentException if the input parameters are null
4761
*/
48-
default <T> ActionRequest getRemoteModelInferenceRequest(Map<String, String> parameters, String modelId) {
62+
default <T> ActionRequest getRemoteModelInferenceRequest(Map<String, String> parameters, String modelId, String functionName) {
63+
MLInput mlInput = new MLInput();
4964
if (parameters == null) {
5065
throw new IllegalArgumentException("wrong input. The model input cannot be empty.");
5166
}
52-
RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet.builder().parameters(parameters).build();
67+
if (functionName.equals("remote")) {
68+
RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet.builder().parameters(parameters).build();
69+
mlInput = MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build();
70+
} else if (functionName.equals("text_embedding") || functionName.equals("sparse_encoding")) {
71+
Gson gson = new Gson();
72+
String textDocs = parameters.getOrDefault("text_docs", "");
73+
if (!textDocs.startsWith("[") || !textDocs.endsWith("]") ) {
74+
textDocs = "[\"" + textDocs + "\"]";
75+
}
76+
List<String> docs = gson.fromJson(textDocs, List.class);
77+
Boolean returnBytes = gson.fromJson(parameters.getOrDefault("return_bytes", "false"), Boolean.class);
78+
Boolean returnNumber = gson.fromJson(parameters.getOrDefault("return_number", "true"), Boolean.class);
79+
List<String> targetResponse = gson.fromJson(parameters.getOrDefault("target_response", "[]"), List.class);
80+
Type listType = new TypeToken<List<Integer>>() {}.getType();
81+
List<Integer> targetResponsePositions = gson.fromJson(parameters.getOrDefault("target_response_positions", "[]"), listType);
82+
ModelResultFilter resultFilter = new ModelResultFilter(returnBytes, returnNumber, targetResponse, targetResponsePositions);
5383

54-
MLInput mlInput = MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build();
84+
TextDocsInputDataSet inputDataSet = TextDocsInputDataSet.builder().docs(docs).resultFilter(resultFilter).build();
85+
mlInput = MLInput.builder().algorithm(FunctionName.TEXT_EMBEDDING).inputDataset(inputDataSet).build();
86+
}
5587

5688
ActionRequest request = new MLPredictionTaskRequest(modelId, mlInput, null);
5789

‎plugin/src/test/java/org/opensearch/ml/processor/InferenceProcessorAttributesTests.java

+3
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,12 @@
1313
import java.util.List;
1414
import java.util.Map;
1515

16+
import org.junit.Ignore;
1617
import org.junit.Test;
1718

1819
public class InferenceProcessorAttributesTests {
1920

21+
@Ignore
2022
public void testConstructor() {
2123
String modelId = "my_model";
2224
List<Map<String, String>> inputMap = new ArrayList<>();
@@ -34,6 +36,7 @@ public void testConstructor() {
3436

3537
InferenceProcessorAttributes mlModelUtil = new InferenceProcessorAttributes(
3638
modelId,
39+
"text_embedding",
3740
inputMap,
3841
outputMap,
3942
modelConfig,

0 commit comments

Comments
 (0)