Merge pull request #3577 from BerriAI/litellm_add_triton_server

[Feat] Add Triton Embeddings to LiteLLM
This commit is contained in:
Ishaan Jaff 2024-05-10 19:20:23 -07:00 committed by GitHub
commit b09075da53
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 265 additions and 0 deletions

View file

@ -4814,6 +4814,12 @@ def get_optional_params_embeddings(
status_code=500,
message=f"Setting dimensions is not supported for OpenAI `text-embedding-3` and later models. To drop it from the call, set `litellm.drop_params = True`.",
)
if custom_llm_provider == "triton":
keys = list(non_default_params.keys())
for k in keys:
non_default_params.pop(k, None)
final_params = {**non_default_params, **kwargs}
return final_params
if custom_llm_provider == "vertex_ai":
if len(non_default_params.keys()) > 0:
if litellm.drop_params is True: # drop the unsupported non-default values