Enable pass query string to input_map in ml inference search response processor

[mingshl]

Description

Enable pass query string to input_map in ml inference search response processor

setting cluster

PUT _cluster/settings
{
  "persistent": {
    "plugins": {
      "ml_commons": {
        "only_run_on_ml_node": "false",
        "model_access_control_enabled": "true",
        "native_memory_threshold": "99"
      }
    }
  }
}

register cross-encoders local model

POST /_plugins/_ml/models/_register
{
  "name": "huggingface/cross-encoders/ms-marco-MiniLM-L-6-v2",
  "version": "1.0.2",
  "model_format": "TORCH_SCRIPT"
}

{
  "task_id": "tQ5p1ZEB4iWlnHsIf2Xw",
  "status": "CREATED"
}

get register task status

GET /_plugins/_ml/tasks/tQ5p1ZEB4iWlnHsIf2Xw

{
  "model_id": "tg5p1ZEB4iWlnHsIh2U9",
  "task_type": "REGISTER_MODEL",
  "function_name": "TEXT_SIMILARITY",
  "state": "COMPLETED",
  "worker_node": [
    "AahcbpI9R2OtId7Wnt0cYA"
  ],
  "create_time": 1725862346467,
  "last_update_time": 1725862356009,
  "is_async": true
}

deploy cross-encoders local model

`POST /_plugins/_ml/models/tg5p1ZEB4iWlnHsIh2U9/_deploy

`

{
  "task_id": "tw5q1ZEB4iWlnHsIo2WX",
  "task_type": "DEPLOY_MODEL",
  "status": "CREATED"
}

get deploy task status

GET /_plugins/_ml/tasks/tw5q1ZEB4iWlnHsIo2WX

 {
  "model_id": "tg5p1ZEB4iWlnHsIh2U9",
  "task_type": "DEPLOY_MODEL",
  "function_name": "TEXT_SIMILARITY",
  "state": "RUNNING",
  "worker_node": [
    "AahcbpI9R2OtId7Wnt0cYA"
  ],
  "create_time": 1725862421392,
  "last_update_time": 1725862421555,
  "is_async": true
}

wait until completed

{
  "model_id": "tg5p1ZEB4iWlnHsIh2U9",
  "task_type": "DEPLOY_MODEL",
  "function_name": "TEXT_SIMILARITY",
  "state": "COMPLETED",
  "worker_node": [
    "AahcbpI9R2OtId7Wnt0cYA"
  ],
  "create_time": 1725862421392,
  "last_update_time": 1725862466568,
  "is_async": true
}

test model predict

POST _plugins/_ml/models/tg5p1ZEB4iWlnHsIh2U9/_predict
{
    "query_text": "today is sunny",
    "text_docs": [
        "how are you",
        "today is sunny",
        "today is july fifth",
        "it is winter"
    ]
}

{
  "inference_results": [
    {
      "output": [
        {
          "name": "similarity",
          "data_type": "FLOAT32",
          "shape": [
            1
          ],
          "data": [
            -11.055183
          ],
          "byte_buffer": {
            "array": "COIwwQ==",
            "order": "LITTLE_ENDIAN"
          }
        }
      ]
    },
    {
      "output": [
        {
          "name": "similarity",
          "data_type": "FLOAT32",
          "shape": [
            1
          ],
          "data": [
            8.969885
          ],
          "byte_buffer": {
            "array": "poQPQQ==",
            "order": "LITTLE_ENDIAN"
          }
        }
      ]
    },
    {
      "output": [
        {
          "name": "similarity",
          "data_type": "FLOAT32",
          "shape": [
            1
          ],
          "data": [
            -5.736347
          ],
          "byte_buffer": {
            "array": "KJC3wA==",
            "order": "LITTLE_ENDIAN"
          }
        }
      ]
    },
    {
      "output": [
        {
          "name": "similarity",
          "data_type": "FLOAT32",
          "shape": [
            1
          ],
          "data": [
            -10.0452175
          ],
          "byte_buffer": {
            "array": "NrkgwQ==",
            "order": "LITTLE_ENDIAN"
          }
        }
      ]
    }
  ]
}

upload index

 PUT /demo-index-0/_doc/1
{
  "dairy": "how are you"
}

 PUT /demo-index-0/_doc/2
{
  "dairy": "today is sunny"
}

 PUT /demo-index-0/_doc/3
{
  "dairy": "today is july fifth"
}

{
  "_index": "demo-index-0",
  "_id": "4",
  "_version": 1,
  "result": "created",
  "_shards": {
    "total": 2,
    "successful": 1,
    "failed": 0
  },
  "_seq_no": 3,
  "_primary_term": 1
}

create search pipeline with query text passing in input_map

PUT /_search/pipeline/my_pipeline
{
  "response_processors": [
    {
      "ml_inference": {
        "tag": "ml_inference",
        "description": "This processor is going to run ml inference during search response",
        "model_id": "tg5p1ZEB4iWlnHsIh2U9",
        "model_input":"{ \"text_docs\": ${input_map.text_docs}, \"query_text\": \"${input_map.query_text}\" }",
        "function_name": "TEXT_SIMILARITY",
        "input_map": [
          {
            "text_docs": "dairy",
             "query_text": "$.query.term.dairy.value"          
         }
        ],
        "output_map": [
          {
            "rank_score": "$.inference_results[*].output[*].data"
          }
        ],
        "full_response_path":false,
        "model_config": {
        },
        "ignore_missing": false,
        "ignore_failure": false
      }
    }
  ]
}

