-
Notifications
You must be signed in to change notification settings - Fork 1.2k
refactor: use extra_body
to pass in input_type
params for asymmetric embedding models for NVIDIA Inference Provider
#3804
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
refactor: use extra_body
to pass in input_type
params for asymmetric embedding models for NVIDIA Inference Provider
#3804
Conversation
98c19d1
to
f4f203a
Compare
…dding models for NVIDIA Inference Provider
f4f203a
to
1d4d263
Compare
"model": embedding_model_id, | ||
"input": [], | ||
} | ||
if is_asymmetric_model(client_with_models, embedding_model_id): |
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.
you can always add extra body unconditionally? if it is null, it is null no harm made. you don't need the double check?
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.
Good point! I also consolidated the verification and retrieval of extra_body
for asymmetric models into a single helper function that's used consistently across all openai_embeddings
test cases.
For other models, return None. | ||
""" | ||
is_asymmetric = is_asymmetric_model(client_with_models, model_id) | ||
if is_asymmetric: |
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.
nah I don't think we need to be that defensive. i'd just simplify this a bunch more. simply make the callsites be:
client.embeddings.create(..., extra_body=get_extra_body())
that's it!
c8d54e3
to
562ad43
Compare
return providers[provider_id] | ||
|
||
|
||
def is_asymmetric_model(client_with_models, model_id): |
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.
where is this used now?
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.
Update to put it in get_extra_body_for_model
: if it is not an asymmetric model, return None
for the extra_body
.
f651ece
to
81e04c7
Compare
81e04c7
to
33190c1
Compare
What does this PR do?
Previously, the NVIDIA inference provider implemented a custom
openai_embeddings
method with a hardcodedinput_type="query"
parameter, which is required by NVIDIA asymmetric embedding models(#3205).Recently
extra_body
parameter is added to the embeddings API (#3794). So, this PR updates the NVIDIA inference provider to use the baseOpenAIMixin.openai_embeddings
method instead and pass theinput_type
through theextra_body
parameter for asymmetric embedding models.Test Plan
Run the following command for the
embedding_model
:nvidia/llama-3.2-nv-embedqa-1b-v2
,nvidia/nv-embedqa-e5-v5
,nvidia/nv-embedqa-mistral-7b-v2
, andsnowflake/arctic-embed-l
.