Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add modelId and modelText to KnnVectorQueryBuilder #106068

Merged
merged 36 commits into from Mar 18, 2024

Conversation

tteofili
Copy link
Contributor

@tteofili tteofili commented Mar 7, 2024

Make it possible to perform KNN queries by supplying model_text and model_id instead of the query_vector.

This makes use of a QueryVectorBuilder. Supplying a text_embedding query_vector_builder with model_text and model_id instead of the query_vector will result in the generation of a query_vector by calling inference (during query rewrite) on the specified model_id with the supplied model_text.
This is consistent with the way query vectors are built from model_id / model_text in KnnSearchBuilder (DFS phase).

Sample query:

{
    "query": {
        "knn" : {
            "field": "embedding",
            "num_candidates": 10,
            "query_vector_builder": {
                "text_embedding": {
                    "model_id": "bert_base",
                    "model_text": "lucene is all you need"
                }
            }
        }
    }
}

See also https://docs.google.com/document/d/12SYyHbPbCzhPYQ65HiRMesANObvzWDrBDm0Qn9QSwFQ/edit#heading=h.r3mn4wd2it4e

Use QueryVectorBuilder within KnnVectorQueryBuilder to make it
possible to perform knn queries also when a query vector is not
immediately available. Supplying a text_embedding query_vector_builder
with model_text and model_id instead of the query_vector will result
in the generation of a query_vector by calling inference on the
specified model_id with the supplied model_text (during query
rewrite). This is consistent with the way query vectors are built
from model_id / model_text in KnnSearchBuilder (DFS phase).
@tteofili tteofili added the :Search/Search Search-related issues that do not fall into other categories label Mar 7, 2024
@elasticsearchmachine
Copy link
Collaborator

Hi @tteofili, I've created a changelog YAML for you.

Comment on lines +116 to +119
query_vector_builder:
text_embedding:
model_id: text_embedding_model
model_text: "the octopus comforter smells"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I realize that this is how the top level knn is done, but now that inference_id is becoming a thing, and being used in the core server code, maybe we can switch it it?

@davidkyle what do you think? Can we make an inference call directly from server now? so our query DSL turns into

"knn": {
  "inference_id": "foo", 
  "my_dense_vector_field": "what did the fox jump over?"
}

?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the reasoning for the query_vector_builder stuff (QueryVectorBuilder interface, same as top level knn) is that we cannot introduce a direct dependency between KnnVectorQueryBuilder (from server) and TextEmbeddingQueryVectorBuilder (from xpack/plugin/ml) or InferenceAction (from xpach/plugin/core).
Also it seems desirable to have consistent behaviors between top level and DSL knn queries.

I now see we have SemanticTextModelSettings in server, so if we want to "switch" to inference_id to be picked up by some ml code at runtime (which is anyway the case with query_vector_builder) we can probably do that, but I'd consider doing that for both top level and DSL queries eventually.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the near future we will want to rename model_id -> inference_id and model_text to ?inference_text but there are many places it needs to be done and should change all uses at the same time. I expect to keep the model_id option for compatibility.

server now contains the InferenceServiceRegistry which you can use to look up inference objects configured in _inference and call infer on those objects directly, but it doesn't have access to the models in _ml/trained_models. I'm assuming the purpose of this PR is to create parity between the top level knn and the knn query, in which case it is best to use the same implementation for now.

In future we will automatically create _inference configurations for text embedding models uploaded with Eland, at that point using the InferenceServiceRegistry becomes an option

@tteofili tteofili marked this pull request as ready for review March 11, 2024 08:37
@elasticsearchmachine elasticsearchmachine added the Team:Search Meta label for search team label Mar 11, 2024
@elasticsearchmachine
Copy link
Collaborator

Pinging @elastic/es-search (Team:Search)

not both. Refer to <<knn-semantic-search>> to learn more.
end::knn-query-vector-builder[]


Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you need to update docs/reference/query-dsl/knn-query.asciidoc as well to include knn-query-vector-builder

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can I use include directive inside docs/reference/query-dsl/knn-query.asciidoc ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should be able to use include

Comment on lines 124 to 131
(Required, array of floats)
(Optional, array of floats)
include::{es-repo-dir}/rest-api/common-parms.asciidoc[tag=knn-query-vector]
====

`query_vector_builder`::
(Optional, object)
include::{es-repo-dir}/rest-api/common-parms.asciidoc[tag=knn-query-vector-builder]

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This API is deprecated. We shouldn't really add features to it. Maybe it occurs by proxy by adding it to the knn query.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

on top of modifying docs/reference/query-dsl/knn-query.asciidoc, should I just revert this change, then?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think so @tteofili. this API is deprecated and folks shouldn't be using it. Everything is under _search now. Either in the knn clause or the knn query.

Comment on lines +116 to +119
query_vector_builder:
text_embedding:
model_id: text_embedding_model
model_text: "the octopus comforter smells"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the near future we will want to rename model_id -> inference_id and model_text to ?inference_text but there are many places it needs to be done and should change all uses at the same time. I expect to keep the model_id option for compatibility.

server now contains the InferenceServiceRegistry which you can use to look up inference objects configured in _inference and call infer on those objects directly, but it doesn't have access to the models in _ml/trained_models. I'm assuming the purpose of this PR is to create parity between the top level knn and the knn query, in which case it is best to use the same implementation for now.

In future we will automatically create _inference configurations for text embedding models uploaded with Eland, at that point using the InferenceServiceRegistry becomes an option

@tteofili
Copy link
Contributor Author

run elasticsearch-ci/part-4

Copy link
Member

@davidkyle davidkyle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM thanks for reorganising the tests

Comment on lines 278 to +280
protected void doXContent(XContentBuilder builder, Params params) throws IOException {
if (queryVectorSupplier != null) {
throw new IllegalStateException("missing a rewriteAndFetch?");
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to be done in wire serialization (StreamOutput) as well.

Copy link
Member

@benwtrent benwtrent left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good.

Could you test inference with nested vectors & inner hits? I suspect we will infer the model twice, once for gathering the nearest docs and another time for inner_hits. This is OK for now and fixing it will be a future optimization.

I just want to make sure this works OK when querying nested vectors & gathering inner_hits.

@benwtrent
Copy link
Member

@tteofili do you mind adding a highlight to the release note associated with this PR?

@elasticsearchmachine
Copy link
Collaborator

@tteofili according to this PR's labels, I need to update the changelog YAML, but I can't because the PR is closed. Please either update the changelog yourself on the appropriate branch, or adjust the labels. Specifically:

  • The PR is labelled release highlight but the changelog has no highlight section

tteofili added a commit to tteofili/elasticsearch that referenced this pull request Apr 12, 2024
tteofili added a commit that referenced this pull request Apr 13, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
>enhancement release highlight :Search/Search Search-related issues that do not fall into other categories Team:Search Meta label for search team v8.14.0
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

5 participants