search with search pipeline, scores are added in the response

GET /demo-index-0/_search?search_pipeline=my_pipeline
{
  "query": {
    "term": {
      "dairy": {
        "value": "today"
      }
    }
  }
}

{
  "took": 400,
  "timed_out": false,
  "_shards": {
    "total": 1,
    "successful": 1,
    "skipped": 0,
    "failed": 0
  },
  "hits": {
    "total": {
      "value": 2,
      "relation": "eq"
    },
    "max_score": 0.71566814,
    "hits": [
      {
        "_index": "demo-index-0",
        "_id": "2",
        "_score": 0.71566814,
        "_source": {
          "dairy": "today is sunny",
          "rank_score": [
            3.6144485
          ]
        }
      },
      {
        "_index": "demo-index-0",
        "_id": "3",
        "_score": 0.6333549,
        "_source": {
          "dairy": "today is july fifth",
          "rank_score": [
            3.6144485
          ]
        }
      }
    ]
  }
}

ToDo

currently ml inference processor only support single tensor for local model, need to support multiple tensor parsing as well.

Related Issues

#2897
#2878

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[ylwu-amzn]

Add more details/examples to description ?

[ylwu-amzn]

> Task :opensearch-ml-common:test

StringUtilsTest > testisValidJSONPath_InvalidInputs FAILED
    java.lang.AssertionError at StringUtilsTest.java:476

StringUtilsTest > testisValidJSONPath_EmptyInput FAILED
    java.lang.AssertionError at StringUtilsTest.java:490

1008 tests completed, 2 failed, 2 skipped

[mingshl]

#2897

there are more details in the RFC #2897

[mingshl]

CI failed

OpenSearchConversationalMemoryHandlerTests > classMethod FAILED
Error: Exception in thread "Thread-4" java.lang.NoClassDefFoundError: Could not initialize class org.opensearch.test.OpenSearchTestCase
	at java.base/java.lang.Thread.run(Thread.java:1583)
	Suppressed: java.lang.IllegalStateException: No context information for thread: Thread[id=29, name=Thread-4, state=RUNNABLE, group=TGRP-ConversationalMemoryHandlerITTests]. Is this thread running under a class com.carrotsearch.randomizedtesting.RandomizedRunner runner context? Add @RunWith(class com.carrotsearch.randomizedtesting.RandomizedRunner.class) to your test class. Make sure your code accesses random contexts within @BeforeClass and @AfterClass boundary (for example, static test class initializers are not permitted to access random contexts).
		at com.carrotsearch.randomizedtesting.RandomizedContext.context(RandomizedContext.java:249)
		at com.carrotsearch.randomizedtesting.RandomizedContext.current(RandomizedContext.java:134)
		at com.carrotsearch.randomizedtesting.RandomizedRunner.augmentStackTrace(RandomizedRunner.java:1885)
		at com.carrotsearch.randomizedtesting.RunnerThreadGroup.uncaughtException(RunnerThreadGroup.java:20)
		at java.base/java.lang.Thread.dispatchUncaughtException(Thread.java:2901)
Caused by: java.lang.ExceptionInInitializerError: Exception java.lang.ExceptionInInitializerError [in thread "SUITE-ConversationalMemoryHandlerITTests-seed#[7B60246E9722A278]"]
	at org.opensearch.test.OpenSearchTestCase.<clinit>(OpenSearchTestCase.java:285)
	at java.base/java.lang.Class.forName0(Native Method)
	at java.base/java.lang.Class.forName(Class.java:534)
	at java.base/java.lang.Class.forName(Class.java:513)
	at com.carrotsearch.randomizedtesting.RandomizedRunner$2.run(RandomizedRunner.java:623)
    java.lang.NoClassDefFoundError at Class.java:-2
        Caused by: java.lang.ExceptionInInitializerError at OpenSearchTestCase.java:285

[mingshl]

after discussion, decided to read the query string from input_map, similar to search request processor. requested review again @ylwu-amzn

[ylwu-amzn]

Does this also have size sort , or just contain query part ?

[mingshl]

yes. added test testProcessResponseSuccessReadRequestMetaFieldFromInputMap

[opensearch-trigger-bot]

The backport to 2.x failed:

The process '/usr/bin/git' failed with exit code 1

To backport manually, run these commands in your terminal:

# Fetch latest updates from GitHub
git fetch
# Create a new working tree
git worktree add .worktrees/backport-2.x 2.x
# Navigate to the new working tree
cd .worktrees/backport-2.x
# Create a new branch
git switch --create backport/backport-2899-to-2.x
# Cherry-pick the merged commit of this pull request and resolve the conflicts
git cherry-pick -x --mainline 1 083abad726933629557028047cb27c482fd950ec
# Push it to GitHub
git push --set-upstream origin backport/backport-2899-to-2.x
# Go back to the original working tree
cd ../..
# Delete the working tree
git worktree remove .worktrees/backport-2.x

Then, create a pull request where the base branch is 2.x and the compare/head branch is backport/backport-2899-to-2.x.

[mingshl]

backport failed because of the previous backport 2.x is not merged yet. #3127